[MLIR][python bindings] Add `register_value_builder` that can be used to wrap results of ops automatically.

[makslevental]

Summary

This diff adds functionality akin to register_attribute_builder but for Values, i.e., users can define functions that map to Python classes that wrap results from ops. Currently the selection is done in a metaclass using a regex match on the stringified type of the value but safer ways are possible (see this diff.

[makslevental]

remove unnecessary includes

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

py::cpp_function newCf(
                                          origRepr + ")")
                                             .str());
           }
           py::object self = superCls.attr("__new__")(cls, otherType);
           return self;
         },
         py::name("__new__"), py::arg("cls"), py::arg("cast_from_type"));
     thisClass.attr("__new__") = newCf;
+    thisClass.attr("__name__") = captureTypeName;
 
     // 'isinstance' method.
     def_staticmethod(
         "isinstance",
         [isaFunction](MlirType other) { return isaFunction(other); },
         py::arg("other_type"));
   }
 };

mlir/lib/Bindings/Python/Globals.h

public:
   pybind11::object registerDialectDecorator(pybind11::object pyClass);
 
   /// Adds a user-friendly Attribute builder.
   /// Raises an exception if the mapping already exists.
   /// This is intended to be called by implementation code.
   void registerAttributeBuilder(const std::string &attributeKind,
                                 pybind11::function pyFunc);
 
+  void registerValueBuilder(const std::string &typeRegex,
+                            pybind11::function pyFunc);
+
   /// Adds a concrete implementation dialect class.
   /// Raises an exception if the mapping already exists.
   /// 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);
 
   /// Returns the custom Attribute builder for Attribute kind.
   std::optional<pybind11::function>
   lookupAttributeBuilder(const std::string &attributeKind);
 
+  std::optional<pybind11::function> lookupValueBuilder(const std::string &type);
+
   /// 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 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>
   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 attribute ODS name to custom builder.
   llvm::StringMap<pybind11::object> attributeBuilderMap;
 
+  std::map<std::string, pybind11::object> valueBuilderMap;
+
   /// Set of dialect namespaces that we have attempted to import implementation
   /// modules for.
   llvm::StringSet<> loadedDialectModulesCache;
   /// 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

 #include "mlir-c/Debug.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include <optional>
+#include <pybind11/pytypes.h>
 #include <utility>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
 using llvm::StringRef;
struct PyAttrBuilderMap {
   static void bind(py::module &m) {
     py::class_<PyAttrBuilderMap>(m, "AttrBuilder", py::module_local())
         .def_static("contains", &PyAttrBuilderMap::dunderContains)
         .def_static("get", &PyAttrBuilderMap::dundeGetItemNamed)
         .def_static("insert", &PyAttrBuilderMap::dundeSetItemNamed);
   }
 };
 
