MLIR  20.0.0git
IRModule.cpp
Go to the documentation of this file.
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 
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 
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::registerValueCaster(MlirTypeID mlirTypeID,
92  pybind11::function valueCaster,
93  bool replace) {
94  pybind11::object &found = valueCasterMap[mlirTypeID];
95  if (found && !replace)
96  throw std::runtime_error("Value caster is already registered: " +
97  py::repr(found).cast<std::string>());
98  found = std::move(valueCaster);
99 }
100 
101 void PyGlobals::registerDialectImpl(const std::string &dialectNamespace,
102  py::object pyClass) {
103  py::object &found = dialectClassMap[dialectNamespace];
104  if (found) {
105  throw std::runtime_error((llvm::Twine("Dialect namespace '") +
106  dialectNamespace + "' is already registered.")
107  .str());
108  }
109  found = std::move(pyClass);
110 }
111 
112 void PyGlobals::registerOperationImpl(const std::string &operationName,
113  py::object pyClass, bool replace) {
114  py::object &found = operationClassMap[operationName];
115  if (found && !replace) {
116  throw std::runtime_error((llvm::Twine("Operation '") + operationName +
117  "' is already registered.")
118  .str());
119  }
120  found = std::move(pyClass);
121 }
122 
123 std::optional<py::function>
124 PyGlobals::lookupAttributeBuilder(const std::string &attributeKind) {
125  const auto foundIt = attributeBuilderMap.find(attributeKind);
126  if (foundIt != attributeBuilderMap.end()) {
127  assert(foundIt->second && "attribute builder is defined");
128  return foundIt->second;
129  }
130  return std::nullopt;
131 }
132 
133 std::optional<py::function> PyGlobals::lookupTypeCaster(MlirTypeID mlirTypeID,
134  MlirDialect dialect) {
135  // Try to load dialect module.
137  const auto foundIt = typeCasterMap.find(mlirTypeID);
138  if (foundIt != typeCasterMap.end()) {
139  assert(foundIt->second && "type caster is defined");
140  return foundIt->second;
141  }
142  return std::nullopt;
143 }
144 
145 std::optional<py::function> PyGlobals::lookupValueCaster(MlirTypeID mlirTypeID,
146  MlirDialect dialect) {
147  // Try to load dialect module.
149  const auto foundIt = valueCasterMap.find(mlirTypeID);
150  if (foundIt != valueCasterMap.end()) {
151  assert(foundIt->second && "value caster is defined");
152  return foundIt->second;
153  }
154  return std::nullopt;
155 }
156 
157 std::optional<py::object>
158 PyGlobals::lookupDialectClass(const std::string &dialectNamespace) {
159  // Make sure dialect module is loaded.
160  if (!loadDialectModule(dialectNamespace))
161  return std::nullopt;
162  const auto foundIt = dialectClassMap.find(dialectNamespace);
163  if (foundIt != dialectClassMap.end()) {
164  assert(foundIt->second && "dialect class is defined");
165  return foundIt->second;
166  }
167  // Not found and loading did not yield a registration.
168  return std::nullopt;
169 }
170 
171 std::optional<pybind11::object>
172 PyGlobals::lookupOperationClass(llvm::StringRef operationName) {
173  // Make sure dialect module is loaded.
174  auto split = operationName.split('.');
175  llvm::StringRef dialectNamespace = split.first;
176  if (!loadDialectModule(dialectNamespace))
177  return std::nullopt;
178 
179  auto foundIt = operationClassMap.find(operationName);
180  if (foundIt != operationClassMap.end()) {
181  assert(foundIt->second && "OpView is defined");
182  return foundIt->second;
183  }
184  // Not found and loading did not yield a registration.
185  return std::nullopt;
186 }
#define MAKE_MLIR_PYTHON_QUALNAME(local)
Definition: Interop.h:57
Globals that are always accessible once the extension has been initialized.
Definition: Globals.h:28
bool loadDialectModule(llvm::StringRef dialectNamespace)
Loads a python module corresponding to the given dialect namespace.
Definition: IRModule.cpp:39
std::optional< pybind11::function > lookupAttributeBuilder(const std::string &attributeKind)
Returns the custom Attribute builder for Attribute kind.
Definition: IRModule.cpp:124
std::optional< pybind11::function > lookupTypeCaster(MlirTypeID mlirTypeID, MlirDialect dialect)
Returns the custom type caster for MlirTypeID mlirTypeID.
Definition: IRModule.cpp:133
void registerAttributeBuilder(const std::string &attributeKind, pybind11::function pyFunc, bool replace=false)
Adds a user-friendly Attribute builder.
Definition: IRModule.cpp:68
std::optional< pybind11::object > lookupDialectClass(const std::string &dialectNamespace)
Looks up a registered dialect class by namespace.
Definition: IRModule.cpp:158
void registerTypeCaster(MlirTypeID mlirTypeID, pybind11::function typeCaster, bool replace=false)
Adds a user-friendly type caster.
Definition: IRModule.cpp:81
void registerValueCaster(MlirTypeID mlirTypeID, pybind11::function valueCaster, bool replace=false)
Adds a user-friendly value caster.
Definition: IRModule.cpp:91
void registerOperationImpl(const std::string &operationName, pybind11::object pyClass, bool replace=false)
Adds a concrete implementation operation class.
Definition: IRModule.cpp:112
void registerDialectImpl(const std::string &dialectNamespace, pybind11::object pyClass)
Adds a concrete implementation dialect class.
Definition: IRModule.cpp:101
std::optional< pybind11::object > lookupOperationClass(llvm::StringRef operationName)
Looks up a registered operation class (deriving from OpView) by operation name.
Definition: IRModule.cpp:172
std::optional< pybind11::function > lookupValueCaster(MlirTypeID mlirTypeID, MlirDialect dialect)
Returns the custom value caster for MlirTypeID mlirTypeID.
Definition: IRModule.cpp:145
mlir::Diagnostic & unwrap(MlirDiagnostic diagnostic)
Definition: Diagnostics.h:19
MLIR_CAPI_EXPORTED MlirStringRef mlirDialectGetNamespace(MlirDialect dialect)
Returns the namespace of the given dialect.
Definition: IR.cpp:127
Include the generated interface declarations.