1 //===- IRModule.cpp - IR 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 "IRModule.h" 10 #include "Globals.h" 11 #include "PybindUtils.h" 12 13 #include "mlir-c/Bindings/Python/Interop.h" 14 #include "mlir-c/Support.h" 15 16 #include <optional> 17 #include <vector> 18 19 namespace py = pybind11; 20 using namespace mlir; 21 using namespace mlir::python; 22 23 // ----------------------------------------------------------------------------- 24 // PyGlobals 25 // ----------------------------------------------------------------------------- 26 27 PyGlobals *PyGlobals::instance = nullptr; 28 29 PyGlobals::PyGlobals() { 30 assert(!instance && "PyGlobals already constructed"); 31 instance = this; 32 // The default search path include {mlir.}dialects, where {mlir.} is the 33 // package prefix configured at compile time. 34 dialectSearchPrefixes.emplace_back(MAKE_MLIR_PYTHON_QUALNAME("dialects")); 35 } 36 37 PyGlobals::~PyGlobals() { instance = nullptr; } 38 39 bool PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) { 40 if (loadedDialectModules.contains(dialectNamespace)) 41 return true; 42 // Since re-entrancy is possible, make a copy of the search prefixes. 43 std::vector<std::string> localSearchPrefixes = dialectSearchPrefixes; 44 py::object loaded = py::none(); 45 for (std::string moduleName : localSearchPrefixes) { 46 moduleName.push_back('.'); 47 moduleName.append(dialectNamespace.data(), dialectNamespace.size()); 48 49 try { 50 loaded = py::module::import(moduleName.c_str()); 51 } catch (py::error_already_set &e) { 52 if (e.matches(PyExc_ModuleNotFoundError)) { 53 continue; 54 } 55 throw; 56 } 57 break; 58 } 59 60 if (loaded.is_none()) 61 return false; 62 // Note: Iterator cannot be shared from prior to loading, since re-entrancy 63 // may have occurred, which may do anything. 64 loadedDialectModules.insert(dialectNamespace); 65 return true; 66 } 67 68 void PyGlobals::registerAttributeBuilder(const std::string &attributeKind, 69 py::function pyFunc, bool replace) { 70 py::object &found = attributeBuilderMap[attributeKind]; 71 if (found && !replace) { 72 throw std::runtime_error((llvm::Twine("Attribute builder for '") + 73 attributeKind + 74 "' is already registered with func: " + 75 py::str(found).operator std::string()) 76 .str()); 77 } 78 found = std::move(pyFunc); 79 } 80 81 void PyGlobals::registerTypeCaster(MlirTypeID mlirTypeID, 82 pybind11::function typeCaster, 83 bool replace) { 84 pybind11::object &found = typeCasterMap[mlirTypeID]; 85 if (found && !replace) 86 throw std::runtime_error("Type caster is already registered with caster: " + 87 py::str(found).operator std::string()); 88 found = std::move(typeCaster); 89 } 90 91 void PyGlobals::registerDialectImpl(const std::string &dialectNamespace, 92 py::object pyClass) { 93 py::object &found = dialectClassMap[dialectNamespace]; 94 if (found) { 95 throw std::runtime_error((llvm::Twine("Dialect namespace '") + 96 dialectNamespace + "' is already registered.") 97 .str()); 98 } 99 found = std::move(pyClass); 100 } 101 102 void PyGlobals::registerOperationImpl(const std::string &operationName, 103 py::object pyClass, bool replace) { 104 py::object &found = operationClassMap[operationName]; 105 if (found && !replace) { 106 throw std::runtime_error((llvm::Twine("Operation '") + operationName + 107 "' is already registered.") 108 .str()); 109 } 110 found = std::move(pyClass); 111 } 112 113 std::optional<py::function> 114 PyGlobals::lookupAttributeBuilder(const std::string &attributeKind) { 115 const auto foundIt = attributeBuilderMap.find(attributeKind); 116 if (foundIt != attributeBuilderMap.end()) { 117 assert(foundIt->second && "attribute builder is defined"); 118 return foundIt->second; 119 } 120 return std::nullopt; 121 } 122 123 std::optional<py::function> PyGlobals::lookupTypeCaster(MlirTypeID mlirTypeID, 124 MlirDialect dialect) { 125 // Make sure dialect module is loaded. 126 if (!loadDialectModule(unwrap(mlirDialectGetNamespace(dialect)))) 127 return std::nullopt; 128 129 const auto foundIt = typeCasterMap.find(mlirTypeID); 130 if (foundIt != typeCasterMap.end()) { 131 assert(foundIt->second && "type caster is defined"); 132 return foundIt->second; 133 } 134 return std::nullopt; 135 } 136 137 std::optional<py::object> 138 PyGlobals::lookupDialectClass(const std::string &dialectNamespace) { 139 // Make sure dialect module is loaded. 140 if (!loadDialectModule(dialectNamespace)) 141 return std::nullopt; 142 const auto foundIt = dialectClassMap.find(dialectNamespace); 143 if (foundIt != dialectClassMap.end()) { 144 assert(foundIt->second && "dialect class is defined"); 145 return foundIt->second; 146 } 147 // Not found and loading did not yield a registration. 148 return std::nullopt; 149 } 150 151 std::optional<pybind11::object> 152 PyGlobals::lookupOperationClass(llvm::StringRef operationName) { 153 // Make sure dialect module is loaded. 154 auto split = operationName.split('.'); 155 llvm::StringRef dialectNamespace = split.first; 156 if (!loadDialectModule(dialectNamespace)) 157 return std::nullopt; 158 159 auto foundIt = operationClassMap.find(operationName); 160 if (foundIt != operationClassMap.end()) { 161 assert(foundIt->second && "OpView is defined"); 162 return foundIt->second; 163 } 164 // Not found and loading did not yield a registration. 165 return std::nullopt; 166 } 167