+struct PyValueBuilderMap {
+  static bool dunderContains(py::object type) {
+    return PyGlobals::get().lookupValueBuilder(py::str(type)).has_value();
+  }
+  static py::function dundeGetItemNamed(py::object type) {
+    auto builder = PyGlobals::get().lookupValueBuilder(py::str(type));
+    if (!builder)
+      throw py::key_error();
+    return *builder;
+  }
+  static void dundeSetItemNamed(pybind11::object typeRegex, py::function func) {
+    PyGlobals::get().registerValueBuilder(py::str(typeRegex), std::move(func));
+  }
+
+  static void bind(py::module &m) {
+    py::class_<PyValueBuilderMap>(m, "ValueBuilder", py::module_local())
+        .def_static("contains", &PyValueBuilderMap::dunderContains)
+        .def_static("get", &PyValueBuilderMap::dundeGetItemNamed)
+        .def_static("insert", &PyValueBuilderMap::dundeSetItemNamed);
+  }
+};
+
 //------------------------------------------------------------------------------
 // Collections.
 //------------------------------------------------------------------------------
 
 namespace {
 
 class PyRegionIterator {
 public:

   PyRegionIterator::bind(m);
   PyRegionList::bind(m);
 
   // Debug bindings.
   PyGlobalDebugFlag::bind(m);
 
   // Attribute builder getter.
   PyAttrBuilderMap::bind(m);
+  PyValueBuilderMap::bind(m);
 
   py::register_local_exception_translator([](std::exception_ptr p) {
     // We can't define exceptions with custom fields through pybind, so instead
     // the exception class is defined in python and imported here.
     try {
       if (p)
         std::rethrow_exception(p);
     } catch (const MLIRError &e) {
       py::object obj = py::module_::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
                            .attr("MLIRError")(e.message, e.errorDiagnostics);
       PyErr_SetObject(PyExc_Exception, obj.ptr());
     }
   });
 }

mlir/lib/Bindings/Python/IRModule.cpp

 //===- IRModule.cpp - IR 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 "IRModule.h"
 #include "Globals.h"
 #include "PybindUtils.h"
 
-#include <vector>
 #include <optional>
+#include <pybind11/pytypes.h>
+#include <regex>
+#include <utility>
+#include <vector>
 
 #include "mlir-c/Bindings/Python/Interop.h"
+#include "llvm/ADT/StringRef.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 // -----------------------------------------------------------------------------
 // PyGlobals
 // -----------------------------------------------------------------------------
void PyGlobals::registerAttributeBuilder(const std::string &attributeKind,
   if (found) {
     throw std::runtime_error((llvm::Twine("Attribute builder for '") +
                               attributeKind + "' is already registered")
                                  .str());
   }
   found = std::move(pyFunc);
 }
 
+void PyGlobals::registerValueBuilder(const std::string &typeRegex,
+                                     py::function pyFunc) {
+  py::object &found = valueBuilderMap[typeRegex];
+  if (found) {
+    throw std::runtime_error(
+        (llvm::Twine("Value builder for '") + "' is already registered").str());
+  }
+  found = std::move(pyFunc);
+}
+
 void PyGlobals::registerDialectImpl(const std::string &dialectNamespace,
                                     py::object pyClass) {
   py::object &found = dialectClassMap[dialectNamespace];
   if (found) {
     throw SetPyError(PyExc_RuntimeError, llvm::Twine("Dialect namespace '") +
                                              dialectNamespace +
                                              "' is already registered.");
   }
if (foundIt != attributeBuilderMap.end()) {
     return foundIt->second;
   }
 
   // Not found and loading did not yield a registration. Negative cache.
   attributeBuilderMap[attributeKind] = py::none();
   return std::nullopt;
 }
 
+std::optional<py::function>
+PyGlobals::lookupValueBuilder(const std::string &type) {
+  const auto foundIt = std::find_if(
+      valueBuilderMap.begin(), valueBuilderMap.end(), [&type](auto pair) {
+        return std::regex_match(type, std::regex(pair.first));
+      });
+  if (foundIt != valueBuilderMap.end()) {
+    if (foundIt->second.is_none())
+      return std::nullopt;
+    assert(foundIt->second && "py::function is defined");
+    return foundIt->second;
+  }
+
+  // Not found and loading did not yield a registration. Negative cache.
+  valueBuilderMap[type] = py::none();
+  return std::nullopt;
+}
+
 std::optional<py::object>
 PyGlobals::lookupDialectClass(const std::string &dialectNamespace) {
   loadDialectModule(dialectNamespace);
   // Fast match against the class map first (common case).
   const auto foundIt = dialectClassMap.find(dialectNamespace);
   if (foundIt != dialectClassMap.end()) {
     if (foundIt->second.is_none())
       return std::nullopt;

mlir/python/mlir/dialects/_ods_common.py

 #  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
 
 # Provide a convenient name for sub-packages to resolve the main C-extension
 # with a relative import.
+from ..ir import Value, ValueBuilder
 from .._mlir_libs import _mlir as _cext
+import inspect
 from typing import Sequence as _Sequence, Union as _Union
 
 __all__ = [
     "equally_sized_accessor",
     "extend_opview_class",
     "get_default_loc_context",
     "get_op_result_or_value",
     "get_op_results_or_values",
     "segmented_accessor",
 ]
 
 
+class ValueMeta(type(Value)):
+
+  def __call__(cls, *args, **kwargs):
+    cls_obj = cls.__new__(cls, *args, **kwargs)
+    cls.__init__(cls_obj, *args, **kwargs)
+    if len(cls_obj.results) == 1:
+      val = get_op_result_or_value(cls_obj)
+      if ValueBuilder.contains(val.type):
+        return ValueBuilder.get(val.type)(val)
+    return cls_obj
+
+
+def rebuild_with_meta(parent_opview_cls, mixin=False):
+  v = ValueMeta(f"{parent_opview_cls.__name__}",
+                parent_opview_cls.__bases__, dict(parent_opview_cls.__dict__))
+
+  # mixins (extensions) for some reasons don't suffer from this problem
+  # i.e., the __class__ is the correctly patched/hacked one
+  if not mixin:
+    # some ops don't have __init__ but one is inherited from OpView (as an instancemethod)
+    if not inspect.ismethoddescriptor(v.__init__):
+      v.__init__.__closure__[0].cell_contents = v
+  return v
+
+
 def extend_opview_class(ext_module):
   """Decorator to extend an OpView class from an extension module.
 
   Extension modules can expose various entry-points:
     Stand-alone class with the same name as a parent OpView class (i.e.
     "ReturnOp"). A name-based match is attempted first before falling back
     to a below mechanism.
 
     def select_opview_mixin(parent_opview_cls):
       If defined, allows an appropriate mixin class to be selected dynamically
       based on the parent OpView class. Should return NotImplemented if a
       decision is not made.
 
   Args:
     ext_module: A module from which to locate extensions. Can be None if not
       available.
 
   Returns:
     A decorator that takes an OpView subclass and further extends it as
     needed.
   """
 
   def class_decorator(parent_opview_cls: type):
     if ext_module is None:
-      return parent_opview_cls
+      return rebuild_with_meta(parent_opview_cls)
     mixin_cls = NotImplemented
     # First try to resolve by name.
     try:
       mixin_cls = getattr(ext_module, parent_opview_cls.__name__)
     except AttributeError:
       # Fall back to a select_opview_mixin hook.
       try:
         select_mixin = getattr(ext_module, "select_opview_mixin")
       except AttributeError:
         pass
       else:
         mixin_cls = select_mixin(parent_opview_cls)
 
     if mixin_cls is NotImplemented or mixin_cls is None:
-      return parent_opview_cls
+      return rebuild_with_meta(parent_opview_cls)
 
     # Have a mixin_cls. Create an appropriate subclass.
     try:
 
       class LocalOpView(mixin_cls, parent_opview_cls):
         pass
     except TypeError as e:
       raise TypeError(
-          f"Could not mixin {mixin_cls} into {parent_opview_cls}") from e
     LocalOpView.__name__ = parent_opview_cls.__name__
     LocalOpView.__qualname__ = parent_opview_cls.__qualname__
-    return LocalOpView
+    return rebuild_with_meta(LocalOpView, mixin=True)
 
   return class_decorator
 
 
 def segmented_accessor(elements, raw_segments, idx):
   """
   Returns a slice of elements corresponding to the idx-th segment.
 

mlir/python/mlir/ir.py

def _indexElementsAttr(x, context):
     return DenseElementsAttr.get(
         np.array(x, dtype=np.int64),
         type=IndexType.get(context=context),
         context=context,
     )
 
 except ImportError:
   pass
+
+
+def register_value_builder(kind):
+
+  def decorator_builder(func):
+    ValueBuilder.insert(kind, func)
+    return func
+
+  return decorator_builder
 No newline at end of file

mlir/test/python/ir/value.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
-from mlir.dialects import func
+from mlir.dialects import func, arith, tensor
 
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
   return f
with Location.unknown(ctx):
     # CHECK:     return
     # CHECK:   }
     # CHECK: }
     print(module)
 
     value2.owner.detach_from_parent()
     # CHECK: %0
     print(value2.get_name())
+
+
+# CHECK-LABEL: TEST: testValueBuilders
+@run
+def testValueBuilders():
+
+  class ScalarOrTensor:
+    def __init__(self, val):
+      self.val = val
+
+    def __add__(self, other):
+      return arith.AddFOp(self.val, other.val)
+
+    def __str__(self):
+      return str(self.val).replace("Value", "ScalarOrTensor")
+
+  @register_value_builder("f64")
+  @register_value_builder("tensor<.*>")
+  def _buildF64Value(val):
+    return ScalarOrTensor(val)
+
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  with Location.unknown(ctx):
+    f64 = F64Type.get()
+    module = Module.create()
+    with InsertionPoint(module.body):
+      one_f64 = arith.ConstantOp(f64, 1.0)
+      # CHECK: ScalarOrTensor(%[[ONE:.*]] = arith.constant 1.000000e+00 : f64)
+      print(one_f64)
+
+      two_f64 = arith.ConstantOp(f64, 2.0)
+      # CHECK: ScalarOrTensor(%[[TWO:.*]] = arith.constant 2.000000e+00 : f64)
+      print(two_f64)
+
+      three_f64 = one_f64 + two_f64
+      # CHECK: ScalarOrTensor(%[[THREE:.*]] = arith.addf %[[ONE:.*]], %[[TWO:.*]] : f64)
+      print(three_f64)
+
+      a_ten_f64 = tensor.EmptyOp([10, 10], f64)
+      # CHECK: ScalarOrTensor(%[[ONE:.*]] = tensor.empty() : tensor<10x10xf64>)
+      print(a_ten_f64)
+      b_ten_f64 = tensor.EmptyOp([10, 10], f64)
+      # CHECK: ScalarOrTensor(%[[ONE:.*]] = tensor.empty() : tensor<10x10xf64>)
+      print(b_ten_f64)
+      c_ten_f64 = a_ten_f64 + b_ten_f64
+      # CHECK: ScalarOrTensor(%[[ONE:.*]] = arith.addf %1, %2 : tensor<10x10xf64>)
+      print(c_ten_f64)
+
+