[mlir][python] Remove "Raw" OpView classes

[rkayaith]

Summary

The raw OpView classes are used to bypass the constructors of OpView
subclasses, but having a separate class can create some confusing
behaviour, e.g.:

op = MyOp(...)
# fails, lhs is 'MyOp', rhs is '_MyOp'
assert type(op) == type(op.operation.opview)

Instead we can use __new__ to achieve the same thing without a
separate class:

my_op = MyOp.__new__(MyOp)
OpView.__init__(my_op, op)

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

1643–1645

Not sure what the original concern with __new__ was, it seems valid to use as far as I can tell. I didn't see any discussion about this on the original review either: https://reviews.llvm.org/D89932

[rkayaith]

friendly ping

[stellaraccident]

Thanks. I had started reviewing this but got side tracked with the archaeology. Thanks for simplifying -- this was always not great.

mlir/lib/Bindings/Python/IRCore.cpp

1643–1645

I don't remember either but I do remember struggling with it. I'm +1 on simplifying since this seems to work.

mlir/lib/Bindings/Python/Globals.h

public:
   /// This is intended to be called by implementation code.
   void registerDialectImpl(const std::string &dialectNamespace,
                            pybind11::object pyClass);
 
   /// Adds a concrete implementation operation class.
   /// Raises an exception if the mapping already exists.
   /// This is intended to be called by implementation code.
   void registerOperationImpl(const std::string &operationName,
-                             pybind11::object pyClass,
-                             pybind11::object rawOpViewClass);
+                             pybind11::object pyClass);
 
   /// Returns the custom Attribute builder for Attribute kind.
   std::optional<pybind11::function>
   lookupAttributeBuilder(const std::string &attributeKind);
 
   /// Looks up a registered dialect class by namespace. Note that this may
   /// trigger loading of the defining module and can arbitrarily re-enter.
   std::optional<pybind11::object>
   lookupDialectClass(const std::string &dialectNamespace);
 
-  /// Looks up a registered raw OpView class by operation name. Note that this
-  /// may trigger a load of the dialect, which can arbitrarily re-enter.
+  /// Looks up a registered operation class (deriving from OpView) by operation
+  /// name. Note that this may trigger a load of the dialect, which can
+  /// arbitrarily re-enter.
   std::optional<pybind11::object>
-  lookupRawOpViewClass(llvm::StringRef operationName);
+  lookupOperationClass(llvm::StringRef operationName);
 
 private:
   static PyGlobals *instance;
   /// Module name prefixes to search under for dialect implementation modules.
   std::vector<std::string> dialectSearchPrefixes;
   /// Map of dialect namespace to external dialect class object.
   llvm::StringMap<pybind11::object> dialectClassMap;
   /// Map of full operation name to external operation class object.
   llvm::StringMap<pybind11::object> operationClassMap;
-  /// Map of operation name to custom subclass that directly initializes
-  /// the OpView base class (bypassing the user class constructor).
-  llvm::StringMap<pybind11::object> rawOpViewClassMap;
   /// Map of attribute ODS name to custom builder.
   llvm::StringMap<pybind11::object> attributeBuilderMap;
 
   /// Set of dialect namespaces that we have attempted to import implementation
   /// modules for.
   llvm::StringSet<> loadedDialectModulesCache;
-  /// Cache of operation name to custom OpView subclass that directly
-  /// initializes the OpView base class (or an undefined object for negative
-  /// lookup). This is maintained on loopup as a shadow of rawOpViewClassMap
-  /// in order for repeat lookups of the OpView classes to only incur the cost
-  /// of one hashtable lookup.
-  llvm::StringMap<pybind11::object> rawOpViewClassMapCache;
+  /// Cache of operation name to external operation class object. This is
+  /// maintained on lookup as a shadow of operationClassMap in order for repeat
+  /// lookups of the classes to only incur the cost of one hashtable lookup.
+  llvm::StringMap<pybind11::object> operationClassMapCache;
 };
 
 } // namespace python
 } // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_GLOBALS_H

mlir/lib/Bindings/Python/IRCore.cpp

