[mlir][Linalg][Python] Create the body of builtin named Linalg ops

[nicolasvasilache]

Summary

This revision adds support to properly add the body of registered builtin named linalg ops.
At this time, indexing_map and iterator_type support is still missing so the op is not executable yet.

[nicolasvasilache]

Step.

[nicolasvasilache]

Step.

[nicolasvasilache]

Step.

[nicolasvasilache]

Step.

[nicolasvasilache]

Step.

[nicolasvasilache]

Step.

[nicolasvasilache]

@stellaraccident @mehdi_amini @ftynse

This is still WIP as I am having some issues around:

# TODO: this fails with error:
 # dialect has no registered type printing hook
 # UNREACHABLE executed at /usr/local/google/home/ntv/github/llvm-project/mlir/include/mlir/IR/Dialect.h:103!
 print("AAA")
 linalgDialect = Context.current.get_dialect_descriptor("linalg")
 named_op.operation.print(print_generic_op_form=True)

as well as:

#if 0
  // TODO: this fails with op not castable to LinalgOp
  auto linalgOp = cast<LinalgOp>(op);

If you see what I'm doing wrong that would be super useful.

Thanks!

[mehdi_amini]

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

108

This error shows up when there is an MLIR type that belongs to a dialect but the dialect does not implement support for custom types.
Are there linalg types? Do you see how what you're doing around all this could lead to a type that belongs to linalg somehow?

[mehdi_amini]

mlir/include/mlir-c/Dialect/Linalg.h

21

I think the core binding could expose "mlirIsRegisteredOp(MlirContext ctx, MlirStringRef name) and you could use is with linalg.xyz` directly.

[stellaraccident]

mlir/include/mlir-c/Dialect/Linalg.h

21

Public CAPI functions must be annotated with MLIR_CAPI_EXPORTED so that they have default (public/exported) visibility.

mlir/lib/Bindings/Python/CMakeLists.txt

81

This being necessary is a symptom of something having gone wrong. By linking the static lib here, you are inadvertently crossing a shared library boundary that gets the layering messed up (the python extension must only depend on the shared library that exports the API -- at least until a better shared library scheme lands upstream).

The reason you were feeling the need to do this was because you didn't declare your new CAPI functions with MLIR_CAPI_EXPORTED, and they were therefore being hidden (not exported) from the MLIRPublicAPI.so shared library. Removing this making the change noted in Linalg.h fixes it.

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

108

It also shows up when certain link time invariants are violated such that the TypeID singletons that establish equality of the dialects are not in fact singleton.

[stellaraccident]

Assuming you take/incorporate https://reviews.llvm.org/D99638, left some comments on structure. Once you get things landed, I would like to take a pass through and do some renames now that we can see the whole namespace. Some things are needlessly awkward but I'd rather fix them as an NFC.

mlir/include/mlir-c/Dialect/Linalg.h

21

Added in: https://reviews.llvm.org/D99638

Nicolas, if off hours and it looks good/helps, feel free to land it and/or incorporate.

mlir/lib/Bindings/Python/MainModule.cpp

231

Can I suggest a different organization/naming of things?

I think that the _mlir cext module should mirror the public package structure where possible. Could we do something like:

auto dialectsModule = m.def_submodule("dialects");
auto linalgModule = dialectsModule.def_submodule("linalg");
populateDialectsLinalgSubmodule(linalgModule);

Then, can you also rename LinalgBuiltin.cpp (and header) -> DialectLinalg.cpp. There is still a slight inconsistency in that we call it dialects in Python and dialect in C++, but at least it is consistent within language.

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

13

I think you don't need this anymore with the new API for checking op registration.

71

Any reason why not just if emit_generic:?

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

8

Given the name collisions, for this could we do one of:

1. Import just what you need by name: from _mlir.dialect.linalg import fill_builtin_region
2. Add to linalg/__init__.py: from .._cext.dialect.linalg import * to bring the native API into the python side API.

[nicolasvasilache]

Address review.

[nicolasvasilache]

drop spurious print

[nicolasvasilache]

Rebase.

[stellaraccident]

Nice! Looks good in principle but need to look closer at the LinalgTypes code in the morning.

mlir/lib/Bindings/Python/mlir/linalg_builtin.py

2

I think you don't need this file anymore?

[ftynse]

Looks good. It would be nice to have DialectLinalg in a separate extension and generally start structuring the bindings more than the current flat representation, e.g. by adding lib/Bindings/Python/Dialect/Linalg. This could be a separate restructuring commit though.

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

41

Nit: do we need std::function (storing the function instance) or would llvm::function_ref suffice (a reference to a free or static member function)?

mlir/lib/Bindings/Python/CMakeLists.txt

72

I'm not sure this should go to CoreBindings. Can't this be a separate extension? Is it the layering problem referred to below?

