MLIR  20.0.0git
Globals.h
Go to the documentation of this file.
1 //===- Globals.h - MLIR Python extension globals --------------------------===//
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 #ifndef MLIR_BINDINGS_PYTHON_GLOBALS_H
10 #define MLIR_BINDINGS_PYTHON_GLOBALS_H
11 
12 #include <optional>
13 #include <string>
14 #include <vector>
15 
16 #include "NanobindUtils.h"
17 #include "mlir-c/IR.h"
18 #include "mlir/CAPI/Support.h"
19 #include "llvm/ADT/DenseMap.h"
20 #include "llvm/ADT/StringRef.h"
21 #include "llvm/ADT/StringSet.h"
22 
23 namespace mlir {
24 namespace python {
25 
26 /// Globals that are always accessible once the extension has been initialized.
27 class PyGlobals {
28 public:
29  PyGlobals();
30  ~PyGlobals();
31 
32  /// Most code should get the globals via this static accessor.
33  static PyGlobals &get() {
34  assert(instance && "PyGlobals is null");
35  return *instance;
36  }
37 
38  /// Get and set the list of parent modules to search for dialect
39  /// implementation classes.
40  std::vector<std::string> &getDialectSearchPrefixes() {
41  return dialectSearchPrefixes;
42  }
43  void setDialectSearchPrefixes(std::vector<std::string> newValues) {
44  dialectSearchPrefixes.swap(newValues);
45  }
46 
47  /// Loads a python module corresponding to the given dialect namespace.
48  /// No-ops if the module has already been loaded or is not found. Raises
49  /// an error on any evaluation issues.
50  /// Note that this returns void because it is expected that the module
51  /// contains calls to decorators and helpers that register the salient
52  /// entities. Returns true if dialect is successfully loaded.
53  bool loadDialectModule(llvm::StringRef dialectNamespace);
54 
55  /// Adds a user-friendly Attribute builder.
56  /// Raises an exception if the mapping already exists and replace == false.
57  /// This is intended to be called by implementation code.
58  void registerAttributeBuilder(const std::string &attributeKind,
59  nanobind::callable pyFunc,
60  bool replace = false);
61 
62  /// Adds a user-friendly type caster. Raises an exception if the mapping
63  /// already exists and replace == false. This is intended to be called by
64  /// implementation code.
65  void registerTypeCaster(MlirTypeID mlirTypeID, nanobind::callable typeCaster,
66  bool replace = false);
67 
68  /// Adds a user-friendly value caster. Raises an exception if the mapping
69  /// already exists and replace == false. This is intended to be called by
70  /// implementation code.
71  void registerValueCaster(MlirTypeID mlirTypeID,
72  nanobind::callable valueCaster,
73  bool replace = false);
74 
75  /// Adds a concrete implementation dialect class.
76  /// Raises an exception if the mapping already exists.
77  /// This is intended to be called by implementation code.
78  void registerDialectImpl(const std::string &dialectNamespace,
79  nanobind::object pyClass);
80 
81  /// Adds a concrete implementation operation class.
82  /// Raises an exception if the mapping already exists and replace == false.
83  /// This is intended to be called by implementation code.
84  void registerOperationImpl(const std::string &operationName,
85  nanobind::object pyClass, bool replace = false);
86 
87  /// Returns the custom Attribute builder for Attribute kind.
88  std::optional<nanobind::callable>
89  lookupAttributeBuilder(const std::string &attributeKind);
90 
91  /// Returns the custom type caster for MlirTypeID mlirTypeID.
92  std::optional<nanobind::callable> lookupTypeCaster(MlirTypeID mlirTypeID,
93  MlirDialect dialect);
94 
95  /// Returns the custom value caster for MlirTypeID mlirTypeID.
96  std::optional<nanobind::callable> lookupValueCaster(MlirTypeID mlirTypeID,
97  MlirDialect dialect);
98 
99  /// Looks up a registered dialect class by namespace. Note that this may
100  /// trigger loading of the defining module and can arbitrarily re-enter.
101  std::optional<nanobind::object>
102  lookupDialectClass(const std::string &dialectNamespace);
103 
104  /// Looks up a registered operation class (deriving from OpView) by operation
105  /// name. Note that this may trigger a load of the dialect, which can
106  /// arbitrarily re-enter.
107  std::optional<nanobind::object>
108  lookupOperationClass(llvm::StringRef operationName);
109 
110 private:
111  static PyGlobals *instance;
112  /// Module name prefixes to search under for dialect implementation modules.
113  std::vector<std::string> dialectSearchPrefixes;
114  /// Map of dialect namespace to external dialect class object.
115  llvm::StringMap<nanobind::object> dialectClassMap;
116  /// Map of full operation name to external operation class object.
117  llvm::StringMap<nanobind::object> operationClassMap;
118  /// Map of attribute ODS name to custom builder.
119  llvm::StringMap<nanobind::callable> attributeBuilderMap;
120  /// Map of MlirTypeID to custom type caster.
122  /// Map of MlirTypeID to custom value caster.
124  /// Set of dialect namespaces that we have attempted to import implementation
125  /// modules for.
126  llvm::StringSet<> loadedDialectModules;
127 };
128 
129 } // namespace python
130 } // namespace mlir
131 
132 #endif // MLIR_BINDINGS_PYTHON_GLOBALS_H
Globals that are always accessible once the extension has been initialized.
Definition: Globals.h:27
bool loadDialectModule(llvm::StringRef dialectNamespace)
Loads a python module corresponding to the given dialect namespace.
Definition: IRModule.cpp:40
void registerTypeCaster(MlirTypeID mlirTypeID, nanobind::callable typeCaster, bool replace=false)
Adds a user-friendly type caster.
Definition: IRModule.cpp:82
void registerOperationImpl(const std::string &operationName, nanobind::object pyClass, bool replace=false)
Adds a concrete implementation operation class.
Definition: IRModule.cpp:111
void registerAttributeBuilder(const std::string &attributeKind, nanobind::callable pyFunc, bool replace=false)
Adds a user-friendly Attribute builder.
Definition: IRModule.cpp:69
void registerValueCaster(MlirTypeID mlirTypeID, nanobind::callable valueCaster, bool replace=false)
Adds a user-friendly value caster.
Definition: IRModule.cpp:91
std::optional< nanobind::callable > lookupValueCaster(MlirTypeID mlirTypeID, MlirDialect dialect)
Returns the custom value caster for MlirTypeID mlirTypeID.
Definition: IRModule.cpp:144
void setDialectSearchPrefixes(std::vector< std::string > newValues)
Definition: Globals.h:43
std::optional< nanobind::object > lookupDialectClass(const std::string &dialectNamespace)
Looks up a registered dialect class by namespace.
Definition: IRModule.cpp:157
std::vector< std::string > & getDialectSearchPrefixes()
Get and set the list of parent modules to search for dialect implementation classes.
Definition: Globals.h:40
std::optional< nanobind::object > lookupOperationClass(llvm::StringRef operationName)
Looks up a registered operation class (deriving from OpView) by operation name.
Definition: IRModule.cpp:171
static PyGlobals & get()
Most code should get the globals via this static accessor.
Definition: Globals.h:33
std::optional< nanobind::callable > lookupAttributeBuilder(const std::string &attributeKind)
Returns the custom Attribute builder for Attribute kind.
Definition: IRModule.cpp:123
void registerDialectImpl(const std::string &dialectNamespace, nanobind::object pyClass)
Adds a concrete implementation dialect class.
Definition: IRModule.cpp:100
std::optional< nanobind::callable > lookupTypeCaster(MlirTypeID mlirTypeID, MlirDialect dialect)
Returns the custom type caster for MlirTypeID mlirTypeID.
Definition: IRModule.cpp:132
Include the generated interface declarations.