MLIR  16.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 <string>
13 #include <vector>
14 
15 #include "PybindUtils.h"
16 
17 #include "llvm/ADT/Optional.h"
18 #include "llvm/ADT/StringRef.h"
19 #include "llvm/ADT/StringSet.h"
20 
21 namespace mlir {
22 namespace python {
23 
24 /// Globals that are always accessible once the extension has been initialized.
25 class PyGlobals {
26 public:
27  PyGlobals();
28  ~PyGlobals();
29 
30  /// Most code should get the globals via this static accessor.
31  static PyGlobals &get() {
32  assert(instance && "PyGlobals is null");
33  return *instance;
34  }
35 
36  /// Get and set the list of parent modules to search for dialect
37  /// implementation classes.
38  std::vector<std::string> &getDialectSearchPrefixes() {
39  return dialectSearchPrefixes;
40  }
41  void setDialectSearchPrefixes(std::vector<std::string> newValues) {
42  dialectSearchPrefixes.swap(newValues);
43  }
44 
45  /// Clears positive and negative caches regarding what implementations are
46  /// available. Future lookups will do more expensive existence checks.
47  void clearImportCache();
48 
49  /// Loads a python module corresponding to the given dialect namespace.
50  /// No-ops if the module has already been loaded or is not found. Raises
51  /// an error on any evaluation issues.
52  /// Note that this returns void because it is expected that the module
53  /// contains calls to decorators and helpers that register the salient
54  /// entities.
55  void loadDialectModule(llvm::StringRef dialectNamespace);
56 
57  /// Decorator for registering a custom Dialect class. The class object must
58  /// have a DIALECT_NAMESPACE attribute.
59  pybind11::object registerDialectDecorator(pybind11::object pyClass);
60 
61  /// Adds a concrete implementation dialect class.
62  /// Raises an exception if the mapping already exists.
63  /// This is intended to be called by implementation code.
64  void registerDialectImpl(const std::string &dialectNamespace,
65  pybind11::object pyClass);
66 
67  /// Adds a concrete implementation operation class.
68  /// Raises an exception if the mapping already exists.
69  /// This is intended to be called by implementation code.
70  void registerOperationImpl(const std::string &operationName,
71  pybind11::object pyClass,
72  pybind11::object rawOpViewClass);
73 
74  /// Looks up a registered dialect class by namespace. Note that this may
75  /// trigger loading of the defining module and can arbitrarily re-enter.
77  lookupDialectClass(const std::string &dialectNamespace);
78 
79  /// Looks up a registered raw OpView class by operation name. Note that this
80  /// may trigger a load of the dialect, which can arbitrarily re-enter.
82  lookupRawOpViewClass(llvm::StringRef operationName);
83 
84 private:
85  static PyGlobals *instance;
86  /// Module name prefixes to search under for dialect implementation modules.
87  std::vector<std::string> dialectSearchPrefixes;
88  /// Map of dialect namespace to external dialect class object.
89  llvm::StringMap<pybind11::object> dialectClassMap;
90  /// Map of full operation name to external operation class object.
91  llvm::StringMap<pybind11::object> operationClassMap;
92  /// Map of operation name to custom subclass that directly initializes
93  /// the OpView base class (bypassing the user class constructor).
94  llvm::StringMap<pybind11::object> rawOpViewClassMap;
95 
96  /// Set of dialect namespaces that we have attempted to import implementation
97  /// modules for.
98  llvm::StringSet<> loadedDialectModulesCache;
99  /// Cache of operation name to custom OpView subclass that directly
100  /// initializes the OpView base class (or an undefined object for negative
101  /// lookup). This is maintained on loopup as a shadow of rawOpViewClassMap
102  /// in order for repeat lookups of the OpView classes to only incur the cost
103  /// of one hashtable lookup.
104  llvm::StringMap<pybind11::object> rawOpViewClassMapCache;
105 };
106 
107 } // namespace python
108 } // namespace mlir
109 
110 #endif // MLIR_BINDINGS_PYTHON_GLOBALS_H
Include the generated interface declarations.
Globals that are always accessible once the extension has been initialized.
Definition: Globals.h:25
void setDialectSearchPrefixes(std::vector< std::string > newValues)
Definition: Globals.h:41
void loadDialectModule(llvm::StringRef dialectNamespace)
Loads a python module corresponding to the given dialect namespace.
Definition: IRModule.cpp:37
void clearImportCache()
Clears positive and negative caches regarding what implementations are available. ...
Definition: IRModule.cpp:138
void registerOperationImpl(const std::string &operationName, pybind11::object pyClass, pybind11::object rawOpViewClass)
Adds a concrete implementation operation class.
Definition: IRModule.cpp:74
std::vector< std::string > & getDialectSearchPrefixes()
Get and set the list of parent modules to search for dialect implementation classes.
Definition: Globals.h:38
llvm::Optional< pybind11::object > lookupDialectClass(const std::string &dialectNamespace)
Looks up a registered dialect class by namespace.
Definition: IRModule.cpp:88
pybind11::object registerDialectDecorator(pybind11::object pyClass)
Decorator for registering a custom Dialect class.
llvm::Optional< pybind11::object > lookupRawOpViewClass(llvm::StringRef operationName)
Looks up a registered raw OpView class by operation name.
Definition: IRModule.cpp:105
void registerDialectImpl(const std::string &dialectNamespace, pybind11::object pyClass)
Adds a concrete implementation dialect class.
Definition: IRModule.cpp:63