[mlir][python] Add facility for extending generated python ODS.

[stellaraccident]

Summary

• This isn't exclusive with other mechanisms for more ODS centric op definitions, but based on discussions, we feel that we will always benefit from a python escape hatch, and that is the most natural way to write things that don't fit the mold.
• I suspect this facility needs further tweaking, and once it settles, I'll document it and add more tests.
• Added extensions for linalg, since it is unusable without them and continued to evolve my e2e example.

[ftynse]

mlir/lib/Bindings/Python/mlir/dialects/_linalg.py

22–23

Isn't there a trait? Not sure how to expose traits to Python though

[stellaraccident]

mlir/lib/Bindings/Python/mlir/dialects/_linalg.py

22–23

I'm sure there is a way. I'm using TODOs guide the next step work on these things to keep them more bite-sized.

[stellaraccident]

Rebase

mlir/examples/python/.style.yapf

+[style]
+  based_on_style = google
+  column_limit = 80
+  indent_width = 2

mlir/examples/python/linalg_matmul.py

 
 from mlir.ir import *
 from mlir.dialects import linalg
 from mlir.dialects import std
 
 
 # TODO: This should be in the core API.
 def FuncOp(name: str, func_type: Type) -> Tuple[Operation, Block]:
-    """Creates a |func| op.
     TODO: This should really be in the MLIR API.
     Returns:
       (operation, entry_block)
     """
-    attrs = {
-        "type": TypeAttr.get(func_type),
-        "sym_name": StringAttr.get(name),
-    }
-    op = Operation.create("func", regions=1, attributes=attrs)
-    body_region = op.regions[0]
-    entry_block = body_region.blocks.append(*func_type.inputs)
-    return op, entry_block
 
 
-# TODO: Generate customs builder vs patching one in.
-def PatchMatmulOpInit(self, lhs, rhs, result, loc=None, ip=None):
-    super(linalg.MatmulOp, self).__init__(
-        self._ods_build_default(operands=[[lhs, rhs], [result]],
-                                results=[],
-                                loc=loc,
-                                ip=ip))
-
-
-linalg.MatmulOp.__init__ = PatchMatmulOpInit
-
-
-def build_matmul_func(func_name, m, k, n, dtype):
-    lhs_type = MemRefType.get(dtype, [m, k])
-    rhs_type = MemRefType.get(dtype, [k, n])
-    result_type = MemRefType.get(dtype, [m, n])
-    # TODO: There should be a one-liner for this.
-    func_type = FunctionType.get([lhs_type, rhs_type, result_type], [])
-    _, entry = FuncOp(func_name, func_type)
-    lhs, rhs, result = entry.arguments
-    with InsertionPoint(entry):
-        linalg.MatmulOp(lhs, rhs, result)
-        std.ReturnOp([])
+def build_matmul_buffers_func(func_name, m, k, n, dtype):
+    linalg.MatmulOp([lhs, rhs], [result])
 
 
+def build_matmul_tensors_func(func_name, m, k, n, dtype):
+  # TODO: MemRefType and TensorTypes should not have inverted dtype/shapes
+  # from each other.
+  lhs_type = RankedTensorType.get([m, k], dtype)
+  rhs_type = RankedTensorType.get([k, n], dtype)
+  result_type = RankedTensorType.get([m, n], dtype)
+  # TODO: There should be a one-liner for this.
+  func_type = FunctionType.get([lhs_type, rhs_type], [result_type])
+  _, entry = FuncOp(func_name, func_type)
+  lhs, rhs = entry.arguments
+  with InsertionPoint(entry):
+    result = linalg.MatmulOp([lhs, rhs], results=[result_type]).result
+    std.ReturnOp([result])
+
+
 def run():
-    with Context() as c, Location.unknown():
-        module = Module.create()
-        # TODO: This at_block_terminator vs default construct distinction feels
-        # wrong and is error-prone.
-        with InsertionPoint.at_block_terminator(module.body):
-            build_matmul_func('main', 18, 32, 96, F32Type.get())
+      build_matmul_buffers_func('main_buffers', 18, 32, 96, F32Type.get())
+      build_matmul_tensors_func('main_tensors', 18, 32, 96, F32Type.get())
 
-        print(module)
 
 
-if __name__ == '__main__': run()
+if __name__ == '__main__':
+  run()

mlir/lib/Bindings/Python/.style.yapf

+[style]
+  based_on_style = google
+  column_limit = 80
+  indent_width = 2

mlir/lib/Bindings/Python/CMakeLists.txt

 include(AddMLIRPythonExtension)
 add_custom_target(MLIRBindingsPythonExtension)
 
 ################################################################################
 # Copy python source tree.
 ################################################################################
 
 set(PY_SRC_FILES
   mlir/__init__.py
   mlir/_dlloader.py
   mlir/ir.py
   mlir/dialects/__init__.py
+  mlir/dialects/_linalg.py
   mlir/ir.py
   mlir/passmanager.py
   mlir/transforms/__init__.py
 )
 
 add_custom_target(MLIRBindingsPythonSources ALL
   DEPENDS ${PY_SRC_FILES}
 )

mlir/lib/Bindings/Python/mlir/dialects/__init__.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
 
 # Re-export the parent _cext so that every level of the API can get it locally.
 from .. import _cext
 
