[mlir][Linalg] Allow calling named ops when available and make it the default.

[nicolasvasilache]

Formatting

[mehdi_amini]

Sweet!

mlir/lib/Bindings/Python/mlir/dialects/linalg/__init__.py

24

Strictly speaking, PyRepr -> YAML does not happen at compiler time at the moment.

52

Thanks for the effort to put an extensive doc!

[stellaraccident]

Nice! I think this needs a fair bit more iteration to get it all laid out right, but I like that everything is now unified. We can iterate from here.

mlir/lib/Bindings/Python/mlir/dialects/linalg/__init__.py

6

Strictly speaking, these are the backing OpView *classes* generated from the linalg tablegen definitions, including things like linalg.generic and those imported via the DSL/YAML path. These will all be named like GenericOp, MatmulOp, CopyOp, etc. They can be used to either construct the corresponding op in the IR, or during traversal, instances of these op classes will be returned.

21

Let's be consistent on named. "... then turn into the core C++ Op classes and Python OpView classes."

48

Note to self: once this lands, I want to look at all of the namespaces and terminology and see how we can better make it consistent. I suspect some extra explicitness in the core_named_ops versions (i.e. delineating them as part of the linalg named op namespace will let us remove this warning.

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/dsl.py

45

This should already be populated as self.model.cpp_op_name (which should likely be renamed cpp_class_name and be explicitly configurable at op definition time (right now it is done with some string/case munging).

69

Something like from .... import linalg as linalg_ops. (may have mis-counted number of dots). Agreed this should just be client in the parent namespace.

I am still vaguely trying to keep the namespace relocatable.

70

Nit: Prefer to add outer parens so it can reflow vs a line continuation \.

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/emitter.py

53

Note to self: I almost factored this as a class in anticipation of needing to carry a bunch of state like this but wanted to see how bad it was. Now I see and want to clean this up.

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

79

Can you add a TODO here: I believe this will be created without its region populated whereas using the standalone builder function below will populate it. I can unify that.

97

Note to self: This reads more like named "standalone builder functions" which then return an instance of the corresponding OpView class. I expect that if the mechanism is working, a further check on isinstance(created_op, linalg.MatmulOp) should return True. Want to verify that is working and then document this correspondance between standalone builder methods and OpView classes.

[nicolasvasilache]

mlir/lib/Bindings/Python/mlir/dialects/linalg/__init__.py

6

Not 100% sure how to phrase and incorporate all this, please improve and land as you see fit.

48

ack

52

ack

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/dsl.py

69

thanks!

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/emitter.py

53

ack, letting you do it as you see fit :)

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

79

yes, it's def. broken atm, will try to fix in the n+2 commit once the n+1 commit executes properly, unless you get to it first.

97

ack, that would be nice, leaving to you as a followup.

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

op = self.build_generic(results=[result_type],
                             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.")
     super().__init__(
         self.build_generic(results=list(results),
                            operands=[list(inputs), list(outputs)],
                            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/lib/Bindings/Python/mlir/dialects/linalg/__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
 
+# These are the backing OpView classes generated from the linalg tablegen
+# definitions following these steps:
+#   DSL -> YAML -> tblgen -> pytblgen -> build/.../_linalg_ops_gen.py.
 from .._linalg_ops_gen import *
+
+# These are the ground truth functions defined as:
+# ```
+#    @linalg_structured_op
+#    def matmul(A=TensorDef(T1, S.M, S.K),
+#               B=TensorDef(T2, S.K, S.N),
+#               C=TensorDef(U, S.M, S.N, output=True)):
+# ```
+# using the linalg-py eDSL.
+# The linalg-py eDSL builds a python representation (PyRepr) that is 
+# used in following ways:
+#  1. PyRepr -> YAML to generate the C++ and Python .td files. These
+#     then turn into the core C++ Op classes and Python OpView classes
+#     respectively (made available in _linalg_ops_gen). The generic OpView class 
+#     mechanism makes the C++ classes available to python through the CAPI.
+#     PyRepr -> YAML currently occurs before compiler compile time.
+#     The other steps in this category occur at compiler compile time.
+#  2. PyRepr -> linalg.core_named_ops calls: piggybacks on the 
+#     _linalg_ops_gen classes and the OpView mechanism to build IR at
+#     runtime in python:
+#       a. by default, the Named Op Form is emitted, e.g.:
+#          `linalg.matmul(lhs, rhs, outs=[out])` creates the following IR:
+#          ```
+#             %1 = linalg.matmul ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>) 
+#                               outs(%0 : tensor<4x8xf32>)
+#                  -> tensor<4x8xf32>   
+#          ```
+#       b. by setting emit_generic=True, the Generic Op Form is emitted, e.g.:
+#           `linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)` creates the following IR:
+#          ```
+#             %1 = linalg.generic {indexing_maps = [...], iterator_types = [...]} 
+#               ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>) 
+#              outs(%0 : tensor<4x8xf32>) {
+#               ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):  
+#                  ...
+#                  linalg.yield %3 : f32
+#             } -> tensor<4x8xf32>  
+#          ```
+#  3. PyRepr -> Runtime Custom Op definitions: directly generates a
+#     linalg.generic form like in 2.b.
+#     !!!WARNING!!!: if one creates a runtime custom op with the same name 
+#     as an existing core named op, step 2. will likely take precedence.
+#     TODO: guard against surprises and fail create Runtime Custom Ops with 
+#     the same name as existing Core Named Ops.
+from .opdsl.ops.core_named_ops import *

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/comprehension.py

 
 ContractionOpInterface = OpInterfaceDef("LinalgContractionOpInterface")
 
 
 class OpMetadataDef(YAMLObject):
   """Metadata about the op (generally not behavior impacting)."""
   yaml_tag = "!LinalgOpMetadata"
 
-  def __init__(self, name: str, cpp_op_name: Optional[str], doc: Optional[str]):
+  def __init__(self, name: str, cpp_class_name: Optional[str], doc: Optional[str]):
     self.name = name
-    self.cpp_op_name = cpp_op_name if cpp_op_name is not None else name
+    self.cpp_class_name = cpp_class_name if cpp_class_name is not None else name
     self.doc = doc
     self.implements = []  # type: List[OpInterfaceDef]
 
   def to_yaml_custom_dict(self):
     d = dict(
         name=self.name,
-        cpp_op_name=self.cpp_op_name,
+        cpp_class_name=self.cpp_class_name,
         doc=self.doc,
     )
     if self.implements:
       d["implements"] = [intr.cpp_name for intr in self.implements]
     return d
 
 
 class LinalgOpDef:
   """Definition of a linalg op."""
 
   def __init__(self,
                name: str,
-               cpp_op_name: Optional[str] = None,
+               cpp_class_name: Optional[str] = None,
                doc: Optional[str] = None):
-    self.metadata = OpMetadataDef(name=name, cpp_op_name=cpp_op_name, doc=doc)
+    self.metadata = OpMetadataDef(name=name, cpp_class_name=cpp_class_name, doc=doc)
     self.registered_tensors = dict()  # type: Dict[str, TensorDef]
     self.comprehensions = list()  # type: List[Comprehension]
     self._affine_state = AffineBuildState()
 
   @property
   def inputs(self) -> Sequence[TensorDef]:
     return [t for t in self.registered_tensors.values() if not t.output]
 
def tensor(self, name):
     """Gets a registered tensor by name."""
     try:
       return self.registered_tensors[name]
     except KeyError:
       raise KeyError(f"Tensor {name} is not registered")
 
   def __repr__(self):
     lines = [
-        f"LinalgOpDef({self.metadata.name} -> {self.metadata.cpp_op_name},"
+        f"LinalgOpDef({self.metadata.name} -> {self.metadata.cpp_class_name},"
     ]
     for name, tensor in self.registered_tensors.items():
       lines.append(f"  {tensor}")
     if self.comprehensions:
       lines[-1] += " {"
       for comprehension in self.comprehensions:
         lines.append(f"    {comprehension}")
       lines.append("}")
     return "\n".join(lines)

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/dsl.py

raise ValueError(
         "but none is set. Did you mean to call this in an op definition?")
 
 
 class DefinedOpCallable:
   """Callable that wraps any defined op function."""
 
   def __init__(self, op_name: str, model: LinalgOpDef):
     self.op_name = op_name
     self.model = model
 
-  def __call__(self, *args, emit_generic: bool = True, **kwargs):
+  def __call__(self, *args, emit_generic: bool = False, **kwargs):
     """Emits the corresponding op definition as IR.
 
     Most arguments are passed through to the underlying emitter. The following
     are interpreted here:
       emit_generic: Emits a generic form as appropriate (default True). If
         False, a named form is emitted (which must have been built in to the
         compiler).
     """
     op_configs = LinalgOpConfig.from_linalg_op_def(self.model,
                                                    context=ir.Context.current)
 
     if len(op_configs) != 1:
       # TODO: Support composite ops.
       raise NotImplementedError(
           f"Emission of composite linalg ops not supported: {op_configs}")
 
+    # TODO: this file should probably not be called dsl.py but rather is a client
+    # of the dsl.py.
+    from .... import linalg as linalg_ops
+    emit_generic = (emit_generic or 
+      (not self.model.metadata.cpp_class_name in linalg_ops.__dict__.keys()))
+
     op_config = op_configs[0]
     if op_config.structured_op:
       if emit_generic:
         return emit_generic_structured_op(op_config.structured_op, *args,
                                           **kwargs)
       else:
-        return emit_named_structured_op(op_config.structured_op, *args,
-                                        **kwargs)
+        return emit_named_structured_op(
+          op_config.structured_op, self.op_name,
+          self.model.metadata.cpp_class_name, *args, **kwargs)
 
     raise NotImplementedError(
         f"Emission of linalg op type not supported: {op_config}")
 
 
 def linalg_structured_op(dsl_func=None,
                          *,
                          op_name=None,
                          op_class_name=None) -> DefinedOpCallable:
   if dsl_func is None:
     # Curry the keyword args in for delayed application.
     return functools.partial(tc_def_op,
                              op_name=op_name,
                              op_class_name=op_class_name)
   # Determine default names by introspecting the function.
   if op_name is None:
     op_name = dsl_func.__name__
   if op_class_name is None:
     # Camel case it.
     op_class_name = f"{''.join(x.title() for x in op_name.split('_'))}Op"
 
   tc_model = LinalgOpDef(name=op_name,
-                         cpp_op_name=op_class_name,
+                         cpp_class_name=op_class_name,
                          doc=inspect.getdoc(dsl_func))
 
   # Extract arguments and TensorDefs from the signature.
   dsl_func_args = list()
   sig = inspect.signature(dsl_func)
   for param_name, param in sig.parameters.items():
     param_default = param.default
     if not isinstance(param_default, TensorDef):

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/emitter.py

 from .config import *
 
 __all__ = [
     "emit_generic_structured_op",
     "emit_named_structured_op",
 ]
 
 
-def emit_generic_structured_op(op_config: LinalgStructuredOpConfig,
+def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
                                  *ins: Value,
-                               outs: Value = ()):
+                                 outs: Value):
   all_arg_defs = op_config.ordered_tensor_args
   in_arg_defs = [arg for arg in all_arg_defs if arg.usage == "input"]
   out_arg_defs = [arg for arg in all_arg_defs if arg.usage == "output"]
 
   # Arity validation.
   if len(ins) != len(in_arg_defs):
     raise ValueError(f"Expected {len(in_arg_defs)} inputs but got "
                      f"{len(ins)} for {op_config}")
for arg_def, arg_element_type in zip(
     type_mapping[tv_name] = arg_element_type
 
   # Emit the generic op.
   # TODO: Support emission of pure memref form.
   indexing_maps_attr = ArrayAttr.get(
       [AffineMapAttr.get(am) for am in op_config.indexing_maps])
   iterator_types_attr = ArrayAttr.get(
       [StringAttr.get(s) for s in op_config.iterator_types])
+
+  return (all_arg_defs, in_arg_defs, out_arg_defs, outs, out_types,
+          type_mapping, indexing_maps_attr, iterator_types_attr)
+
+
+def emit_generic_structured_op(op_config: LinalgStructuredOpConfig,
+                               *ins: Value,
+                               outs: Value = ()):
+  all_arg_defs, in_arg_defs, out_arg_defs, outs, out_types, \
+  type_mapping, indexing_maps_attr, iterator_types_attr =   \
+     prepare_common_structured_op(op_config, *ins, outs = outs)
+
   generic_op = linalg.GenericOp(
       result_tensors=out_types,
       inputs=ins,
       outputs=outs,
       indexing_maps=indexing_maps_attr,
       iterator_types=iterator_types_attr,
       doc=None,  # TODO: Make optional.
       library_call=None,  # TODO: Make optional.
def emit_generic_structured_op(op_config: LinalgStructuredOpConfig,
 
   if len(out_arg_defs) == 1:
     return generic_op.result
   else:
     return generic_op.results
 
 
 def emit_named_structured_op(op_config: LinalgStructuredOpConfig,
+                             op_name: str,
+                             op_class_name: str,
                              *ins: Value,
                              outs: Value = ()):
+  all_arg_defs, in_arg_defs, out_arg_defs, outs, out_types, \
+  type_mapping, indexing_maps_attr, iterator_types_attr =   \
+     prepare_common_structured_op(op_config, *ins, outs = outs)
+
+  if not op_class_name in linalg.__dict__.keys():
     raise NotImplementedError(
-      f"Emission of named structured ops is not supported: {op_config}")
+        f"Unknown named op_name / op_class_name: {op_name} / {op_class_name}")
+
+  named_op = getattr(linalg, op_class_name)(ins, outs, out_types)
+  if len(out_arg_defs) == 1:
+    return named_op.result
+  else:
+    return named_op.results
 
 
 class _BodyBuilder:
   """Constructs a structured op body by evaluating assignments."""
 
   def __init__(self, type_mapping: Dict[str, Type],
                block_arg_mapping: Dict[str, Value]):
     self.type_mapping = type_mapping

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

with Context() as ctx, Location.unknown():
     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
+        # TODO: prperly hook up the region.
         linalg.MatmulOp([lhs, rhs], outputs=[result])
         std.ReturnOp([])
 
   # CHECK: linalg.matmul ins(%arg0, %arg1 : memref<2x3x4xf32>, memref<2x3x4xf32>) outs(%arg2 : memref<2x3x4xf32>)
   print(module)
+
+# CHECK-LABEL: TEST: testNamedStructuredOp
+@run
+def testNamedStructuredOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    f32 = F32Type.get()
+    with InsertionPoint(module.body):
+      @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
+                                   RankedTensorType.get((16, 8), f32))
+      def named_form(lhs, rhs):
+        init_result = linalg.InitTensorOp([4, 8], f32)
+        # CHECK: linalg.matmul
+        # TODO: prperly hook up the region.
+        return linalg.matmul(lhs, rhs, outs=[init_result.result])
+
+      @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
+                                   RankedTensorType.get((16, 8), f32))
+      def generic_form(lhs, rhs):
+        init_result = linalg.InitTensorOp([4, 8], f32)
+        # CHECK: linalg.generic
+        return linalg.matmul(lhs, rhs, outs=[init_result.result], emit_generic=True)
+
+  print(module)