/// For the case of a python __init__ (py::init) method, pybind11 is quite /// strict about needing to return a pointer that is not yet associated to /// an py::object. Since the forContext() method acts like a pool, possibly /// returning a recycled context, it does not satisfy this need. The usual /// way in python to accomplish such a thing is to override __new__, but /// that is also not supported by pybind11. Instead, we use this entry /// point which always constructs a fresh context (which cannot alias an /// existing one because it is fresh). static PyMlirContext *createNewContextForInit();
/// Returns a context reference for the singleton PyMlirContext wrapper for /// the given context. static PyMlirContextRef forContext(MlirContext context); ~PyMlirContext();
/// Accesses the underlying MlirContext. MlirContext get(){ return context; }
/// Gets a strong reference to this context, which will ensure it is kept /// alive for the life of the reference. PyMlirContextRef getRef(){ returnPyMlirContextRef(this, pybind11::cast(this)); }
/// Gets a capsule wrapping the void* within the MlirContext. pybind11::object getCapsule();
/// Creates a PyMlirContext from the MlirContext wrapped by a capsule. /// Note that PyMlirContext instances are uniqued, so the returned object /// may be a pre-existing object. Ownership of the underlying MlirContext /// is taken by calling this function. static pybind11::object createFromCapsule(pybind11::object capsule);
/// Gets the count of live context objects. Used for testing. staticsize_tgetLiveCount();
/// Gets the count of live operations associated with this context. /// Used for testing. size_tgetLiveOperationCount();
/// Clears the live operations map, returning the number of entries which were /// invalidated. To be used as a safety mechanism so that API end-users can't /// corrupt by holding references they shouldn't have accessed in the first /// place. size_tclearLiveOperations();
/// Gets the count of live modules associated with this context. /// Used for testing. size_tgetLiveModuleCount();
/// Enter and exit the context manager. pybind11::object contextEnter(); voidcontextExit(const pybind11::object &excType, const pybind11::object &excVal, const pybind11::object &excTb);
/// Attaches a Python callback as a diagnostic handler, returning a /// registration object (internally a PyDiagnosticHandler). pybind11::object attachDiagnosticHandler(pybind11::object callback);
/// Controls whether error diagnostics should be propagated to diagnostic /// handlers, instead of being captured by `ErrorCapture`. voidsetEmitErrorDiagnostics(bool value){ emitErrorDiagnostics = value; } structErrorCapture;
private: PyMlirContext(MlirContext context); // Interns the mapping of live MlirContext::ptr to PyMlirContext instances, // preserving the relationship that an MlirContext maps to a single // PyMlirContext wrapper. This could be replaced in the future with an // extension mechanism on the MlirContext for stashing user pointers. // Note that this holds a handle, which does not imply ownership. // Mappings will be removed when the context is destructed. using LiveContextMap = llvm::DenseMap<void *, PyMlirContext *>; static LiveContextMap &getLiveContexts();
// Interns all live modules associated with this context. Modules tracked // in this map are valid. When a module is invalidated, it is removed // from this map, and while it still exists as an instance, any // attempt to access it will raise an error. using LiveModuleMap = llvm::DenseMap<constvoid *, std::pair<pybind11::handle, PyModule *>>; LiveModuleMap liveModules;
// Interns all live operations associated with this context. Operations // tracked in this map are valid. When an operation is invalidated, it is // removed from this map, and while it still exists as an instance, any // attempt to access it will raise an error. using LiveOperationMap = llvm::DenseMap<void *, std::pair<pybind11::handle, PyOperation *>>; LiveOperationMap liveOperations;
voidPyThreadContextEntry::push(FrameKind frameKind, py::object context, py::object insertionPoint, py::object location){ auto &stack = getStack(); stack.emplace_back(frameKind, std::move(context), std::move(insertionPoint), std::move(location)); // If the new stack has more than one entry and the context of the new top // entry matches the previous, copy the insertionPoint and location from the // previous entry if missing from the new top entry. if (stack.size() > 1) { auto &prev = *(stack.rbegin() + 1); auto ¤t = stack.back(); if (current.context.is(prev.context)) { // Default non-context objects from the previous entry. if (!current.insertionPoint) current.insertionPoint = prev.insertionPoint; if (!current.location) current.location = prev.location; } } }
defprocess_initializer_module(module_name): try: m = importlib.import_module(f".{module_name}", __name__) except ModuleNotFoundError: returnFalse except ImportError: message = ( f"Error importing mlir initializer {module_name}. This may " "happen in unclean incremental builds but is likely a real bug if " "encountered otherwise and the MLIR Python API may not function." ) logger.warning(message, exc_info=True)
logger.debug("Initializing MLIR with module: %s", module_name) ifhasattr(m, "register_dialects"): logger.debug("Registering dialects from initializer %r", m) m.register_dialects(registry) ifhasattr(m, "context_init_hook"): logger.debug("Adding context init hook from %r", m) post_init_hooks.append(m.context_init_hook) returnTrue
# If _mlirRegisterEverything is built, then include it as an initializer # module. process_initializer_module("_mlirRegisterEverything")
# Load all _site_initialize_{i} modules, where 'i' is a number starting # at 0. for i in itertools.count(): module_name = f"_site_initialize_{i}" ifnot process_initializer_module(module_name): break
classContext(ir._BaseContext): def__init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.append_dialect_registry(registry) for hook in post_init_hooks: hook(self) # TODO: There is some debate about whether we should eagerly load # all dialects. It is being done here in order to preserve existing # behavior. See: https://github.com/llvm/llvm-project/issues/56037 self.load_all_available_dialects()
ir.Context = Context
classMLIRError(Exception): """ An exception with diagnostic information. Has the following fields: message: str error_diagnostics: List[ir.DiagnosticInfo] """
def__str__(self): s = self.message if self.error_diagnostics: s += ":" for diag in self.error_diagnostics: s += ( "\nerror: " + str(diag.location)[4:-1] + ": " + diag.message.replace("\n", "\n ") ) for note in diag.notes: s += ( "\n note: " + str(note.location)[4:-1] + ": " + note.message.replace("\n", "\n ") ) return s
/// Checks whether the given attribute is a bool attribute. MLIR_CAPI_EXPORTED boolmlirAttributeIsABool(MlirAttribute attr);
/// Creates a bool attribute in the given context with the given value. MLIR_CAPI_EXPORTED MlirAttribute mlirBoolAttrGet(MlirContext ctx, int value);
// auto cast to Concrete type .def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR, [](PyAttribute &self) { MlirTypeID mlirTypeID = mlirAttributeGetTypeID(self); assert(!mlirTypeIDIsNull(mlirTypeID) && "mlirTypeID was expected to be non-null."); std::optional<pybind11::function> typeCaster = PyGlobals::get().lookupTypeCaster(mlirTypeID, mlirAttributeGetDialect(self)); if (!typeCaster) return py::cast(self); return typeCaster.value()(self); });