
This PR replaces the mixin `OpView` extension mechanism with the standard inheritance mechanism. Why? Firstly, mixins are not very pythonic (inheritance is usually used for this), a little convoluted, and too "tight" (can only be used in the immediately adjacent `_ext.py`). Secondly, it (mixins) are now blocking are correct implementation of "value builders" (see [here](https://github.com/llvm/llvm-project/pull/68764)) where the problem becomes how to choose the correct base class that the value builder should call. This PR looks big/complicated but appearances are deceiving; 4 things were needed to make this work: 1. Drop `skipDefaultBuilders` in `OpPythonBindingGen::emitDefaultOpBuilders` 2. Former mixin extension classes are converted to inherit from the generated `OpView` instead of being "mixins" a. extension classes that simply were calling into an already generated `super().__init__` continue to do so b. (almost all) extension classes that were calling `self.build_generic` because of a lack of default builder being generated can now also just call `super().__init__` 3. To handle the [lone single use-case](https://sourcegraph.com/search?q=context%3Aglobal+select_opview_mixin&patternType=standard&sm=1&groupBy=repo) of `select_opview_mixin`, namely [linalg](https://github.com/llvm/llvm-project/blob/main/mlir/python/mlir/dialects/_linalg_ops_ext.py#L38), only a small change was necessary in `opdsl/lang/emitter.py` (thanks to the emission/generation of default builders/`__init__`s) 4. since the `extend_opview_class` decorator is removed, we need a way to register extension classes as the desired `OpView` that `op.opview` conjures into existence; so we do the standard thing and just enable replacing the existing registered `OpView` i.e., `register_operation(_Dialect, replace=True)`. Note, the upgrade path for the common case is to change an extension to inherit from the generated builder and decorate it with `register_operation(_Dialect, replace=True)`. In the slightly more complicated case where `super().__init(self.build_generic(...))` is called in the extension's `__init__`, this needs to be updated to call `__init__` in `OpView`, i.e., the grandparent (see updated docs). Note, also `<DIALECT>_ext.py` files/modules will no longer be automatically loaded. Note, the PR has 3 base commits that look funny but this was done for the purpose of tracking the line history of moving the `<DIALECT>_ops_ext.py` class into `<DIALECT>.py` and updating (commit labeled "fix").
105 lines
4.0 KiB
C++
105 lines
4.0 KiB
C++
//===- MainModule.cpp - Main pybind module --------------------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include <tuple>
|
|
|
|
#include "PybindUtils.h"
|
|
|
|
#include "Globals.h"
|
|
#include "IRModule.h"
|
|
#include "Pass.h"
|
|
|
|
namespace py = pybind11;
|
|
using namespace mlir;
|
|
using namespace py::literals;
|
|
using namespace mlir::python;
|
|
|
|
// -----------------------------------------------------------------------------
|
|
// Module initialization.
|
|
// -----------------------------------------------------------------------------
|
|
|
|
PYBIND11_MODULE(_mlir, m) {
|
|
m.doc() = "MLIR Python Native Extension";
|
|
|
|
py::class_<PyGlobals>(m, "_Globals", py::module_local())
|
|
.def_property("dialect_search_modules",
|
|
&PyGlobals::getDialectSearchPrefixes,
|
|
&PyGlobals::setDialectSearchPrefixes)
|
|
.def(
|
|
"append_dialect_search_prefix",
|
|
[](PyGlobals &self, std::string moduleName) {
|
|
self.getDialectSearchPrefixes().push_back(std::move(moduleName));
|
|
self.clearImportCache();
|
|
},
|
|
"module_name"_a)
|
|
.def("_register_dialect_impl", &PyGlobals::registerDialectImpl,
|
|
"dialect_namespace"_a, "dialect_class"_a,
|
|
"Testing hook for directly registering a dialect")
|
|
.def("_register_operation_impl", &PyGlobals::registerOperationImpl,
|
|
"operation_name"_a, "operation_class"_a, "replace"_a = false,
|
|
"Testing hook for directly registering an operation");
|
|
|
|
// Aside from making the globals accessible to python, having python manage
|
|
// it is necessary to make sure it is destroyed (and releases its python
|
|
// resources) properly.
|
|
m.attr("globals") =
|
|
py::cast(new PyGlobals, py::return_value_policy::take_ownership);
|
|
|
|
// Registration decorators.
|
|
m.def(
|
|
"register_dialect",
|
|
[](py::object pyClass) {
|
|
std::string dialectNamespace =
|
|
pyClass.attr("DIALECT_NAMESPACE").cast<std::string>();
|
|
PyGlobals::get().registerDialectImpl(dialectNamespace, pyClass);
|
|
return pyClass;
|
|
},
|
|
"dialect_class"_a,
|
|
"Class decorator for registering a custom Dialect wrapper");
|
|
m.def(
|
|
"register_operation",
|
|
[](const py::object &dialectClass, bool replace) -> py::cpp_function {
|
|
return py::cpp_function(
|
|
[dialectClass, replace](py::object opClass) -> py::object {
|
|
std::string operationName =
|
|
opClass.attr("OPERATION_NAME").cast<std::string>();
|
|
PyGlobals::get().registerOperationImpl(operationName, opClass,
|
|
replace);
|
|
|
|
// Dict-stuff the new opClass by name onto the dialect class.
|
|
py::object opClassName = opClass.attr("__name__");
|
|
dialectClass.attr(opClassName) = opClass;
|
|
return opClass;
|
|
});
|
|
},
|
|
"dialect_class"_a, "replace"_a = false,
|
|
"Produce a class decorator for registering an Operation class as part of "
|
|
"a dialect");
|
|
m.def(
|
|
MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR,
|
|
[](MlirTypeID mlirTypeID, py::function typeCaster, bool replace) {
|
|
PyGlobals::get().registerTypeCaster(mlirTypeID, std::move(typeCaster),
|
|
replace);
|
|
},
|
|
"typeid"_a, "type_caster"_a, "replace"_a = false,
|
|
"Register a type caster for casting MLIR types to custom user types.");
|
|
|
|
// Define and populate IR submodule.
|
|
auto irModule = m.def_submodule("ir", "MLIR IR Bindings");
|
|
populateIRCore(irModule);
|
|
populateIRAffine(irModule);
|
|
populateIRAttributes(irModule);
|
|
populateIRInterfaces(irModule);
|
|
populateIRTypes(irModule);
|
|
|
|
// Define and populate PassManager submodule.
|
|
auto passModule =
|
|
m.def_submodule("passmanager", "MLIR Pass Management Bindings");
|
|
populatePassManagerSubmodule(passModule);
|
|
}
|