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 "Dialects.h" 14 #include "Globals.h" 15 #include "IRModule.h" 16 #include "Pass.h" 17 18 namespace py = pybind11; 19 using namespace mlir; 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 py::arg("module_name")) 40 .def("_register_dialect_impl", &PyGlobals::registerDialectImpl, 41 py::arg("dialect_namespace"), py::arg("dialect_class"), 42 "Testing hook for directly registering a dialect") 43 .def("_register_operation_impl", &PyGlobals::registerOperationImpl, 44 py::arg("operation_name"), py::arg("operation_class"), 45 py::arg("raw_opview_class"), 46 "Testing hook for directly registering an operation"); 47 48 // Aside from making the globals accessible to python, having python manage 49 // it is necessary to make sure it is destroyed (and releases its python 50 // resources) properly. 51 m.attr("globals") = 52 py::cast(new PyGlobals, py::return_value_policy::take_ownership); 53 54 // Registration decorators. 55 m.def( 56 "register_dialect", 57 [](py::object pyClass) { 58 std::string dialectNamespace = 59 pyClass.attr("DIALECT_NAMESPACE").cast<std::string>(); 60 PyGlobals::get().registerDialectImpl(dialectNamespace, pyClass); 61 return pyClass; 62 }, 63 py::arg("dialect_class"), 64 "Class decorator for registering a custom Dialect wrapper"); 65 m.def( 66 "register_operation", 67 [](py::object dialectClass) -> py::cpp_function { 68 return py::cpp_function( 69 [dialectClass](py::object opClass) -> py::object { 70 std::string operationName = 71 opClass.attr("OPERATION_NAME").cast<std::string>(); 72 auto rawSubclass = PyOpView::createRawSubclass(opClass); 73 PyGlobals::get().registerOperationImpl(operationName, opClass, 74 rawSubclass); 75 76 // Dict-stuff the new opClass by name onto the dialect class. 77 py::object opClassName = opClass.attr("__name__"); 78 dialectClass.attr(opClassName) = opClass; 79 80 // Now create a special "Raw" subclass that passes through 81 // construction to the OpView parent (bypasses the intermediate 82 // child's __init__). 83 opClass.attr("_Raw") = rawSubclass; 84 return opClass; 85 }); 86 }, 87 py::arg("dialect_class"), 88 "Produce a class decorator for registering an Operation class as part of " 89 "a dialect"); 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 // Define and populate dialect submodules. 105 auto dialectsModule = m.def_submodule("dialects"); 106 auto linalgModule = dialectsModule.def_submodule("linalg"); 107 populateDialectLinalgSubmodule(linalgModule); 108 populateDialectSparseTensorSubmodule( 109 dialectsModule.def_submodule("sparse_tensor"), irModule); 110 } 111