[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

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

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

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

[nicolasvasilache]

Address comment.

[nicolasvasilache]

mlir/include/mlir-c/AffineMap.h

181

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/include/mlir-c/AffineMap.h

 
 /// Returns the affine map consisting of the most minor `numResults` results.
 /// Returns the null AffineMap if the `numResults` is equal to zero.
 /// Returns the `affineMap` if `numResults` is greater or equals to number of
 /// results of the given affine map.
 MLIR_CAPI_EXPORTED MlirAffineMap
 mlirAffineMapGetMinorSubMap(MlirAffineMap affineMap, intptr_t numResults);
 
+/// Returns the simplified affine map resulting from dropping the symbols that
+/// do not appear in any of the individual maps in `affineMaps`.
+/// Asserts that all maps in `affineMaps` are normalized to the same number of
+/// dims and symbols.
+/// Takes a callback `populateResult` to fill the `res` container with value
+/// `m` at entry `idx`. This allows returning without worrying about ownership
+/// considerations.
+MLIR_CAPI_EXPORTED void mlirAffineMapCompressUnusedSymbols(
+    MlirAffineMap *affineMaps, intptr_t size, void *result,
+    void (*populateResult)(void *res, intptr_t idx, MlirAffineMap m));
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_AFFINEMAP_H

mlir/include/mlir/IR/AffineMap.h

 };
 
 /// Simplifies an affine map by simplifying its underlying AffineExpr results.
 AffineMap simplifyAffineMap(AffineMap map);
 
 /// Drop the dims that are not used.
 AffineMap compressUnusedDims(AffineMap map);
 
+/// Drop the dims that are not used by any of the individual maps in `maps`.
+/// Asserts that all maps in `maps` are normalized to the same number of
+/// dims and symbols.
+SmallVector<AffineMap> compressUnusedDims(ArrayRef<AffineMap> maps);
+
 /// Drop the dims that are not listed in `unusedDims`.
 AffineMap compressDims(AffineMap map,
                        const llvm::SmallDenseSet<unsigned> &unusedDims);
 
 /// Drop the symbols that are not used.
 AffineMap compressUnusedSymbols(AffineMap map);
 
+/// Drop the symbols that are not used by any of the individual maps in `maps`.
+/// Asserts that all maps in `maps` are normalized to the same number of
+/// dims and symbols.
+SmallVector<AffineMap> compressUnusedSymbols(ArrayRef<AffineMap> maps);
+
 /// Drop the symbols that are not listed in `unusedSymbols`.
 AffineMap compressSymbols(AffineMap map,
                           const llvm::SmallDenseSet<unsigned> &unusedSymbols);
 
 /// Returns a map with the same dimension and symbol count as `map`, but whose
 /// results are the unique affine expressions of `map`.
 AffineMap removeDuplicateExprs(AffineMap map);
 

mlir/lib/Bindings/Python/IRAffine.cpp

py::class_<PyAffineMap>(m, "AffineMap")
            [](PyAffineMap &self) {
              PyPrintAccumulator printAccum;
              printAccum.parts.append("AffineMap(");
              mlirAffineMapPrint(self, printAccum.getCallback(),
                                 printAccum.getUserData());
              printAccum.parts.append(")");
              return printAccum.join();
            })
+      .def_static("compress_unused_symbols",
+                  [](py::list affineMaps, DefaultingPyMlirContext context) {
+                    SmallVector<MlirAffineMap> maps;
+                    pyListToVector<PyAffineMap, MlirAffineMap>(
+                        affineMaps, maps, "attempting to create an AffineMap");
+                    std::vector<MlirAffineMap> compressed(affineMaps.size());
+                    auto populate = [](void *result, intptr_t idx,
+                                       MlirAffineMap m) {
+                      static_cast<MlirAffineMap *>(result)[idx] = (m);
+                    };
+                    mlirAffineMapCompressUnusedSymbols(
+                        maps.data(), maps.size(), compressed.data(), populate);
+                    std::vector<PyAffineMap> res;
+                    for (auto m : compressed)
+                      res.push_back(PyAffineMap(context->getRef(), m));
+                    return res;
+                  })
       .def_property_readonly(
           "context",
           [](PyAffineMap &self) { return self.getContext().getObject(); },
           "Context that owns the Affine Map")
       .def(
           "dump", [](PyAffineMap &self) { mlirAffineMapDump(self); },
           kDumpDocstring)
       .def_static(

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])
+      [AffineMapAttr.get(am)
+       # TODO: linalg verification does not currently allow symbols.
+       # Compress them for now.
+       for am in AffineMap.compress_unused_symbols(op_config.indexing_maps, Context.current)])
   iterator_types_attr = ArrayAttr.get(
       [StringAttr.get(s) for s in op_config.iterator_types])
+  sparse_attr = ArrayAttr.get(
+      [BoolAttr.get(False) for s in list(ins) + list(outs) if isa(RankedTensorType, s.type)])
+  if len(sparse_attr) == 0:
+    sparse_attr = None
 
