1 //===- MainModule.cpp - Main pybind module --------------------------------===// 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-exception 6 // 7 //===----------------------------------------------------------------------===// 8 9 #include <tuple> 10 11 #include "PybindUtils.h" 12 13 #include "Globals.h" 14 #include "IRModule.h" 15 #include "Pass.h" 16 17 namespace py = pybind11; 18 using namespace mlir; 19 using namespace py::literals; 20 using namespace mlir::python; 21 22 // ----------------------------------------------------------------------------- 23 // Module initialization. 24 // ----------------------------------------------------------------------------- 25 26 PYBIND11_MODULE(_mlir, m) { 27 m.doc() = "MLIR Python Native Extension"; 28 29 py::class_<PyGlobals>(m, "_Globals", py::module_local()) 30 .def_property("dialect_search_modules", 31 &PyGlobals::getDialectSearchPrefixes, 32 &PyGlobals::setDialectSearchPrefixes) 33 .def( 34 "append_dialect_search_prefix", 35 [](PyGlobals &self, std::string moduleName) { 36 self.getDialectSearchPrefixes().push_back(std::move(moduleName)); 37 self.clearImportCache(); 38 }, 39 "module_name"_a) 40 .def("_register_dialect_impl", &PyGlobals::registerDialectImpl, 41 "dialect_namespace"_a, "dialect_class"_a, 42 "Testing hook for directly registering a dialect") 43 .def("_register_operation_impl", &PyGlobals::registerOperationImpl, 44 "operation_name"_a, "operation_class"_a, 45 "Testing hook for directly registering an operation"); 46 47 // Aside from making the globals accessible to python, having python manage 48 // it is necessary to make sure it is destroyed (and releases its python 49 // resources) properly. 50 m.attr("globals") = 51 py::cast(new PyGlobals, py::return_value_policy::take_ownership); 52 53 // Registration decorators. 54 m.def( 55 "register_dialect", 56 [](py::object pyClass) { 57 std::string dialectNamespace = 58 pyClass.attr("DIALECT_NAMESPACE").cast<std::string>(); 59 PyGlobals::get().registerDialectImpl(dialectNamespace, pyClass); 60 return pyClass; 61 }, 62 "dialect_class"_a, 63 "Class decorator for registering a custom Dialect wrapper"); 64 m.def( 65 "register_operation", 66 [](const py::object &dialectClass) -> py::cpp_function { 67 return py::cpp_function( 68 [dialectClass](py::object opClass) -> py::object { 69 std::string operationName = 70 opClass.attr("OPERATION_NAME").cast<std::string>(); 71 PyGlobals::get().registerOperationImpl(operationName, opClass); 72 73 // Dict-stuff the new opClass by name onto the dialect class. 74 py::object opClassName = opClass.attr("__name__"); 75 dialectClass.attr(opClassName) = opClass; 76 return opClass; 77 }); 78 }, 79 "dialect_class"_a, 80 "Produce a class decorator for registering an Operation class as part of " 81 "a dialect"); 82 m.def( 83 MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR, 84 [](MlirTypeID mlirTypeID, py::function typeCaster, bool replace) { 85 PyGlobals::get().registerTypeCaster(mlirTypeID, std::move(typeCaster), 86 replace); 87 }, 88 "typeid"_a, "type_caster"_a, "replace"_a = false, 89 "Register a type caster for casting MLIR types to custom user types."); 90 91 // Define and populate IR submodule. 92 auto irModule = m.def_submodule("ir", "MLIR IR Bindings"); 93 populateIRCore(irModule); 94 populateIRAffine(irModule); 95 populateIRAttributes(irModule); 96 populateIRInterfaces(irModule); 97 populateIRTypes(irModule); 98 99 // Define and populate PassManager submodule. 100 auto passModule = 101 m.def_submodule("passmanager", "MLIR Pass Management Bindings"); 102 populatePassManagerSubmodule(passModule); 103 } 104