151 lines · cpp
1//===- PythonTestModuleNanobind.cpp - PythonTest dialect extension --------===//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// This is the nanobind edition of the PythonTest dialect module.9//===----------------------------------------------------------------------===//10 11#include "PythonTestCAPI.h"12#include "mlir-c/BuiltinAttributes.h"13#include "mlir-c/BuiltinTypes.h"14#include "mlir-c/Diagnostics.h"15#include "mlir-c/IR.h"16#include "mlir/Bindings/Python/Diagnostics.h"17#include "mlir/Bindings/Python/Nanobind.h"18#include "mlir/Bindings/Python/NanobindAdaptors.h"19#include "nanobind/nanobind.h"20 21namespace nb = nanobind;22using namespace mlir::python::nanobind_adaptors;23 24static bool mlirTypeIsARankedIntegerTensor(MlirType t) {25 return mlirTypeIsARankedTensor(t) &&26 mlirTypeIsAInteger(mlirShapedTypeGetElementType(t));27}28 29NB_MODULE(_mlirPythonTestNanobind, m) {30 m.def(31 "register_python_test_dialect",32 [](MlirContext context, bool load) {33 MlirDialectHandle pythonTestDialect =34 mlirGetDialectHandle__python_test__();35 mlirDialectHandleRegisterDialect(pythonTestDialect, context);36 if (load) {37 mlirDialectHandleLoadDialect(pythonTestDialect, context);38 }39 },40 nb::arg("context"), nb::arg("load") = true,41 // clang-format off42 nb::sig("def register_python_test_dialect(context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") ", load: bool = True) -> None"));43 // clang-format on44 45 m.def(46 "register_dialect",47 [](MlirDialectRegistry registry) {48 MlirDialectHandle pythonTestDialect =49 mlirGetDialectHandle__python_test__();50 mlirDialectHandleInsertDialect(pythonTestDialect, registry);51 },52 nb::arg("registry"),53 // clang-format off54 nb::sig("def register_dialect(registry: " MAKE_MLIR_PYTHON_QUALNAME("ir.DialectRegistry") ") -> None"));55 // clang-format on56 57 m.def(58 "test_diagnostics_with_errors_and_notes",59 [](MlirContext ctx) {60 mlir::python::CollectDiagnosticsToStringScope handler(ctx);61 mlirPythonTestEmitDiagnosticWithNote(ctx);62 throw nb::value_error(handler.takeMessage().c_str());63 },64 // clang-format off65 nb::sig("def test_diagnostics_with_errors_and_notes(arg: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") ", /) -> None"));66 // clang-format on67 68 mlir_attribute_subclass(m, "TestAttr",69 mlirAttributeIsAPythonTestTestAttribute,70 mlirPythonTestTestAttributeGetTypeID)71 .def_classmethod(72 "get",73 [](const nb::object &cls, MlirContext ctx) {74 return cls(mlirPythonTestTestAttributeGet(ctx));75 },76 // clang-format off77 nb::sig("def get(cls: object, context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> object"),78 // clang-format on79 nb::arg("cls"), nb::arg("context").none() = nb::none());80 81 mlir_type_subclass(m, "TestType", mlirTypeIsAPythonTestTestType,82 mlirPythonTestTestTypeGetTypeID)83 .def_classmethod(84 "get",85 [](const nb::object &cls, MlirContext ctx) {86 return cls(mlirPythonTestTestTypeGet(ctx));87 },88 // clang-format off89 nb::sig("def get(cls: object, context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> object"),90 // clang-format on91 nb::arg("cls"), nb::arg("context").none() = nb::none());92 93 auto typeCls =94 mlir_type_subclass(m, "TestIntegerRankedTensorType",95 mlirTypeIsARankedIntegerTensor,96 nb::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))97 .attr("RankedTensorType"))98 .def_classmethod(99 "get",100 [](const nb::object &cls, std::vector<int64_t> shape,101 unsigned width, MlirContext ctx) {102 MlirAttribute encoding = mlirAttributeGetNull();103 return cls(mlirRankedTensorTypeGet(104 shape.size(), shape.data(), mlirIntegerTypeGet(ctx, width),105 encoding));106 },107 // clang-format off108 nb::sig("def get(cls: object, shape: collections.abc.Sequence[int], width: int, context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> object"),109 // clang-format on110 nb::arg("cls"), nb::arg("shape"), nb::arg("width"),111 nb::arg("context").none() = nb::none());112 113 assert(nb::hasattr(typeCls.get_class(), "static_typeid") &&114 "TestIntegerRankedTensorType has no static_typeid");115 116 MlirTypeID mlirRankedTensorTypeID = mlirRankedTensorTypeGetTypeID();117 118 nb::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))119 .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(120 mlirRankedTensorTypeID, nb::arg("replace") = true)(121 nanobind::cpp_function([typeCls](const nb::object &mlirType) {122 return typeCls.get_class()(mlirType);123 }));124 125 auto valueCls = mlir_value_subclass(m, "TestTensorValue",126 mlirTypeIsAPythonTestTestTensorValue)127 .def("is_null", [](MlirValue &self) {128 return mlirValueIsNull(self);129 });130 131 nb::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))132 .attr(MLIR_PYTHON_CAPI_VALUE_CASTER_REGISTER_ATTR)(133 mlirRankedTensorTypeID)(134 nanobind::cpp_function([valueCls](const nb::object &valueObj) {135 std::optional<nb::object> capsule =136 mlirApiObjectToCapsule(valueObj);137 assert(capsule.has_value() && "capsule is not null");138 MlirValue v = mlirPythonCapsuleToValue(capsule.value().ptr());139 MlirType t = mlirValueGetType(v);140 // This is hyper-specific in order to exercise/test registering a141 // value caster from cpp (but only for a single test case; see142 // testTensorValue python_test.py).143 if (mlirShapedTypeHasStaticShape(t) &&144 mlirShapedTypeGetDimSize(t, 0) == 1 &&145 mlirShapedTypeGetDimSize(t, 1) == 2 &&146 mlirShapedTypeGetDimSize(t, 2) == 3)147 return valueCls.get_class()(valueObj);148 return valueObj;149 }));150}151