[mlir][linalg] Add an InitTensorOp python builder.

[stellaraccident]

Summary

• This has the API I want but I am not thrilled with the implementation. There are various things that could be improved both about the way that Python builders are mapped and the way the Linalg ops are factored to increase code sharing between C++/Python.
• Landing this as-is since it at least makes the InitTensorOp usable with the right API. Will refactor underneath in follow-ons.

[mehdi_amini]

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

10

Can you add a class doc? Even just to redirect to https://mlir.llvm.org/docs/Dialects/Linalg/#linalginit_tensor-mlirlinalginittensorop or at minimum mentioned that it maps to the linalg.init_tensor op.

[stellaraccident]

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

10

Done. (also realized that I forgot to document the extension mechanism. will do that in a followup)

[stellaraccident]

Comments and rebase.

mlir/lib/Bindings/Python/mlir/dialects/_linalg_ops_ext.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
 
+from typing import Optional, Sequence, Union
+from ..ir import *
+from ._ods_common import get_default_loc_context
+
+
+class InitTensorOp:
+  """Extends the linalg.init_tensor op."""
+
+  def __init__(self,
+               sizes: Union[Sequence[int], Sequence[Value]],
+               element_type: Type,
+               *,
+               loc=None,
+               ip=None):
+    """Constructs an `init_tensor` with either static or dynamic sizes."""
+    context = get_default_loc_context(loc)
+    operands = []
+    attributes = {}
+    # TODO: Refactor the InitTensorOp to take an element type attribute and
+    # then use normal result type inference, unifying the Python and C++ side
+    # with a standard mechanism (versus stashing that in builders).
+    if sizes and isinstance(sizes[0], Value):
+      # Dynamic sizes.
+      operands.extend(sizes)
+      static_size_ints = [-1] * len(sizes)
+      result_type = RankedTensorType.get(static_size_ints, element_type)
+    else:
+      # Static sizes.
+      result_type = RankedTensorType.get(sizes, element_type)
+      static_size_ints = sizes
+
+    index_type = IndexType.get(context)
+    attributes["static_sizes"] = ArrayAttr.get(
+        [IntegerAttr.get(index_type, s) for s in static_size_ints],
+        context=context)
+    op = self.build_generic(results=[result_type],
+                            operands=operands,
+                            attributes=attributes,
+                            loc=loc,
+                            ip=ip)
+    OpView.__init__(self, op)
+
 
 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.")

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

__all__ = [
     "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:
+    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.
 
-    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
+    # First try to resolve by name.
     try:
-      select_mixin = getattr(ext_module, "select_opview_mixin")
+      mixin_cls = getattr(ext_module, parent_opview_cls.__name__)
     except AttributeError:
-      # Try to default resolve it.
+      # Fall back to a select_opview_mixin hook.
       try:
-        mixin_cls = getattr(ext_module, parent_opview_cls.__name__)
+        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
 
     # Have a mixin_cls. Create an appropriate subclass.
     try:
 
       class LocalOpView(mixin_cls, parent_opview_cls):
         pass

mlir/test/Bindings/Python/dialects/linalg/opdsl/emit_structured_generic.py

with InsertionPoint(module.body):
     # CHECK-SAME: indexing_maps = [#[[$MAPA]], #[[$MAPB]], #[[$MAPC]]]
     # CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
     # CHECK-SAME: ins(%[[A]], %[[B]]
     # CHECK-SAME: outs(%[[INITC]]
 
     @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
                                  RankedTensorType.get((16, 8), f32))
     def test_matmul_mono(lhs, rhs):
-      # TODO: Enable outs inference and add sugar for InitTensorOp
-      # construction.
-      init_result = linalg.InitTensorOp(result=RankedTensorType.get((4, 8),
-                                                                    f32),
-                                        static_sizes=ArrayAttr.get([
-                                            IntegerAttr.get(IndexType.get(), 4),
-                                            IntegerAttr.get(IndexType.get(), 8)
-                                        ]),
-                                        sizes=[])
+      init_result = linalg.InitTensorOp([4, 8], f32)
       return matmul_mono(lhs, rhs, outs=[init_result.result])
 
     # CHECK-LABEL: @test_i8i8i32_matmul
     # CHECK:      ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: i32)
     # CHECK-NEXT:   %[[A_CAST:.+]] = sexti %[[A_ARG]] : i8 to i32
     # CHECK-NEXT:   %[[B_CAST:.+]] = sexti %[[B_ARG]] : i8 to i32
     # CHECK-NEXT:   %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i32
     # CHECK-NEXT:   %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32

mlir/test/Bindings/Python/dialects/linalg/ops.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 from mlir.dialects import builtin
 from mlir.dialects import linalg
 from mlir.dialects import std
 
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
+  return f
+
+
+# CHECK-LABEL: TEST: testInitTensor
+@run
+def testInitTensor():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    f32 = F32Type.get()
+    with InsertionPoint(module.body):
+      # CHECK-LABEL: func @static_sizes
+      # CHECK: %0 = linalg.init_tensor [3, 4] : tensor<3x4xf32>
+      @builtin.FuncOp.from_py_func()
+      def static_sizes():
+        return linalg.InitTensorOp([3, 4], f32)
+
+      # CHECK-LABEL: func @dynamic_sizes
+      # CHECK: %0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
+      @builtin.FuncOp.from_py_func(IndexType.get(), IndexType.get())
+      def dynamic_sizes(d0, d1):
+        return linalg.InitTensorOp([d0, d1], f32)
+
+      # CHECK-LABEL: func @zero_d
+      # CHECK: %0 = linalg.init_tensor [] : tensor<f32>
+      @builtin.FuncOp.from_py_func()
+      def zero_d():
+        return linalg.InitTensorOp([], f32)
+
+  print(module)
 
 
 # CHECK-LABEL: TEST: testStructuredOpOnTensors
+@run
 def testStructuredOpOnTensors():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     tensor_type = RankedTensorType.get((2, 3, 4), f32)
     with InsertionPoint(module.body):
       func = builtin.FuncOp(name="matmul_test",
                             type=FunctionType.get(
                                 inputs=[tensor_type, tensor_type],
                                 results=[tensor_type]))
       with InsertionPoint(func.add_entry_block()):
         lhs, rhs = func.entry_block.arguments
         result = linalg.MatmulOp([lhs, rhs], results=[tensor_type]).result
         std.ReturnOp([result])
 
   # CHECK: %[[R:.*]] = linalg.matmul ins(%arg0, %arg1 : tensor<2x3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
   print(module)
 
 
-run(testStructuredOpOnTensors)
-
-
 # CHECK-LABEL: TEST: testStructuredOpOnBuffers
+@run
 def testStructuredOpOnBuffers():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     memref_type = MemRefType.get((2, 3, 4), f32)
     with InsertionPoint(module.body):
       func = builtin.FuncOp(name="matmul_test",
                             type=FunctionType.get(
                                 inputs=[memref_type, memref_type, memref_type],
                                 results=[]))
       with InsertionPoint(func.add_entry_block()):
         lhs, rhs, result = func.entry_block.arguments
         linalg.MatmulOp([lhs, rhs], outputs=[result])
         std.ReturnOp([])
 
   # CHECK: linalg.matmul ins(%arg0, %arg1 : memref<2x3x4xf32>, memref<2x3x4xf32>) outs(%arg2 : memref<2x3x4xf32>)
   print(module)
-
-
-run(testStructuredOpOnBuffers)