92 lines · cpp
1//===- DialectGPU.cpp - Pybind module for the GPU passes ------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===---------------------------------------------------------------------===//8 9#include "mlir-c/Dialect/GPU.h"10#include "mlir-c/IR.h"11#include "mlir-c/Support.h"12#include "mlir/Bindings/Python/Nanobind.h"13#include "mlir/Bindings/Python/NanobindAdaptors.h"14 15namespace nb = nanobind;16using namespace nanobind::literals;17 18using namespace mlir;19using namespace mlir::python;20using namespace mlir::python::nanobind_adaptors;21 22// -----------------------------------------------------------------------------23// Module initialization.24// -----------------------------------------------------------------------------25 26NB_MODULE(_mlirDialectsGPU, m) {27 m.doc() = "MLIR GPU Dialect";28 //===-------------------------------------------------------------------===//29 // AsyncTokenType30 //===-------------------------------------------------------------------===//31 32 auto mlirGPUAsyncTokenType =33 mlir_type_subclass(m, "AsyncTokenType", mlirTypeIsAGPUAsyncTokenType);34 35 mlirGPUAsyncTokenType.def_classmethod(36 "get",37 [](const nb::object &cls, MlirContext ctx) {38 return cls(mlirGPUAsyncTokenTypeGet(ctx));39 },40 "Gets an instance of AsyncTokenType in the same context", nb::arg("cls"),41 nb::arg("ctx") = nb::none());42 43 //===-------------------------------------------------------------------===//44 // ObjectAttr45 //===-------------------------------------------------------------------===//46 47 mlir_attribute_subclass(m, "ObjectAttr", mlirAttributeIsAGPUObjectAttr)48 .def_classmethod(49 "get",50 [](const nb::object &cls, MlirAttribute target, uint32_t format,51 const nb::bytes &object,52 std::optional<MlirAttribute> mlirObjectProps,53 std::optional<MlirAttribute> mlirKernelsAttr) {54 MlirStringRef objectStrRef = mlirStringRefCreate(55 static_cast<char *>(const_cast<void *>(object.data())),56 object.size());57 return cls(mlirGPUObjectAttrGetWithKernels(58 mlirAttributeGetContext(target), target, format, objectStrRef,59 mlirObjectProps.has_value() ? *mlirObjectProps60 : MlirAttribute{nullptr},61 mlirKernelsAttr.has_value() ? *mlirKernelsAttr62 : MlirAttribute{nullptr}));63 },64 "cls"_a, "target"_a, "format"_a, "object"_a,65 "properties"_a = nb::none(), "kernels"_a = nb::none(),66 "Gets a gpu.object from parameters.")67 .def_property_readonly(68 "target",69 [](MlirAttribute self) { return mlirGPUObjectAttrGetTarget(self); })70 .def_property_readonly(71 "format",72 [](MlirAttribute self) { return mlirGPUObjectAttrGetFormat(self); })73 .def_property_readonly(74 "object",75 [](MlirAttribute self) {76 MlirStringRef stringRef = mlirGPUObjectAttrGetObject(self);77 return nb::bytes(stringRef.data, stringRef.length);78 })79 .def_property_readonly("properties",80 [](MlirAttribute self) -> nb::object {81 if (mlirGPUObjectAttrHasProperties(self))82 return nb::cast(83 mlirGPUObjectAttrGetProperties(self));84 return nb::none();85 })86 .def_property_readonly("kernels", [](MlirAttribute self) -> nb::object {87 if (mlirGPUObjectAttrHasKernels(self))88 return nb::cast(mlirGPUObjectAttrGetKernels(self));89 return nb::none();90 });91}92