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. get()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. getDialectSearchPrefixes()38 std::vector<std::string> &getDialectSearchPrefixes() { 39 return dialectSearchPrefixes; 40 } setDialectSearchPrefixes(std::vector<std::string> newValues)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. 76 llvm::Optional<pybind11::object> 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. 81 llvm::Optional<pybind11::object> 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 111