py::object PyOperation::clone(const py::object &maybeIp) {
 
   return cloned->createOpView();
 }
 
 py::object PyOperation::createOpView() {
   checkValid();
   MlirIdentifier ident = mlirOperationGetName(get());
   MlirStringRef identStr = mlirIdentifierStr(ident);
-  auto opViewClass = PyGlobals::get().lookupRawOpViewClass(
+  auto operationCls = PyGlobals::get().lookupOperationClass(
       StringRef(identStr.data, identStr.length));
-  if (opViewClass)
-    return (*opViewClass)(getRef().getObject());
+  if (operationCls)
+    return PyOpView::constructDerived(*operationCls, *getRef().get());
   return py::cast(PyOpView(getRef().getObject()));
 }
 
 void PyOperation::erase() {
   checkValid();
   // TODO: Fix memory hazards when erasing a tree of operations for which a deep
   // Python reference to a child operation is live. All children should also
   // have their `valid` bit set to false.
PyOpView::buildGeneric(const py::object &cls, py::list resultTypeList,
   return PyOperation::create(name,
                              /*results=*/std::move(resultTypes),
                              /*operands=*/std::move(operands),
                              /*attributes=*/std::move(attributes),
                              /*successors=*/std::move(successors),
                              /*regions=*/*regions, location, maybeIp);
 }
 
+pybind11::object PyOpView::constructDerived(const pybind11::object &cls,
+                                            const PyOperation &operation) {
+  // TODO: pybind11 2.6 supports a more direct form.
+  // Upgrade many years from now.
+  //   auto opViewType = py::type::of<PyOpView>();
+  py::handle opViewType = py::detail::get_type_handle(typeid(PyOpView), true);
+  py::object instance = cls.attr("__new__")(cls);
+  opViewType.attr("__init__")(instance, operation);
+  return instance;
+}
+
 PyOpView::PyOpView(const py::object &operationObject)
     // Casting through the PyOperationBase base-class and then back to the
     // Operation lets us accept any PyOperationBase subclass.
     : operation(py::cast<PyOperationBase &>(operationObject).getOperation()),
       operationObject(operation.getRef().getObject()) {}
 
-py::object PyOpView::createRawSubclass(const py::object &userClass) {
-  // This is... a little gross. The typical pattern is to have a pure python
-  // class that extends OpView like:
-  //   class AddFOp(_cext.ir.OpView):
-  //     def __init__(self, loc, lhs, rhs):
-  //       operation = loc.context.create_operation(
-  //           "addf", lhs, rhs, results=[lhs.type])
-  //       super().__init__(operation)
-  //
-  // I.e. The goal of the user facing type is to provide a nice constructor
-  // that has complete freedom for the op under construction. This is at odds
-  // with our other desire to sometimes create this object by just passing an
-  // operation (to initialize the base class). We could do *arg and **kwargs
-  // munging to try to make it work, but instead, we synthesize a new class
-  // on the fly which extends this user class (AddFOp in this example) and
-  // *give it* the base class's __init__ method, thus bypassing the
-  // intermediate subclass's __init__ method entirely. While slightly,
-  // underhanded, this is safe/legal because the type hierarchy has not changed
-  // (we just added a new leaf) and we aren't mucking around with __new__.
-  // Typically, this new class will be stored on the original as "_Raw" and will
-  // be used for casts and other things that need a variant of the class that
-  // is initialized purely from an operation.
-  py::object parentMetaclass =
-      py::reinterpret_borrow<py::object>((PyObject *)&PyType_Type);
-  py::dict attributes;
-  // TODO: pybind11 2.6 supports a more direct form. Upgrade many years from
-  // now.
-  //   auto opViewType = py::type::of<PyOpView>();
-  auto opViewType = py::detail::get_type_handle(typeid(PyOpView), true);
-  attributes["__init__"] = opViewType.attr("__init__");
-  py::str origName = userClass.attr("__name__");
-  py::str newName = py::str("_") + origName;
-  return parentMetaclass(newName, py::make_tuple(userClass), attributes);
-}
-
 //------------------------------------------------------------------------------
 // PyInsertionPoint.
 //------------------------------------------------------------------------------
 
 PyInsertionPoint::PyInsertionPoint(PyBlock &block) : block(block) {}
 
 PyInsertionPoint::PyInsertionPoint(PyOperationBase &beforeOperationBase)
     : refOperation(beforeOperationBase.getOperation().getRef()),
opViewClass.attr("parse") = classmethod(
         std::string clsOpName =
             py::cast<std::string>(cls.attr("OPERATION_NAME"));
         MlirStringRef parsedOpName =
             mlirIdentifierStr(mlirOperationGetName(*parsed.get()));
         if (!mlirStringRefEqual(parsedOpName, toMlirStringRef(clsOpName)))
           throw py::value_error(
               "Expected a '" + clsOpName + "' op, got: '" +
               std::string(parsedOpName.data, parsedOpName.length) + "'");
-        return cls.attr("_Raw")(parsed.getObject());
+        return PyOpView::constructDerived(cls, *parsed.get());
       },
       py::arg("cls"), py::arg("source"), py::kw_only(),
       py::arg("source_name") = "", py::arg("context") = py::none(),
       "Parses a specific, generated OpView based on class level attributes");
 
   //----------------------------------------------------------------------------
   // Mapping of PyRegion.
   //----------------------------------------------------------------------------

mlir/lib/Bindings/Python/IRModule.h

 /// custom ODS-style operation classes. Since this class is subclass on the
 /// python side, it must present an __init__ method that operates in pure
 /// python types.
 class PyOpView : public PyOperationBase {
 public:
   PyOpView(const pybind11::object &operationObject);
   PyOperation &getOperation() override { return operation; }
 
-  static pybind11::object createRawSubclass(const pybind11::object &userClass);
-
   pybind11::object getOperationObject() { return operationObject; }
 
   static pybind11::object
   buildGeneric(const pybind11::object &cls, pybind11::list resultTypeList,
                pybind11::list operandList,
                std::optional<pybind11::dict> attributes,
                std::optional<std::vector<PyBlock *>> successors,
                std::optional<int> regions, DefaultingPyLocation location,
                const pybind11::object &maybeIp);
 
+  /// Construct an instance of a class deriving from OpView, bypassing its
+  /// `__init__` method. The derived class will typically define a constructor
+  /// that provides a convenient builder, but we need to side-step this when
+  /// constructing an `OpView` for an already-built operation.
+  ///
+  /// The caller is responsible for verifying that `operation` is a valid
+  /// operation to construct `cls` with.
+  static pybind11::object constructDerived(const pybind11::object &cls,
+                                           const PyOperation &operation);
+
 private:
   PyOperation &operation;           // For efficient, cast-free access from C++
   pybind11::object operationObject; // Holds the reference.
 };
 
 /// Wrapper around an MlirRegion.
 /// Regions are managed completely by their containing operation. Unlike the
 /// C++ API, the python API does not support detached regions.

mlir/lib/Bindings/Python/IRModule.cpp

if (found) {
     throw SetPyError(PyExc_RuntimeError, llvm::Twine("Dialect namespace '") +
                                              dialectNamespace +
                                              "' is already registered.");
   }
   found = std::move(pyClass);
 }
 
 void PyGlobals::registerOperationImpl(const std::string &operationName,
-                                      py::object pyClass,
-                                      py::object rawOpViewClass) {
+                                      py::object pyClass) {
   py::object &found = operationClassMap[operationName];
   if (found) {
     throw SetPyError(PyExc_RuntimeError, llvm::Twine("Operation '") +
                                              operationName +
                                              "' is already registered.");
   }
   found = std::move(pyClass);
-  rawOpViewClassMap[operationName] = std::move(rawOpViewClass);
 }
 
 std::optional<py::function>
 PyGlobals::lookupAttributeBuilder(const std::string &attributeKind) {
   // Fast match against the class map first (common case).
   const auto foundIt = attributeBuilderMap.find(attributeKind);
   if (foundIt != attributeBuilderMap.end()) {
     if (foundIt->second.is_none())
PyGlobals::lookupDialectClass(const std::string &dialectNamespace) {
   }
 
   // Not found and loading did not yield a registration. Negative cache.
   dialectClassMap[dialectNamespace] = py::none();
   return std::nullopt;
 }
 
 std::optional<pybind11::object>
-PyGlobals::lookupRawOpViewClass(llvm::StringRef operationName) {
+PyGlobals::lookupOperationClass(llvm::StringRef operationName) {
   {
-    auto foundIt = rawOpViewClassMapCache.find(operationName);
-    if (foundIt != rawOpViewClassMapCache.end()) {
+    auto foundIt = operationClassMapCache.find(operationName);
+    if (foundIt != operationClassMapCache.end()) {
       if (foundIt->second.is_none())
         return std::nullopt;
       assert(foundIt->second && "py::object is defined");
       return foundIt->second;
     }
   }
 
   // Not found. Load the dialect namespace.
   auto split = operationName.split('.');
   llvm::StringRef dialectNamespace = split.first;
   loadDialectModule(dialectNamespace);
 
   // Attempt to find from the canonical map and cache.
   {
-    auto foundIt = rawOpViewClassMap.find(operationName);
-    if (foundIt != rawOpViewClassMap.end()) {
+    auto foundIt = operationClassMap.find(operationName);
+    if (foundIt != operationClassMap.end()) {
       if (foundIt->second.is_none())
         return std::nullopt;
       assert(foundIt->second && "py::object is defined");
       // Positive cache.
-      rawOpViewClassMapCache[operationName] = foundIt->second;
+      operationClassMapCache[operationName] = foundIt->second;
       return foundIt->second;
     }
     // Negative cache.
-    rawOpViewClassMap[operationName] = py::none();
+    operationClassMap[operationName] = py::none();
     return std::nullopt;
   }
 }
 
 void PyGlobals::clearImportCache() {
   loadedDialectModulesCache.clear();
-  rawOpViewClassMapCache.clear();
+  operationClassMapCache.clear();
 }

mlir/lib/Bindings/Python/MainModule.cpp

py::class_<PyGlobals>(m, "_Globals", py::module_local())
             self.clearImportCache();
           },
           py::arg("module_name"))
       .def("_register_dialect_impl", &PyGlobals::registerDialectImpl,
            py::arg("dialect_namespace"), py::arg("dialect_class"),
            "Testing hook for directly registering a dialect")
       .def("_register_operation_impl", &PyGlobals::registerOperationImpl,
            py::arg("operation_name"), py::arg("operation_class"),
-           py::arg("raw_opview_class"),
            "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);
 
m.def(
       "Class decorator for registering a custom Dialect wrapper");
   m.def(
       "register_operation",
       [](py::object dialectClass) -> py::cpp_function {
         return py::cpp_function(
             [dialectClass](py::object opClass) -> py::object {
               std::string operationName =
                   opClass.attr("OPERATION_NAME").cast<std::string>();
-              auto rawSubclass = PyOpView::createRawSubclass(opClass);
-              PyGlobals::get().registerOperationImpl(operationName, opClass,
-                                                     rawSubclass);
+              PyGlobals::get().registerOperationImpl(operationName, opClass);
 
               // Dict-stuff the new opClass by name onto the dialect class.
               py::object opClassName = opClass.attr("__name__");
               dialectClass.attr(opClassName) = opClass;
-
-              // Now create a special "Raw" subclass that passes through
-              // construction to the OpView parent (bypasses the intermediate
-              // child's __init__).
-              opClass.attr("_Raw") = rawSubclass;
               return opClass;
             });
       },
       py::arg("dialect_class"),
       "Produce a class decorator for registering an Operation class as part of "
       "a dialect");
 
   // Define and populate IR submodule.

mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi

 from typing import List
 
 globals: "_Globals"
 
 class _Globals:
     dialect_search_modules: List[str]
     def _register_dialect_impl(self, dialect_namespace: str, dialect_class: type) -> None: ...
-    def _register_operation_impl(self, operation_name: str, operation_class: type, raw_opview_class: type) -> None: ...
+    def _register_operation_impl(self, operation_name: str, operation_class: type) -> None: ...
     def append_dialect_search_prefix(self, module_name: str) -> None: ...
 
 def register_dialect(dialect_class: type) -> object: ...
 def register_operation(dialect_class: type) -> object: ...

mlir/test/python/ir/operation.py

module = Module.parse(r"""
       %2 = "custom.f32"() : () -> f32
       %3 = arith.addf %1, %2 : f32
     """)
     print(module)
 
     # addf should map to a known OpView class in the arithmetic dialect.
     # We know the OpView for it defines an 'lhs' attribute.
     addf = module.body.operations[2]
-    # CHECK: <mlir.dialects._arith_ops_gen._AddFOp object
+    # CHECK: <mlir.dialects._arith_ops_gen.AddFOp object
     print(repr(addf))
     # CHECK: "custom.f32"()
     print(addf.lhs)
 
     # One of the custom ops should resolve to the default OpView.
     custom = module.body.operations[0]
     # CHECK: OpView object
     print(repr(custom))