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