MLIR  22.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 <regex>
14 #include <string>
15 #include <unordered_set>
16 #include <vector>
17 
18 #include "NanobindUtils.h"
19 #include "mlir-c/IR.h"
20 #include "mlir/CAPI/Support.h"
21 #include "llvm/ADT/DenseMap.h"
22 #include "llvm/ADT/StringExtras.h"
23 #include "llvm/ADT/StringRef.h"
24 #include "llvm/ADT/StringSet.h"
25 #include "llvm/Support/Regex.h"
26 
27 namespace mlir {
28 namespace python {
29 
30 /// Globals that are always accessible once the extension has been initialized.
31 /// Methods of this class are thread-safe.
32 class PyGlobals {
33 public:
34  PyGlobals();
35  ~PyGlobals();
36 
37  /// Most code should get the globals via this static accessor.
38  static PyGlobals &get() {
39  assert(instance && "PyGlobals is null");
40  return *instance;
41  }
42 
43  /// Get and set the list of parent modules to search for dialect
44  /// implementation classes.
45  std::vector<std::string> getDialectSearchPrefixes() {
46  nanobind::ft_lock_guard lock(mutex);
47  return dialectSearchPrefixes;
48  }
49  void setDialectSearchPrefixes(std::vector<std::string> newValues) {
50  nanobind::ft_lock_guard lock(mutex);
51  dialectSearchPrefixes.swap(newValues);
52  }
53  void addDialectSearchPrefix(std::string value) {
54  nanobind::ft_lock_guard lock(mutex);
55  dialectSearchPrefixes.push_back(std::move(value));
56  }
57 
58  /// Loads a python module corresponding to the given dialect namespace.
59  /// No-ops if the module has already been loaded or is not found. Raises
60  /// an error on any evaluation issues.
61  /// Note that this returns void because it is expected that the module
62  /// contains calls to decorators and helpers that register the salient
63  /// entities. Returns true if dialect is successfully loaded.
64  bool loadDialectModule(llvm::StringRef dialectNamespace);
65 
66  /// Adds a user-friendly Attribute builder.
67  /// Raises an exception if the mapping already exists and replace == false.
68  /// This is intended to be called by implementation code.
69  void registerAttributeBuilder(const std::string &attributeKind,
70  nanobind::callable pyFunc,
71  bool replace = false);
72 
73  /// Adds a user-friendly type caster. Raises an exception if the mapping
74  /// already exists and replace == false. This is intended to be called by
75  /// implementation code.
76  void registerTypeCaster(MlirTypeID mlirTypeID, nanobind::callable typeCaster,
77  bool replace = false);
78 
79  /// Adds a user-friendly value caster. Raises an exception if the mapping
80  /// already exists and replace == false. This is intended to be called by
81  /// implementation code.
82  void registerValueCaster(MlirTypeID mlirTypeID,
83  nanobind::callable valueCaster,
84  bool replace = false);
85 
86  /// Adds a concrete implementation dialect class.
87  /// Raises an exception if the mapping already exists.
88  /// This is intended to be called by implementation code.
89  void registerDialectImpl(const std::string &dialectNamespace,
90  nanobind::object pyClass);
91 
92  /// Adds a concrete implementation operation class.
93  /// Raises an exception if the mapping already exists and replace == false.
94  /// This is intended to be called by implementation code.
95  void registerOperationImpl(const std::string &operationName,
96  nanobind::object pyClass, bool replace = false);
97 
98  /// Returns the custom Attribute builder for Attribute kind.
99  std::optional<nanobind::callable>
100  lookupAttributeBuilder(const std::string &attributeKind);
101 
102  /// Returns the custom type caster for MlirTypeID mlirTypeID.
103  std::optional<nanobind::callable> lookupTypeCaster(MlirTypeID mlirTypeID,
104  MlirDialect dialect);
105 
106  /// Returns the custom value caster for MlirTypeID mlirTypeID.
107  std::optional<nanobind::callable> lookupValueCaster(MlirTypeID mlirTypeID,
108  MlirDialect dialect);
109 
110  /// Looks up a registered dialect class by namespace. Note that this may
111  /// trigger loading of the defining module and can arbitrarily re-enter.
112  std::optional<nanobind::object>
113  lookupDialectClass(const std::string &dialectNamespace);
114 
115  /// Looks up a registered operation class (deriving from OpView) by operation
116  /// name. Note that this may trigger a load of the dialect, which can
117  /// arbitrarily re-enter.
118  std::optional<nanobind::object>
119  lookupOperationClass(llvm::StringRef operationName);
120 
121  class TracebackLoc {
122  public:
123  bool locTracebacksEnabled();
124 
125  void setLocTracebacksEnabled(bool value);
126 
127  size_t locTracebackFramesLimit();
128 
129  void setLocTracebackFramesLimit(size_t value);
130 
131  void registerTracebackFileInclusion(const std::string &file);
132 
133  void registerTracebackFileExclusion(const std::string &file);
134 
135  bool isUserTracebackFilename(llvm::StringRef file);
136 
137  static constexpr size_t kMaxFrames = 512;
138 
139  private:
140  nanobind::ft_mutex mutex;
141  bool locTracebackEnabled_ = false;
142  size_t locTracebackFramesLimit_ = 10;
143  std::unordered_set<std::string> userTracebackIncludeFiles;
144  std::unordered_set<std::string> userTracebackExcludeFiles;
145  std::regex userTracebackIncludeRegex;
146  bool rebuildUserTracebackIncludeRegex = false;
147  std::regex userTracebackExcludeRegex;
148  bool rebuildUserTracebackExcludeRegex = false;
149  llvm::StringMap<bool> isUserTracebackFilenameCache;
150  };
151 
152  TracebackLoc &getTracebackLoc() { return tracebackLoc; }
153 
154 private:
155  static PyGlobals *instance;
156 
157  nanobind::ft_mutex mutex;
158 
159  /// Module name prefixes to search under for dialect implementation modules.
160  std::vector<std::string> dialectSearchPrefixes;
161  /// Map of dialect namespace to external dialect class object.
162  llvm::StringMap<nanobind::object> dialectClassMap;
163  /// Map of full operation name to external operation class object.
164  llvm::StringMap<nanobind::object> operationClassMap;
165  /// Map of attribute ODS name to custom builder.
166  llvm::StringMap<nanobind::callable> attributeBuilderMap;
167  /// Map of MlirTypeID to custom type caster.
169  /// Map of MlirTypeID to custom value caster.
171  /// Set of dialect namespaces that we have attempted to import implementation
172  /// modules for.
173  llvm::StringSet<> loadedDialectModules;
174 
175  TracebackLoc tracebackLoc;
176 };
177 
178 } // namespace python
179 } // namespace mlir
180 
181 #endif // MLIR_BINDINGS_PYTHON_GLOBALS_H
void registerTracebackFileExclusion(const std::string &file)
Definition: IRModule.cpp:233
static constexpr size_t kMaxFrames
Definition: Globals.h:137
bool isUserTracebackFilename(llvm::StringRef file)
Definition: IRModule.cpp:245
void setLocTracebacksEnabled(bool value)
Definition: IRModule.cpp:206
void setLocTracebackFramesLimit(size_t value)
Definition: IRModule.cpp:216
void registerTracebackFileInclusion(const std::string &file)
Definition: IRModule.cpp:221
Globals that are always accessible once the extension has been initialized.
Definition: Globals.h:32
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:87
std::vector< std::string > getDialectSearchPrefixes()
Get and set the list of parent modules to search for dialect implementation classes.
Definition: Globals.h:45
void addDialectSearchPrefix(std::string value)
Definition: Globals.h:53
void registerOperationImpl(const std::string &operationName, nanobind::object pyClass, bool replace=false)
Adds a concrete implementation operation class.
Definition: IRModule.cpp:119
void registerAttributeBuilder(const std::string &attributeKind, nanobind::callable pyFunc, bool replace=false)
Adds a user-friendly Attribute builder.
Definition: IRModule.cpp:73
void registerValueCaster(MlirTypeID mlirTypeID, nanobind::callable valueCaster, bool replace=false)
Adds a user-friendly value caster.
Definition: IRModule.cpp:97
std::optional< nanobind::callable > lookupValueCaster(MlirTypeID mlirTypeID, MlirDialect dialect)
Returns the custom value caster for MlirTypeID mlirTypeID.
Definition: IRModule.cpp:155
void setDialectSearchPrefixes(std::vector< std::string > newValues)
Definition: Globals.h:49
TracebackLoc & getTracebackLoc()
Definition: Globals.h:152
std::optional< nanobind::object > lookupDialectClass(const std::string &dialectNamespace)
Looks up a registered dialect class by namespace.
Definition: IRModule.cpp:169
std::optional< nanobind::object > lookupOperationClass(llvm::StringRef operationName)
Looks up a registered operation class (deriving from OpView) by operation name.
Definition: IRModule.cpp:184
static PyGlobals & get()
Most code should get the globals via this static accessor.
Definition: Globals.h:38
std::optional< nanobind::callable > lookupAttributeBuilder(const std::string &attributeKind)
Returns the custom Attribute builder for Attribute kind.
Definition: IRModule.cpp:132
void registerDialectImpl(const std::string &dialectNamespace, nanobind::object pyClass)
Adds a concrete implementation dialect class.
Definition: IRModule.cpp:107
std::optional< nanobind::callable > lookupTypeCaster(MlirTypeID mlirTypeID, MlirDialect dialect)
Returns the custom type caster for MlirTypeID mlirTypeID.
Definition: IRModule.cpp:142
Include the generated interface declarations.