Maksim Levental a2288a8944
[mlir][python] remove mixins (#68853)
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").
2023-10-19 16:20:14 -05:00

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);
}