-  return (all_arg_defs, in_arg_defs, out_arg_defs, outs, out_types,
-          type_mapping, indexing_maps_attr, iterator_types_attr)
+  return (all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types,
+          type_mapping, indexing_maps_attr, iterator_types_attr, sparse_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 =   \
+  all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, \
+  type_mapping, indexing_maps_attr, iterator_types_attr, sparse_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.
+      sparse=sparse_attr)  # 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, \
-  type_mapping, indexing_maps_attr, iterator_types_attr =   \
+  all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, \
+  type_mapping, indexing_maps_attr, iterator_types_attr, sparse_attr =   \
      prepare_common_structured_op(op_config, *ins, outs = outs)
 
   # If we get here, there must exist a builtin class `op_class_name`.
   ctx = Context.current
   fully_qualified_name = 'linalg.' + op_name
   if (not ctx.is_registered_operation(fully_qualified_name) or
       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)
+  named_op = getattr(linalg, op_class_name)(ins, outs, result_types)
   linalgDialect = ctx.get_dialect_descriptor("linalg")
   fill_builtin_region(linalgDialect, named_op.operation)
+  # Note: mlir-linalg-ods-yaml-gen.cpp uses a special linalg.memoized_indexing_maps
+  # attribute that the non-yaml path does not. The non-yaml path hardcodes the 
+  # indexing_maps in C++ directly.
+  named_op.operation.attributes["linalg.memoized_indexing_maps"] = indexing_maps_attr
+  # iterator_types are hardcoded in C++ both in the yaml and non-yaml path.
 
-  if len(out_arg_defs) == 1:
+  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/lib/CAPI/IR/AffineMap.cpp

MlirAffineMap mlirAffineMapGetMajorSubMap(MlirAffineMap affineMap,
                                           intptr_t numResults) {
   return wrap(unwrap(affineMap).getMajorSubMap(numResults));
 }
 
 MlirAffineMap mlirAffineMapGetMinorSubMap(MlirAffineMap affineMap,
                                           intptr_t numResults) {
   return wrap(unwrap(affineMap).getMinorSubMap(numResults));
 }
+
+void mlirAffineMapCompressUnusedSymbols(
+    MlirAffineMap *affineMaps, intptr_t size, void *result,
+    void (*populateResult)(void *res, intptr_t idx, MlirAffineMap m)) {
+  SmallVector<AffineMap> maps;
+  for (intptr_t idx = 0; idx < size; ++idx)
+    maps.push_back(unwrap(affineMaps[idx]));
+  intptr_t idx = 0;
+  for (auto m : mlir::compressUnusedSymbols(maps))
+    populateResult(result, idx++, wrap(m));
+}

mlir/lib/IR/AffineMap.cpp

AffineMap mlir::compressUnusedDims(AffineMap map) {
   });
   llvm::SmallDenseSet<unsigned> unusedDims;
   for (unsigned d = 0, e = map.getNumDims(); d != e; ++d)
     if (!usedDims.contains(d))
       unusedDims.insert(d);
   return compressDims(map, unusedDims);
 }
 
+static SmallVector<AffineMap>
+compressUnusedImpl(ArrayRef<AffineMap> maps,
+                   llvm::function_ref<AffineMap(AffineMap)> compressionFun) {
+  if (maps.empty())
+    return SmallVector<AffineMap>();
+  SmallVector<AffineExpr> allExprs;
+  allExprs.reserve(maps.size() * maps.front().getNumResults());
+  unsigned numDims = maps.front().getNumDims(),
+           numSymbols = maps.front().getNumSymbols();
+  for (auto m : maps) {
+    assert(numDims == m.getNumDims() && numSymbols == m.getNumSymbols() &&
+           "expected maps with same num dims and symbols");
+    llvm::append_range(allExprs, m.getResults());
+  }
+  AffineMap unifiedMap = compressionFun(
+      AffineMap::get(numDims, numSymbols, allExprs, maps.front().getContext()));
+  unsigned unifiedNumDims = unifiedMap.getNumDims(),
+           unifiedNumSymbols = unifiedMap.getNumSymbols();
+  ArrayRef<AffineExpr> unifiedResults = unifiedMap.getResults();
+  SmallVector<AffineMap> res;
+  res.reserve(maps.size());
+  for (auto m : maps) {
+    res.push_back(AffineMap::get(unifiedNumDims, unifiedNumSymbols,
+                                 unifiedResults.take_front(m.getNumResults()),
+                                 m.getContext()));
+    unifiedResults = unifiedResults.drop_front(m.getNumResults());
+  }
+  return res;
+}
+
+SmallVector<AffineMap> mlir::compressUnusedDims(ArrayRef<AffineMap> maps) {
+  return compressUnusedImpl(maps,
+                            [](AffineMap m) { return compressUnusedDims(m); });
+}
+
 AffineMap
 mlir::compressSymbols(AffineMap map,
                       const llvm::SmallDenseSet<unsigned> &unusedSymbols) {
   unsigned numSymbols = 0;
   SmallVector<AffineExpr> symReplacements;
   symReplacements.reserve(map.getNumSymbols());
   MLIRContext *context = map.getContext();
   for (unsigned sym = 0, e = map.getNumSymbols(); sym < e; ++sym) {
AffineMap mlir::compressUnusedSymbols(AffineMap map) {
   });
   llvm::SmallDenseSet<unsigned> unusedSymbols;
   for (unsigned d = 0, e = map.getNumSymbols(); d != e; ++d)
     if (!usedSymbols.contains(d))
       unusedSymbols.insert(d);
   return compressSymbols(map, unusedSymbols);
 }
 
