ae4bbd0ec6
This PR is quite similiar to #174700. In this PR, I added a C API for each (upstream) MLIR attributes to retrieve its name (for example, `StringAttr -> mlirStringAttrGetName() -> "builtin.string"`), and exposed a corresponding type_name class attribute in the Python bindings (e.g., `StringAttr.attr_name -> "builtin.string"`). This can be used in various places to avoid hard-coded strings, such as eliminating the manual string in `irdl.base("#builtin.string")`. Note that parts of this PR (mainly mechanical changes) were produced via GitHub Copilot and GPT-5.2. I have manually reviewed the changes and verified them with tests to ensure correctness.
124 lines
4.9 KiB
C++
124 lines
4.9 KiB
C++
//===- DialectGPU.cpp - Pybind module for the GPU passes ------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===---------------------------------------------------------------------===//
|
|
|
|
#include "mlir-c/Dialect/GPU.h"
|
|
#include "mlir-c/IR.h"
|
|
#include "mlir-c/Support.h"
|
|
#include "mlir/Bindings/Python/IRAttributes.h"
|
|
#include "mlir/Bindings/Python/IRCore.h"
|
|
#include "mlir/Bindings/Python/Nanobind.h"
|
|
#include "mlir/Bindings/Python/NanobindAdaptors.h"
|
|
|
|
namespace nb = nanobind;
|
|
using namespace nanobind::literals;
|
|
using namespace mlir::python::nanobind_adaptors;
|
|
|
|
namespace mlir {
|
|
namespace python {
|
|
namespace MLIR_BINDINGS_PYTHON_DOMAIN {
|
|
namespace gpu {
|
|
// -----------------------------------------------------------------------------
|
|
// AsyncTokenType
|
|
// -----------------------------------------------------------------------------
|
|
|
|
struct AsyncTokenType : PyConcreteType<AsyncTokenType> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirTypeIsAGPUAsyncTokenType;
|
|
static constexpr const char *pyClassName = "AsyncTokenType";
|
|
static inline const MlirStringRef name = mlirGPUAsyncTokenTypeGetName();
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](DefaultingPyMlirContext context) {
|
|
return AsyncTokenType(context->getRef(),
|
|
mlirGPUAsyncTokenTypeGet(context.get()->get()));
|
|
},
|
|
"Gets an instance of AsyncTokenType in the same context",
|
|
nb::arg("context").none() = nb::none());
|
|
}
|
|
};
|
|
|
|
//===-------------------------------------------------------------------===//
|
|
// ObjectAttr
|
|
//===-------------------------------------------------------------------===//
|
|
|
|
struct ObjectAttr : PyConcreteAttribute<ObjectAttr> {
|
|
static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAGPUObjectAttr;
|
|
static constexpr const char *pyClassName = "ObjectAttr";
|
|
static inline const MlirStringRef name = mlirGPUObjectAttrGetName();
|
|
using Base::Base;
|
|
|
|
static void bindDerived(ClassTy &c) {
|
|
c.def_static(
|
|
"get",
|
|
[](const PyAttribute &target, uint32_t format, const nb::bytes &object,
|
|
std::optional<PyDictAttribute> mlirObjectProps,
|
|
std::optional<PyAttribute> mlirKernelsAttr,
|
|
DefaultingPyMlirContext context) {
|
|
MlirStringRef objectStrRef = mlirStringRefCreate(
|
|
static_cast<char *>(const_cast<void *>(object.data())),
|
|
object.size());
|
|
return ObjectAttr(
|
|
context->getRef(),
|
|
mlirGPUObjectAttrGetWithKernels(
|
|
mlirAttributeGetContext(target), target, format, objectStrRef,
|
|
mlirObjectProps.has_value() ? *mlirObjectProps
|
|
: MlirAttribute{nullptr},
|
|
mlirKernelsAttr.has_value() ? *mlirKernelsAttr
|
|
: MlirAttribute{nullptr}));
|
|
},
|
|
"target"_a, "format"_a, "object"_a, "properties"_a = nb::none(),
|
|
"kernels"_a = nb::none(), "context"_a = nb::none(),
|
|
"Gets a gpu.object from parameters.");
|
|
|
|
c.def_prop_ro("target", [](ObjectAttr &self) {
|
|
return PyAttribute(self.getContext(), mlirGPUObjectAttrGetTarget(self))
|
|
.maybeDownCast();
|
|
});
|
|
c.def_prop_ro("format", [](const ObjectAttr &self) {
|
|
return mlirGPUObjectAttrGetFormat(self);
|
|
});
|
|
c.def_prop_ro("object", [](const ObjectAttr &self) {
|
|
MlirStringRef stringRef = mlirGPUObjectAttrGetObject(self);
|
|
return nb::bytes(stringRef.data, stringRef.length);
|
|
});
|
|
c.def_prop_ro(
|
|
"properties", [](ObjectAttr &self) -> std::optional<PyDictAttribute> {
|
|
if (mlirGPUObjectAttrHasProperties(self))
|
|
return PyDictAttribute(self.getContext(),
|
|
mlirGPUObjectAttrGetProperties(self));
|
|
return std::nullopt;
|
|
});
|
|
c.def_prop_ro("kernels",
|
|
[](ObjectAttr &self)
|
|
-> std::optional<nb::typed<nb::object, PyAttribute>> {
|
|
if (mlirGPUObjectAttrHasKernels(self))
|
|
return PyAttribute(self.getContext(),
|
|
mlirGPUObjectAttrGetKernels(self))
|
|
.maybeDownCast();
|
|
return std::nullopt;
|
|
});
|
|
}
|
|
};
|
|
} // namespace gpu
|
|
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
|
|
} // namespace python
|
|
} // namespace mlir
|
|
|
|
// -----------------------------------------------------------------------------
|
|
// Module initialization.
|
|
// -----------------------------------------------------------------------------
|
|
|
|
NB_MODULE(_mlirDialectsGPU, m) {
|
|
m.doc() = "MLIR GPU Dialect";
|
|
|
|
mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::gpu::AsyncTokenType::bind(m);
|
|
mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::gpu::ObjectAttr::bind(m);
|
|
}
|