mlir/lib/CAPI/Dialect/Linalg.cpp

34

Nit: this isn't necessary, the builder is created in this function and doesn't escape.

[nicolasvasilache]

mlir/lib/Bindings/Python/CMakeLists.txt

72

yes, this was a source of pain and much deeper understanding of the shared_lib, RTTI etc is required than what I have.
The thinking is that these will be future layering improvements driven by @stellaraccident and yourself :)

[nicolasvasilache]

Address review.

[stellaraccident]

mlir/lib/Bindings/Python/CMakeLists.txt

72

We should start factoring the directory structure we want but we need some non trivial build system work to get the dynamic libraries layered right to really separate things. Without a more principled library layout for the core project, we don't really get anything by just separating the leaves -- and in fact, there are a bunch of really fragile linkage issues that emerge unless if we layer it right. The way it is now has the one advantage that it can work on all platforms and all linkage modes of the project (static, shared, dylib). That starts to decay quickly without some careful work. Happy to elaborate more in another forum, as I've got a lot of state on this...

mlir/include/mlir-c/Dialect/Linalg.h

 
 #include "mlir-c/Registration.h"
 #include "mlir-c/Support.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+/// Return true if `name` is the name of a builtin named Linalg op.
+bool mlirLinalgIsBuiltinNamedOp(MlirDialect linalgDialect, MlirStringRef name);
+
+/// Apply the special region builder for the builtin named Linalg op.
+/// Assert that `op` is a builtin named Linalg op.
+void mlirLinalgFillBuiltinNamedOpRegion(MlirDialect linalgDialect,
+                                        MlirOperation op);
+
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Linalg, linalg);
 
 #ifdef __cplusplus
 }
 #endif
 
 #include "mlir/Dialect/Linalg/Passes.capi.h.inc"
 
 #endif // MLIR_C_DIALECT_LINALG_H

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

let description = [{
     Document](https://mlir.llvm.org/docs/Rationale/RationaleLinalgDialect) are
     are also available and should be read first before going in the details of
     the op semantics.
   }];
   let cppNamespace = "::mlir::linalg";
   let dependentDialects = [
     "AffineDialect", "StandardOpsDialect", "tensor::TensorDialect"
   ];
+  let extraClassDeclaration = [{
+    using RegionBuilderFunType = std::function<void(Block &, ValueRange)>;
+    RegionBuilderFunType getRegionBuilder(StringRef name) {
+      return namedStructuredOpRegionBuilders.lookup(name);
+    }
+    private:
+      llvm::StringMap<RegionBuilderFunType> namedStructuredOpRegionBuilders;
+  }];
 }
 
 // Whether a type is a RangeType.
 def LinalgIsRangeTypePred : CPred<"$_self.isa<RangeType>()">;
 def Range : DialectType<Linalg_Dialect, LinalgIsRangeTypePred, "range">;
 
 #endif // LINALG_BASE