+SmallVector<AffineMap> mlir::compressUnusedSymbols(ArrayRef<AffineMap> maps) {
+  return compressUnusedImpl(
+      maps, [](AffineMap m) { return compressUnusedSymbols(m); });
+}
+
 AffineMap mlir::simplifyAffineMap(AffineMap map) {
   SmallVector<AffineExpr, 8> exprs;
   for (auto e : map.getResults()) {
     exprs.push_back(
         simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols()));
   }
   return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs,
                         map.getContext());

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

with Context() as ctx, Location.unknown():
   i8 = IntegerType.get_signless(8)
   i16 = IntegerType.get_signless(16)
   i32 = IntegerType.get_signless(32)
   with InsertionPoint(module.body):
 
     # Note that these all have the same indexing maps. We verify the first and
     # then do more permutation tests on casting and body generation
     # behavior.
-    # CHECK: #[[$MAPA:.+]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>
-    # CHECK: #[[$MAPB:.+]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)>
-    # CHECK: #[[$MAPC:.+]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>
+    # CHECK: #[[$MAPA:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+    # CHECK: #[[$MAPB:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+    # CHECK: #[[$MAPC:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
 
     # CHECK-LABEL: func @test_matmul_mono
     # CHECK-SAME:  %[[A:.+]]: tensor<4x16xf32>
     # CHECK-SAME: %[[B:.+]]: tensor<16x8xf32>
 
     # CHECK: %[[INITC:.+]] = linalg.init_tensor [4, 8] : tensor<4x8xf32>
     # CHECK: linalg.generic
     # CHECK-SAME: indexing_maps = [#[[$MAPA]], #[[$MAPB]], #[[$MAPC]]]

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

with Context() as ctx, Location.unknown():
     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)
         # First check the named form with custom format
         #      CHECK: linalg.matmul
+        #  CHECK-NOT: linalg.memoized_indexing_maps
         # CHECK-SAME:    ins(%{{.*}} : tensor<4x16xf32>, tensor<16x8xf32>)
         # CHECK-SAME:   outs(%{{.*}} : tensor<4x8xf32>)
         # CHECK-SAME:   -> tensor<4x8xf32>
         # CHECK-NEXT: return
         return linalg.matmul(lhs, rhs, outs=[init_result.result])
 
   print(module)
 
 # CHECK-LABEL: TEST: testNamedStructuredOpGenericForm
 @run
 def testNamedStructuredOpGenericForm():
   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"(%{{.*}})
         # CHECK-NEXT:  ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32):
         # CHECK-NEXT:    std.mulf{{.*}} (f32, f32) -> f32
         # CHECK-NEXT:    std.addf{{.*}} (f32, f32) -> f32
         # CHECK-NEXT:    linalg.yield{{.*}} (f32) -> ()
-        # CHECK-NEXT:    {operand_segment_sizes = dense<[2, 1]> : vector<2xi32>} : 
+        # CHECK-NEXT:    {linalg.memoized_indexing_maps{{.*}}operand_segment_sizes = dense<[2, 1]> : vector<2xi32>} : 
         # CHECK-SAME: (tensor<4x16xf32>, tensor<16x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32>
         return linalg.matmul(lhs, rhs, outs=[init_result.result])
 
   module.operation.print(print_generic_op_form=True)
 
 # CHECK-LABEL: TEST: testNamedStructuredAsGenericOp
 @run
 def testNamedStructuredAsGenericOp():

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>) -> ()
+
+  %c0 = constant 0 : index
+  %0 = memref.load %C[%c0, %c0] : memref<4x8xf32>
+
+  // TODO: FFI-based solution to allow testing and printing with python code.
+  return %0 : f32
+}
+"""
+
+def transform(module):
+  import mlir.conversions
+  import mlir.dialects.linalg.passes
+  import mlir.transforms
+
+  # TODO: Allow cloning functions from one module to another.
+  # Atm we have to resort to string concatenation.
+  mod = Module.parse(
+    str(module.operation.regions[0].blocks[0].operations[0].operation) +
+    boilerplate)
+  pm = PassManager.parse("func(convert-linalg-to-loops, convert-scf-to-std)," + 
+                         "convert-vector-to-llvm," + 
+                         "convert-std-to-llvm")
+  pm.run(mod)
+  return mod
+
+def test_builtin():
+  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):
+        linalg.matmul(lhs, rhs, outs=[out])
+    
+    execution_engine = ExecutionEngine(transform(module))
+
+    # TODO: FFI-based solution to allow testing and printing with python code.
+    # 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_builtin()
+
+def test_generic():
+  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):
+        linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)
+    
+    execution_engine = ExecutionEngine(transform(module))
+
+    # TODO: FFI-based solution to allow testing and printing with python code.
+    # 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_generic()