[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

55

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

[nicolasvasilache]

Address comments.

[nicolasvasilache]

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

55

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

 from .scalar_expr import *
 from .config import *
 
 __all__ = [
     "emit_generic_structured_op",
     "emit_named_structured_op",
 ]
 
+def isa(cls : Type, ty : Type):
+  try:
+    cls(ty)
+    return True
+  except ValueError:
+    return False
 
 def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
                                  *ins: 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}")
   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)
 
+  result_types = [t for t in out_types if isa(RankedTensorType, t)]
+
   # 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.
 
   # Construct the body.
   block_arg_names = _get_tensor_def_names(*in_arg_defs, *out_arg_defs)
   block_arg_types = _get_shaped_element_types_from_values(*ins, *outs)
   block = generic_op.regions[0].blocks.append(*block_arg_types)
   block_arg_mapping = dict(zip(block_arg_names, block.arguments))
   with InsertionPoint(block):
     body_builder = _BodyBuilder(type_mapping, block_arg_mapping)
     for assignment in op_config.assignments:
       body_builder.assign(assignment)
     body_builder.yield_outputs(*_get_tensor_def_names(*out_arg_defs))
 
-  if len(out_arg_defs) == 1:
+  if len(result_types) == 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, \
+  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 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 *
+
+# Log everything to stderr and flush so that we have a unified stream to match
+# errors/info emitted by MLIR to stderr.
+def log(*args):
+  print(*args, file=sys.stderr)
+  sys.stderr.flush()
+
+boilerplate = """
+func @main() -> f32 attributes {llvm.emit_c_interface} {
+  %v0 = constant 0.0 : f32
+  %v1 = constant 1.0 : f32
+  %v2 = constant 2.0 : f32
+
+  %A = memref.alloc() : memref<4x16xf32>
+  %B = memref.alloc() : memref<16x8xf32>
+  %C = memref.alloc() : memref<4x8xf32>
+  linalg.fill(%A, %v1) : memref<4x16xf32>, f32
+  linalg.fill(%B, %v2) : memref<16x8xf32>, f32
+  linalg.fill(%C, %v0) : memref<4x8xf32>, f32
+
+  call @matmul_on_buffers(%A, %B, %C) : 
+    (memref<4x16xf32>, memref<16x8xf32>, memref<4x8xf32>) -> ()
+
+  // TODO: Link stuff in with python.
+  // %res = memref.cast %C: memref<4x8xf32> to memref<*xf32>
+  // call @print_memref_f32(%res) : (memref<*xf32>) -> ()
+
+  %c0 = constant 0 : index
+  %0 = memref.load %C[%c0, %c0] : memref<4x8xf32>
+  // TODO: Link stuff in with python.
+  // vector.print %0: f32
+
+  return %0 : f32
+}
+
+// TODO: Link stuff in with python.
+// func private @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }
+"""
+
+def transform(module):
+  import mlir.conversions
+  import mlir.dialects.linalg.passes
+  import mlir.transforms
+  pm = PassManager.parse("func(convert-linalg-to-loops)," +
+                         "convert-scf-to-std," + 
+                         "convert-vector-to-llvm," + 
+                         "convert-std-to-llvm")
+  # TODO: For now, roundtrip through strings for now because linalg IR is invalid
+  # atm but prints correctly.
+  # TODO: invalid linalg IR seems to also break the ability to print the first 
+  # operation, so we hack the string for now.
+  # print(module.body.operations[0])
+  # print(module.body.operations[0].entry_block)
+  program = str(module)[10:-3] + boilerplate
+  mod = Module.parse(program)
+  pm.run(mod)
+  # print(mod)  
+  return mod
+
+def test():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    f32 = F32Type.get()
+    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):
+        # TODO: body is not populated.
+        linalg.matmul(lhs, rhs, outs=[out])
+        # TODO: This does not work because the indexing_map symbols have not been 
+        # simplified away and the linalg.generic fails verification.
+        # linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)
+    
+    execution_engine = ExecutionEngine(transform(module))
+
+    # Prepare arguments: one result f32.
+    # Arguments must be passed as pointers.
+    c_float_p = ctypes.c_float * 1
+    res = c_float_p(-1.)
+    execution_engine.invoke("main", res)
+
+    log('RESULT: ', res[0])
+    # CHECK: RESULT: 32.0
+
+test()