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 <vector> 14 15 #include "mlir-c/Bindings/Python/Interop.h" 16 17 namespace py = pybind11; 18 using namespace mlir; 19 using namespace mlir::python; 20 21 // ----------------------------------------------------------------------------- 22 // PyGlobals 23 // ----------------------------------------------------------------------------- 24 25 PyGlobals *PyGlobals::instance = nullptr; 26 27 PyGlobals::PyGlobals() { 28 assert(!instance && "PyGlobals already constructed"); 29 instance = this; 30 // The default search path include {mlir.}dialects, where {mlir.} is the 31 // package prefix configured at compile time. 32 dialectSearchPrefixes.emplace_back(MAKE_MLIR_PYTHON_QUALNAME("dialects")); 33 } 34 35 PyGlobals::~PyGlobals() { instance = nullptr; } 36 37 void PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) { 38 py::gil_scoped_acquire(); 39 if (loadedDialectModulesCache.contains(dialectNamespace)) 40 return; 41 // Since re-entrancy is possible, make a copy of the search prefixes. 42 std::vector<std::string> localSearchPrefixes = dialectSearchPrefixes; 43 py::object loaded; 44 for (std::string moduleName : localSearchPrefixes) { 45 moduleName.push_back('.'); 46 moduleName.append(dialectNamespace.data(), dialectNamespace.size()); 47 48 try { 49 py::gil_scoped_release(); 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 // Note: Iterator cannot be shared from prior to loading, since re-entrancy 61 // may have occurred, which may do anything. 62 loadedDialectModulesCache.insert(dialectNamespace); 63 } 64 65 void PyGlobals::registerDialectImpl(const std::string &dialectNamespace, 66 py::object pyClass) { 67 py::gil_scoped_acquire(); 68 py::object &found = dialectClassMap[dialectNamespace]; 69 if (found) { 70 throw SetPyError(PyExc_RuntimeError, llvm::Twine("Dialect namespace '") + 71 dialectNamespace + 72 "' is already registered."); 73 } 74 found = std::move(pyClass); 75 } 76 77 void PyGlobals::registerOperationImpl(const std::string &operationName, 78 py::object pyClass, 79 py::object rawOpViewClass) { 80 py::gil_scoped_acquire(); 81 py::object &found = operationClassMap[operationName]; 82 if (found) { 83 throw SetPyError(PyExc_RuntimeError, llvm::Twine("Operation '") + 84 operationName + 85 "' is already registered."); 86 } 87 found = std::move(pyClass); 88 rawOpViewClassMap[operationName] = std::move(rawOpViewClass); 89 } 90 91 llvm::Optional<py::object> 92 PyGlobals::lookupDialectClass(const std::string &dialectNamespace) { 93 py::gil_scoped_acquire(); 94 loadDialectModule(dialectNamespace); 95 // Fast match against the class map first (common case). 96 const auto foundIt = dialectClassMap.find(dialectNamespace); 97 if (foundIt != dialectClassMap.end()) { 98 if (foundIt->second.is_none()) 99 return llvm::None; 100 assert(foundIt->second && "py::object is defined"); 101 return foundIt->second; 102 } 103 104 // Not found and loading did not yield a registration. Negative cache. 105 dialectClassMap[dialectNamespace] = py::none(); 106 return llvm::None; 107 } 108 109 llvm::Optional<pybind11::object> 110 PyGlobals::lookupRawOpViewClass(llvm::StringRef operationName) { 111 { 112 py::gil_scoped_acquire(); 113 auto foundIt = rawOpViewClassMapCache.find(operationName); 114 if (foundIt != rawOpViewClassMapCache.end()) { 115 if (foundIt->second.is_none()) 116 return llvm::None; 117 assert(foundIt->second && "py::object is defined"); 118 return foundIt->second; 119 } 120 } 121 122 // Not found. Load the dialect namespace. 123 auto split = operationName.split('.'); 124 llvm::StringRef dialectNamespace = split.first; 125 loadDialectModule(dialectNamespace); 126 127 // Attempt to find from the canonical map and cache. 128 { 129 py::gil_scoped_acquire(); 130 auto foundIt = rawOpViewClassMap.find(operationName); 131 if (foundIt != rawOpViewClassMap.end()) { 132 if (foundIt->second.is_none()) 133 return llvm::None; 134 assert(foundIt->second && "py::object is defined"); 135 // Positive cache. 136 rawOpViewClassMapCache[operationName] = foundIt->second; 137 return foundIt->second; 138 } 139 // Negative cache. 140 rawOpViewClassMap[operationName] = py::none(); 141 return llvm::None; 142 } 143 } 144 145 void PyGlobals::clearImportCache() { 146 py::gil_scoped_acquire(); 147 loadedDialectModulesCache.clear(); 148 rawOpViewClassMapCache.clear(); 149 } 150