mlir/include/mlir/Dialect/Linalg/IR/LinalgTypes.h

 //===- LinalgTypes.h - Linalg Types ---------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_LINALGTYPES_H_
 #define MLIR_DIALECT_LINALG_LINALGTYPES_H_
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Types.h"
+#include "llvm/ADT/StringMap.h"
 
 #include "mlir/Dialect/Linalg/IR/LinalgOpsDialect.h.inc"
 
 namespace mlir {
 class MLIRContext;
 
 namespace linalg {
 

mlir/lib/Bindings/Python/CMakeLists.txt

 
 ################################################################################
 # Build core python extension
 ################################################################################
 add_mlir_python_extension(MLIRCoreBindingsPythonExtension _mlir
   INSTALL_DIR
     python
   SOURCES
+    LinalgBuiltin.cpp
     MainModule.cpp
     IRAffine.cpp
     IRAttributes.cpp
     IRCore.cpp
     IRTypes.cpp
     PybindUtils.cpp
     Pass.cpp
     ExecutionEngine.cpp
+  LINK_LIBS
+    MLIRCAPILinalg
 )
 add_dependencies(MLIRBindingsPythonExtension MLIRCoreBindingsPythonExtension)
 
 # Note that we copy from the source tree just like for headers because
 # it will not be polluted with py_cache runtime artifacts (from testing and
 # such).
 install(
   DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/mlir

mlir/lib/Bindings/Python/LinalgBuiltin.h

+//===- LinalgBuiltin.h - Linalg builtin submodule of pybind module --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_BINDINGS_PYTHON_LINALGBUILTIN_H
+#define MLIR_BINDINGS_PYTHON_LINALGBUILTIN_H
+
+#include "PybindUtils.h"
+
+namespace mlir {
+namespace python {
+
+void populateLinalgBuiltinSubmodule(pybind11::module &m);
+
+} // namespace python
+} // namespace mlir
+
+#endif // MLIR_BINDINGS_PYTHON_LINALGBUILTIN_H

mlir/lib/Bindings/Python/LinalgBuiltin.cpp

+//===- LinalgBuiltin.cpp - Pybind module for Linalg builting API support --===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "IRModule.h"
+#include "mlir-c/Dialect/Linalg.h"
+#include "mlir-c/IR.h"
+
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+using namespace mlir;
+using namespace mlir::python;
+
+namespace mlir {
+namespace python {
+
+void populateLinalgBuiltinSubmodule(py::module &m) {
+  m.def(
+      "is_builtin",
+      [](PyDialectDescriptor &dialect, std::string opName) {
+        return mlirLinalgIsBuiltinNamedOp(
+            dialect.get(), mlirStringRefCreateFromCString(opName.c_str()));
+      },
+      py::arg("dialect"), py::arg("opname"),
+      "Return true if `name` is the name of a builtin named Linalg op.");
+  m.def(
+      "fill_builtin_region",
+      [](PyDialectDescriptor &dialect, PyOperation &op) {
+        return mlirLinalgFillBuiltinNamedOpRegion(dialect.get(), op.get());
+      },
+      py::arg("dialect"), py::arg("op"),
+      "Fill the region for `op`, which is assumed to be a builtin named Linalg "
+      "op.");
+}
+
+} // namespace python
+} // namespace mlir

mlir/lib/Bindings/Python/MainModule.cpp

 //===- MainModule.cpp - Main pybind module --------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include <tuple>
 
 #include "PybindUtils.h"
 
 #include "ExecutionEngine.h"
 #include "Globals.h"
 #include "IRModule.h"
+#include "LinalgBuiltin.h"
 #include "Pass.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 // -----------------------------------------------------------------------------
 // PyGlobals
PYBIND11_MODULE(_mlir, m) {
   auto passModule =
       m.def_submodule("passmanager", "MLIR Pass Management Bindings");
   populatePassManagerSubmodule(passModule);
 
   // Define and populate ExecutionEngine submodule.
   auto executionEngineModule =
       m.def_submodule("execution_engine", "MLIR JIT Execution Engine");
   populateExecutionEngineSubmodule(executionEngineModule);
+
+  // Define and populate Linalg submodule.
+  auto linalgBuiltinModule =
+      m.def_submodule("linalg_builtin", "MLIR Linalg Builtin Helper Bindings");
+  populateLinalgBuiltinSubmodule(linalgBuiltinModule);
 }

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/dsl.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 Dict, List
 
 from contextlib import contextmanager
 import functools
 import inspect
 import threading
 
 from mlir import ir
+from _mlir import linalg_builtin
 from .comprehension import *
 from .config import *
 from .emitter import *
 
 _CONTEXT = threading.local()
 
 
 @contextmanager
def __call__(self, *args, emit_generic: bool = False, **kwargs):
     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()))
+    linalgDialect = ir.Context.current.get_dialect_descriptor("linalg")
+    fully_qualified_name = 'linalg.' + self.op_name
+    emit_generic = (emit_generic or not linalg_builtin.is_builtin(linalgDialect, fully_qualified_name))
 
     op_config = op_configs[0]
     if op_config.structured_op:
-      if emit_generic:
+      if emit_generic is True:
         return emit_generic_structured_op(op_config.structured_op, *args,
                                           **kwargs)
       else:
         return emit_named_structured_op(
           op_config.structured_op, self.op_name,
           self.model.metadata.cpp_class_name, *args, **kwargs)
 
     raise NotImplementedError(

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/emitter.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 Dict, Sequence
 
 from mlir.ir import *
+from _mlir import linalg_builtin
 from mlir.dialects import linalg
 from mlir.dialects import std
 
 from .scalar_expr import *
 from .config import *
 
 __all__ = [
     "emit_generic_structured_op",
     "emit_named_structured_op",
 ]
 
-
 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.
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 we get here, there must exist a builtin class `op_class_name`.
   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)
+
+  # TODO: this fails with error:
+  # dialect has no registered type printing hook
+  # UNREACHABLE executed at /usr/local/google/home/ntv/github/llvm-project/mlir/include/mlir/IR/Dialect.h:103!
+  print("AAA")
+  linalgDialect = Context.current.get_dialect_descriptor("linalg")
+  named_op.operation.print(print_generic_op_form=True)
+  linalg_builtin.fill_builtin_region(linalgDialect, named_op.operation)
+  print("BBB")
+  named_op.operation.print(print_generic_op_form=True)
+  print("CCC")
+
   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."""

mlir/lib/Bindings/Python/mlir/linalg_builtin.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
+
+# Simply a wrapper around the extension module of the same name.
+from ._cext_loader import _reexport_cext
+_reexport_cext("linalg_builtin", __name__)
+del _reexport_cext

mlir/lib/CAPI/Dialect/Linalg.cpp

 //===- Linalg.cpp - C Interface for Linalg dialect ------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir-c/Dialect/Linalg.h"
 #include "mlir/CAPI/Registration.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 
-MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Linalg, linalg,
-                                      mlir::linalg::LinalgDialect)
+using namespace mlir;
+using namespace mlir::linalg;
+
+/// Return true if `name` is the name of a builtin named Linalg op.
+bool mlirLinalgIsBuiltinNamedOp(MlirDialect linalgDialect, MlirStringRef name) {
+  return static_cast<LinalgDialect *>(unwrap(linalgDialect))
+             ->getRegionBuilder(unwrap(name)) != nullptr;
+}
+
+/// Apply the special region builder for the builtin named Linalg op.
+/// Assert that `op` is a builtin named Linalg op.
+void mlirLinalgFillBuiltinNamedOpRegion(MlirDialect linalgDialect,
+                                        MlirOperation mlirOp) {
+  Operation *op = unwrap(mlirOp);
+  LinalgDialect::RegionBuilderFunType fun =
+      static_cast<LinalgDialect *>(unwrap(linalgDialect))
+          ->getRegionBuilder(op->getName().getStringRef());
+  assert(fun && "Expected a builtin named Linalg op.");
+  assert(op->getNumRegions() == 1 && "Expected Linalg op with 1 region");
+  assert(op->getRegion(0).getBlocks().empty() &&
+         "Expected Linalg op with 0 blocks");
+  // TODO: allow captures.
+  SmallVector<Type, 8> argTypes;
+#if 0
+  // TODO: this fails with op not castable to LinalgOp
+  auto linalgOp = cast<LinalgOp>(op);
+  for (auto t : linalgOp.getShapedOperandTypes())
+    argTypes.push_back(getElementTypeOrSelf(t));
+#else
+  for (auto t : op->getOperandTypes()) {
+    // TODO: this fails with error:
+    // dialect has no registered type printing hook
+    t.dump();
+    getElementTypeOrSelf(t).dump();
+    argTypes.push_back(getElementTypeOrSelf(t));
+  }
+#endif
+  OpBuilder b(op->getContext());
+  OpBuilder::InsertionGuard guard(b);
+  Region &region = op->getRegion(0);
+  Block *body = b.createBlock(&region, /*insertPt=*/{}, argTypes);
+  fun(*body, ValueRange{});
+}
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Linalg, linalg, LinalgDialect)

mlir/lib/Dialect/Linalg/IR/LinalgTypes.cpp

 };
 
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // LinalgDialect
 //===----------------------------------------------------------------------===//
 
+/// Trait to check if T provides a `regionBuilder` method.
+template <typename T, typename... Args>
+using has_region_builder = decltype(T::regionBuilder);
+template <typename T>
+using detect_has_region_builder = llvm::is_detected<has_region_builder, T>;
+
+/// SFINAE helper for single C++ class without a `regionBuilder` method (e.g.
+/// an OpInterface).
+template <typename OpType, typename = std::enable_if_t<
+                               !detect_has_region_builder<OpType>::value>>
+void addNamedOpBuilderImpl(
+    llvm::StringMap<LinalgDialect::RegionBuilderFunType> &map) {
+  // Do nothing.
+}
+
+template <typename OpType,
+          typename = std::enable_if_t<detect_has_region_builder<OpType>::value>,
+          typename = void>
+void addNamedOpBuilderImpl(
+    llvm::StringMap<LinalgDialect::RegionBuilderFunType> &map) {
+  map.insert(std::make_pair(
+      OpType::getOperationName(),
+      static_cast<LinalgDialect::RegionBuilderFunType>(OpType::regionBuilder)));
+}
+
+template <typename... OpTypes>
+void addNamedOpBuilders(
+    llvm::StringMap<LinalgDialect::RegionBuilderFunType> &map) {
+  (void)std::initializer_list<int>{0,
+                                   (addNamedOpBuilderImpl<OpTypes>(map), 0)...};
+}
+
 void mlir::linalg::LinalgDialect::initialize() {
   addTypes<RangeType>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc"
       >();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgSparseOps.cpp.inc"
       >();
 
+  // Fill the Linalg-specific OpName to RegionBuilder map.
+  addNamedOpBuilders<
+#define GET_OP_LIST
+#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
+      >(namedStructuredOpRegionBuilders);
+
   addInterfaces<LinalgInlinerInterface>();
 }
 
 Type mlir::linalg::LinalgDialect::parseType(DialectAsmParser &parser) const {
   // Parse the main keyword for the type.
   StringRef keyword;
   if (parser.parseKeyword(&keyword))
     return Type();

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
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.
+        # TODO: properly 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)
 
+  module.operation.print(print_generic_op_form=True)
   print(module)