[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 ↗

(On Diff #334266)

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

80

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 ↗

(On Diff #334266)

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 ↗

(On Diff #334266)

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 ↗

(On Diff #334266)

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

71 ↗

(On Diff #334266)

Any reason why not just if emit_generic:?

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

8 ↗

(On Diff #334266)

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

1 ↗

(On Diff #334364)

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

51

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 &, ArrayRef<Value>)>;
+    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
     MainModule.cpp
     IRAffine.cpp
     IRAttributes.cpp
     IRCore.cpp
     IRTypes.cpp
     PybindUtils.cpp
     Pass.cpp
     ExecutionEngine.cpp
 )
 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
   DESTINATION python
add_llvm_install_targets(
     install-MLIRBindingsPythonDialects
     DEPENDS MLIRBindingsPythonSources
     COMPONENT MLIRBindingsPythonDialects)
 endif()
 
 add_subdirectory(Transforms)
 add_subdirectory(Conversions)
 
+add_mlir_python_extension(MLIRLinalgBindingsPythonExtension _mlirLinalg
+  INSTALL_DIR
+    python
+  SOURCES
+    Linalg.cpp
+)
+add_dependencies(MLIRBindingsPythonExtension MLIRLinalgBindingsPythonExtension)
+
 add_mlir_python_extension(MLIRLinalgPassesBindingsPythonExtension _mlirLinalgPasses
   INSTALL_DIR
     python
   SOURCES
     LinalgPasses.cpp
 )
 add_dependencies(MLIRBindingsPythonExtension MLIRLinalgPassesBindingsPythonExtension)

mlir/lib/Bindings/Python/Linalg.cpp

+//===- Linalg.cpp - Pybind module for Linalg APIs ------------------------===//
+//
+// 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 "IRModule.h"
+#include "mlir-c/IR.h"
+
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+using namespace mlir;
+using namespace mlir::python;
+
+// -----------------------------------------------------------------------------
+// Module initialization.
+// -----------------------------------------------------------------------------
+
+PYBIND11_MODULE(_mlirLinalg, m) {
+  m.doc() = "MLIR Linalg Dialect APIs";
+
+  m.def(
+      "is_builtin",
+      [](PyDialect &dialect, std::string opName) {
+        return mlirLinalgIsBuiltinNamedOp(
+            py::cast<PyDialectDescriptor>(dialect.getDescriptor()).get(),
+            mlirStringRefCreateFromCString(opName.c_str()));
+      },
+      py::arg("dialect"), py::arg("opname"),
+      "Return true if `name` is the name of a builtin named Linalg op.");
+}

mlir/lib/Bindings/Python/mlir/dialects/linalg/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
+
+from ...._cext_loader import _load_extension
+_cextLinalgPasses = _load_extension("_mlirLinalg")

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"
 
+/// Return true if `name` is the name of a builtin named Linalg op.
+bool mlirLinalgIsBuiltinNamedOp(MlirDialect linalgDialect, MlirStringRef name) {
+  return static_cast<mlir::linalg::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 op) {
+  mlir::linalg::LinalgDialect::RegionBuilderFunType fun =
+      static_cast<mlir::linalg::LinalgDialect *>(unwrap(linalgDialect))
+          ->getRegionBuilder(unwrap(op)->getName().getStringRef());
+  assert(fun && "Expected a builtin named Linalg op.");
+}
+
 MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Linalg, linalg,
                                       mlir::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(OpType::getOperationName(), 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();