-def _segmented_accessor(elements, raw_segments, idx):
+__all__ = [
+    "equally_sized_accessor",
+    "extend_opview_class",
+    "get_default_loc_context",
+    "segmented_accessor",
+]
+
+
+def extend_opview_class(ext_module):
+  """Decorator to extend an OpView class from an extension module.
+
+  Extension modules can expose various entry-points:
+    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.
+
+    Stand-alone class with the same name as a parent OpView class (i.e.
+    "ReturnOp").
+
+  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
+    mixin_cls = NotImplemented
+    try:
+      select_mixin = getattr(ext_module, "select_opview_mixin")
+    except AttributeError:
+      # Try to default resolve it.
+      try:
+        select_mixin = getattr(ext_module, parent_opview_cls.__name__)
+      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
+
+    # 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 class_decorator
+
+
+def segmented_accessor(elements, raw_segments, idx):
   """
   Returns a slice of elements corresponding to the idx-th segment.
 
     elements: a sliceable container (operands or results).
     raw_segments: an mlir.ir.Attribute, of DenseIntElements subclass containing
         sizes of the segments.
     idx: index of the segment.
   """
   segments = _cext.ir.DenseIntElementsAttr(raw_segments)
   start = sum(segments[i] for i in range(idx))
   end = start + segments[idx]
   return elements[start:end]
 
 
-def _equally_sized_accessor(elements, n_variadic, n_preceding_simple,
-                            n_preceding_variadic):
+def equally_sized_accessor(elements, n_variadic, n_preceding_simple,
   """
   Returns a starting position and a number of elements per variadic group
   assuming equally-sized groups and the given numbers of preceding groups.
 
     elements: a sequential container.
     n_variadic: the number of variadic groups in the container.
     n_preceding_simple: the number of non-variadic groups preceding the current
         group.
     n_preceding_variadic: the number of variadic groups preceding the current
         group.
   """
 
   total_variadic_length = len(elements) - n_variadic + 1
   # This should be enforced by the C++-side trait verifier.
   assert total_variadic_length % n_variadic == 0
 
   elements_per_group = total_variadic_length // n_variadic
   start = n_preceding_simple + n_preceding_variadic * elements_per_group
   return start, elements_per_group
 
-def _get_default_loc_context(location = None):
+
+def get_default_loc_context(location=None):
   """
   Returns a context in which the defaulted location is created. If the location
   is None, takes the current location from the stack, raises ValueError if there
   is no location on the stack.
   """
   if location is None:
     # Location.current raises ValueError if there is no current location.
     return _cext.ir.Location.current.context
   return location.context

mlir/lib/Bindings/Python/mlir/dialects/_linalg.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
+
+
+class StructuredOpMixin:
+  """All structured ops use the same mixin class."""
+
+  def __init__(self, inputs, outputs=(), results=(), loc=None, ip=None):
+    if outputs and results:
+      raise ValueError(
+          "Structured ops must have outputs or results, but not both.")
+    super().__init__(
+        self._ods_build_default(operands=[list(inputs),
+                                          list(outputs)],
+                                results=list(results),
+                                loc=loc,
+                                ip=ip))
+
+
+def select_opview_mixin(parent_opview_cls):
+  # TODO: This shouldn't be a heuristic: we should have a way to annotate
+  # the OpView to note that it is a structured op.
+  if ("__init__" not in parent_opview_cls.__dict__ and
+      hasattr(parent_opview_cls, "inputs") and
+      hasattr(parent_opview_cls, "outputs") and
+      hasattr(parent_opview_cls, "result_tensors")):
+    return StructuredOpMixin

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
 
 using namespace mlir;
 using namespace mlir::tblgen;
 
 /// File header and includes.
+///   {0} is the dialect namespace.
 constexpr const char *fileHeader = R"Py(
 # Autogenerated by mlir-tblgen; don't manually edit.
 
 from . import _cext as _ods_cext
-from . import _segmented_accessor as _ods_segmented_accessor, _equally_sized_accessor as _ods_equally_sized_accessor, _get_default_loc_context as _ods_get_default_loc_context
+from . import extend_opview_class as _ods_extend_opview_class, segmented_accessor as _ods_segmented_accessor, equally_sized_accessor as _ods_equally_sized_accessor, get_default_loc_context as _ods_get_default_loc_context
 _ods_ir = _ods_cext.ir
+
+try:
+  from . import _{0} as _ods_ext_module
+except ImportError:
+  _ods_ext_module = None
+
 )Py";
 
 /// Template for dialect class:
 ///   {0} is the dialect namespace.
 constexpr const char *dialectClassTemplate = R"Py(
 @_ods_cext.register_dialect
 class _Dialect(_ods_ir.Dialect):
   DIALECT_NAMESPACE = "{0}"
   pass
 
 )Py";
 
 /// Template for operation class:
 ///   {0} is the Python class name;
 ///   {1} is the operation name.
 constexpr const char *opClassTemplate = R"Py(
 @_ods_cext.register_operation(_Dialect)
+@_ods_extend_opview_class(_ods_ext_module)
 class {0}(_ods_ir.OpView):
   OPERATION_NAME = "{1}"
 )Py";
 
 /// Template for class level declarations of operand and result
 /// segment specs.
 ///   {0} is either "OPERAND" or "RESULT"
 ///   {1} is the segment spec

 /// registration requirements.
 static bool emitAllOps(const llvm::RecordKeeper &records, raw_ostream &os) {
   if (clDialectName.empty())
     llvm::PrintFatalError("dialect name not provided");
 
   AttributeClasses attributeClasses;
   constructAttributeMapping(records, attributeClasses);
 
-  os << fileHeader;
+  os << llvm::formatv(fileHeader, clDialectName.getValue());
   os << llvm::formatv(dialectClassTemplate, clDialectName.getValue());
   for (const llvm::Record *rec : records.getAllDerivedDefinitions("Op")) {
     Operator op(rec);
     if (op.getDialectName() == clDialectName.getValue())
       emitOpBindings(op, attributeClasses, os);
   }
   return false;
 }
 
 static GenRegistration
     genPythonBindings("gen-python-op-bindings",
                       "Generate Python bindings for MLIR Ops", &emitAllOps);