This is an archive of the discontinued LLVM Phabricator instance.

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

Authored by nicolasvasilache on Mar 30 2021, 4:43 AM.

Details

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.

Diff Detail

Event Timeline

nicolasvasilache requested review of this revision.Mar 30 2021, 4:43 AM

@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 added inline comments.Mar 30 2021, 3:14 PM
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 added inline comments.Mar 30 2021, 3:17 PM
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.

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.

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 marked 10 inline comments as done.Mar 30 2021, 11:43 PM
nicolasvasilache retitled this revision from [mlir][Linalg][Python] WIP - Create the body of builtin named Linalg ops to [mlir][Linalg][Python] Create the body of builtin named Linalg ops.
nicolasvasilache edited the summary of this revision. (Show Details)
nicolasvasilache edited the summary of this revision. (Show Details)

Address review.

drop spurious print

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 accepted this revision.Mar 31 2021, 12:28 AM

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.

This revision is now accepted and ready to land.Mar 31 2021, 12:28 AM
nicolasvasilache marked 4 inline comments as done.Mar 31 2021, 12:53 AM
nicolasvasilache added inline comments.
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 marked an inline comment as done.

Address review.

This revision was landed with ongoing or failed builds.Mar 31 2021, 1:01 AM
This revision was automatically updated to reflect the committed changes.
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...