[mlir][Python][Linalg] Add missing attributes to linalg ops

[nicolasvasilache]

Summary

This revision tightens up the handling of attributes for both named
and generic linalg ops.
To demonstrate the IR validity, a working e2e Linalg example is added.

[stellaraccident]

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

27

Oof. What is going on with the body not populating?

(Side note, you can just say str(module))

[nicolasvasilache]

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

27

Well I did not hook it up yet so there is not magic :) it will come.
I'm looking for the fastest DFS to execution first (at least in the ongoing hacks).

[nicolasvasilache]

Make example work.

[nicolasvasilache]

Cleanup.

[nicolasvasilache]

Cleanup.

[nicolasvasilache]

Step.

[nicolasvasilache]

Minor cleanups

[nicolasvasilache]

Fix test and rebase.

[nicolasvasilache]

@gysit @springerm FYI

[ftynse]

mlir/include/mlir-c/AffineMap.h

181 ↗

(On Diff #334448)

How about the first argument being void * instead? This will allow the caller to pass in any container they want, not necessarily a contiguous, and also additional data if desired.

mlir/lib/Bindings/Python/IRAffine.cpp

547–549 ↗

(On Diff #334448)

Nit: name shadowing almost got me here with compressed used twice

[mehdi_amini]

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

54

Can't all the convert- be inside the func(...) group?

[nicolasvasilache]

Address comments.

[nicolasvasilache]

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

54

nope, updated to maximally put stuff under func, the rest is module.

[ftynse]

mlir/include/mlir-c/AffineMap.h

181 ↗

(On Diff #334448)

MlirAffineMap *result should also become void *. The simplest user will just cast.

[nicolasvasilache]

Address comment.

[nicolasvasilache]

mlir/include/mlir-c/AffineMap.h

181 ↗

(On Diff #334448)

missed it, thanks!

[stellaraccident]

It doesn't look like any C-API or Python API unit tests of the new API were added? Would you mind adding something basic there?

[nicolasvasilache]

In D99430#2664203, @stellaraccident wrote:

It doesn't look like any C-API or Python API unit tests of the new API were added? Would you mind adding something basic there?

Landed in 48268aa0a9c18ab244b252b9617b91f6f5055c95.

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

raise ValueError(f"Expected {len(in_arg_defs)} inputs but got "
                      f"{len(ins)} for {op_config}")
   if outs and len(outs) != len(out_arg_defs):
     raise ValueError(f"Expected {len(out_arg_defs)} outputs but got "
                      f"{len(outs)} for {op_config}")
 
   outs, out_types = _infer_structured_outs(op_config, in_arg_defs, ins,
                                            out_arg_defs, outs)
 
+  # TODO: This hack only works on buffers
+  result_types = []
+  # result_types = [t for t in out_types if t == RankedTensorType.get(t.shape())]
+
   # Extract type vars for input/output based types.
   type_mapping = dict()  # type: Dict[str, Type]
   for arg_def, arg_element_type in zip(
       in_arg_defs + out_arg_defs,
       _get_shaped_element_types_from_values(*ins, *outs)):
     tv_name = arg_def.tensor_def.type_var.name
     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,
+  return (all_arg_defs, in_arg_defs, out_arg_defs, outs, result_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, \
+  all_arg_defs, in_arg_defs, out_arg_defs, outs, result_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,
+      result_tensors=result_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.
       sparse=BoolAttr.get(False))  # TODO: Make optional.
 
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, \
+  all_arg_defs, in_arg_defs, out_arg_defs, outs, result_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"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:
+  named_op = getattr(linalg, op_class_name)(ins, outs, result_types)
+  if len(result_types) == 1:
     return named_op.result
   else:
     return named_op.results
 
 
 class _BodyBuilder:
   """Constructs a structured op body by evaluating assignments."""
 

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

 # RUN: %PYTHON %s 2>&1 | FileCheck %s
 
 import gc, sys
 from mlir.ir import *
 from mlir.dialects import builtin
 from mlir.dialects import linalg
 from mlir.dialects import std
 from mlir.passmanager import *
 from mlir.execution_engine import *
 
-def lowerToLLVM(module):
+def transform(module):
+  # print(module)
+  # print(module.body.operations[0].print(print_generic_op_form=True))
+  # print(module.body.operations[0].entry_block.operations[0].regions[0])
+  
   import mlir.conversions
   import mlir.transforms
-  pm = PassManager.parse("print-op-stats")
+  pm = PassManager.parse("print-op-stats, print-op-stats")
   #pm = PassManager.parse("module(func(print-op-stats))")
+  #pm = PassManager.parse("convert-std-to-llvm")
   #pm = PassManager.parse("convert-linalg-to-loops")
-  pm.run(module)
+
+  # Roundtrip through strings for now because lianlg IR is invalid atm but 
+  # prints correctly.
+
+  print(module)
+  pm.run(Module.parse(module.__str__()))
+  # pm.run(module)
+  
   return module
 
 def test():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     memref_type = MemRefType.get((2, 3, 4), f32)
     with InsertionPoint(module.body):
-      # @builtin.FuncOp.from_py_func(MemRefType.get((4, 16), f32),
-      #                              MemRefType.get((16, 8), f32),
-      #                              MemRefType.get((4, 8), f32))
-      # def matmul_on_buffers(lhs, rhs, out):
-      #   linalg.matmul(lhs, rhs, outs=[out])
-
-      @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
-                                   RankedTensorType.get((16, 8), f32))
-      def named_form(lhs, rhs):
-        return linalg.matmul(lhs, rhs, outs=[rhs], emit_generic=True)
+      @builtin.FuncOp.from_py_func(MemRefType.get((4, 16), f32),
+                                   MemRefType.get((16, 8), f32),
+                                   MemRefType.get((4, 8), f32))
+      def matmul_on_buffers(lhs, rhs, out):
+        linalg.matmul(lhs, rhs, outs=[out])
+
+      # @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
+      #                              RankedTensorType.get((16, 8), f32))
+      # def named_form(lhs, rhs):
+      #   return linalg.matmul(lhs, rhs, outs=[rhs], emit_generic=False)
     
-    print(module)
-    print(lowerToLLVM(module))
+    transform(module)
 
     # execution_engine = ExecutionEngine(lowerToLLVM(module))
     # # Prepare arguments: two input floats and one result.
     # # Arguments must be passed as pointers.
     # c_float_p = ctypes.c_float * 1
     # arg0 = c_float_p(42.)
     # arg1 = c_float_p(2.)
     # res = c_float_p(-1.)
     # execution_engine.invoke("add", arg0, arg1, res)
     # # CHECK: 42.0 + 2.0 = 44.0
     # print("{0} + {1} = {2}".format(arg0[0], arg1[0], res[0]))
 
 test()
 No newline at end of file