[mlir][tensor][linalg] Introduce DataLayoutPropagation pass.

[hanchung]

Summary

It introduces a pattern that swaps linalg.generic + tensor.pack to
tensor.pack + linalg.generic. It requires all the iteration types
being parallel; the indexing map of output operand is identiy. They can
all be relaxed in the future.

The user can decide whether the propagation should be applied or not by
passing a control function.

[hanchung]

remove unused variable

update indexing_map

[chelini]

mlir/include/mlir/Dialect/Utils/IndexingUtils.h

83 ↗

(On Diff #478427)

I think we should update the comment here. You get [c, a].

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

151

is llvm::None padding? Can we also propagate if we have padding?

171

We can use:

rewriter.inlineRegionBefore(genericOp.getRegion(), newGenericOp.getRegion(), newGenericOp.getRegion().begin());

[nicolasvasilache]

mlir/include/mlir/Dialect/Utils/IndexingUtils.h

85 ↗

(On Diff #478427)

Why do you need the , N part ?

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

12

please use linalg.map for all the generics in this test, they are simple maps

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/Passes.td

49 ↗

(On Diff #478427)

Should this be a test pass ?
We don't want to start adding (and especially maintaining) heuristics at this point, right?

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

84

this is lacking docs

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

34

Please add this to PackOp

static PackOp::createDestinationTensor(...)

This is where one would expect to find this helper.

59

more docs plz, input IR, output IR etc.

68

The functionality should be implemented in a function that takes relevant pieces of IR (i.e. the packop) and returns the relevant pieces of IR (the new linalg op and the new unpack op).
The pattern should just call those.

This way we can connect a transform op and even replace that test pass.

87

this is lacking docs, and more especially why you need some non-serializable C++ lambda injected here

104

Seems like a very complex piece of code to leave undocumented and inlined here.

What can be hoisted, what becomes a util on some op etc?

[hanchung]

mlir/include/mlir/Dialect/Utils/IndexingUtils.h

83 ↗

(On Diff #478427)

I misunderstood the semantics of outer_dims_perm. It should be either a permutation or empty. The method does not make sense to me now.. We won't need this method anymore. I enhance the verifier in https://reviews.llvm.org/D138936

[hanchung]

address comments

[hanchung]

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

34

The trade-off is that we can't use tensor::createDimValues if we add it to PackOp. Otherwise, there will be cyclic deps. TensorOps -> TensorUtils -> TensorOps

68

Do we want it being able to connect to a transform op? My understanding is that transform ops aims to apply the pattern once, while propagation patterns would be applied until IR gets converged. The pass is at the position similar to element-wise fusion, IMO.

87

It was intended to provide a callback function for users. So the users can decide what to propagate. Since the pass is in early development phase, we can drop the support; and add it back when needed.

151

It could introduce undefined behavior if we unconditionally propagate pack op through all the ops. E.g., if the padding value is zero and there are division ops in a generic op. Some values of padding area could be NaN (0/0).

I disable the propagation in the pattern by default, and thinking to add an option for enabling aggressive propagation in the future.

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

12

IIUC, linalg.map op is named op. I can't simply update the test file because the pattern works on linalg::GenericOp. Are you suggesting me to update the pattern taking linalg::LinalgOp instead? I think that could work for linalg.map, linalg.transpose, etc; I can take a stab at it, just wanna clarify before updating it.

[hanchung]

move the pass to test-only

[mravishankar]

mlir/include/mlir/Dialect/Linalg/Passes.td

49 ↗

(On Diff #478427)

+1 to this point. Might be better to expose the core utility method or patterns as populate methods, and move the pass itself to a test pass.

mlir/include/mlir/Dialect/Utils/IndexingUtils.h

85 ↗

(On Diff #478427)

If you want to update the SmallVector in place, you could use MutableArrayRef<..>.

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

68

It is still worth separating the core implementation into a stand-alone function (even if it is static and not exposed right now). Using fixed point is one way of applying the functionality. Even the elementwise op fusion patterns are just a wrapper around the core functions that implement the transformation.

87

This is the control that is based on call site constraints. Instead of using heuristics in the core method that decides when to apply this, we defer the control to the caller, and make a go/no-go decision....

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

87

Injecting lambda form above allows inversion of control that is convenient very short term and has almost always proven to very quickly turn to technical debt.

Given recent offline discussions I had and other parts of the codebase I have seen, I am going to more seriously push back against this anti-pattern globally.
Injecting C++ callback control from above is a sign of missing abstractions and should almost always be disallowed.
The alternative is often to refactor multiple times until the right abstractions emerge.

In other words, the transformations we add must be functional-style, statically return the pieces of IR (existing created or updated) that make sense for that transform.
This is not something customizable, if you need more information then statically return more information: no backchannels through lambdas.

If you need different behaviors, instead of injecting a dynamic mechanism through a lambda, write another transformation that takes more/different inputs and return more/different outputs.
Refactor the reusable utility helpers to avoid copy-pasta.

These transformations can then be plugged into patterns and transform dialect ops who can be responsible for the switch between different static behaviors.

[hanchung]

remove unneeded changes

move implementations to functions

[hanchung]

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

87

thanks for the details! So we provide methods upstream, and let upstream patterns and downstream projects use the methods. They can always wrap the methods into a custom patterns w/o a callback function. I moved them to functions and the pattern becomes a wrapper of the method.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

87

We discussed offline with Mahesh and this may be one of the cases that I could see as an exception to the rule.

Some loose characteristics of potential green/red flags that seem to emerge are:

1. the information conveyed by the callback cannot be serialized in IR in a reasonable form
2. the callback use remains superficial and does not permeate the transform implementation but is called once at the high level (e.g. in the match part of a pattern or in a transform dialect op)
3. the callback does not capture or bypass the return mechanism
4. maybe others ?

I'd love a more principled discussion on the topic and that we come up with good guidelines we can follow.

Removing the blocker, thanks for your patience.

Edit: I don't know how to do this without accepting or resigning but consider my concern addressed :)

[hanchung]

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

12

It's quite tricky because the builder and clone method can't be reused. The most tricky part is about updating indexing_maps. The builder and clone method does not take indexing_maps into account. Also, there is not a clean way to use setting methods for them.

[chelini]

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

180

We also need to make sure that the result of the current linalg generic operation has no users. Otherwise, the linalg generic remains alive without a body.

230

struct instead of class?

[chelini]

Overall looks good to me. Just left two minor comments.

[hanchung]

inlineRegionBefore -> cloneRegionBefore

class -> struct

[hanchung]

rebase

[mravishankar]

Thanks! Overall looks good to me. Have a few nits.

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

82

Nit: Move the comments before the statement to make it read better.

85

Nice! TIL.

[nicolasvasilache]

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

12

@pifon2a ?

[hanchung]

move comments above the variable declaration

[hanchung]

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

180

I see the point. I think we should use cloneRegionBefore instead of inline*.

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

12

I'm going to land the PR to unblock some work on @chelini side tomorrow. If it's an easy fix, I can give it a shot. If not, I'd prefer fixing it afterwards. Thanks!

[hanchung]

All the comments have been addressed except testing with linalg.map. I did some study and think that's not feasible to address. The pass works on the generic ops similar to ElementWiseFusion. If it can be done in elementwise fusion, it can also be done here. Since the other pass is testing with this form, I think it's not a big concern. I'm going to land the revision, and prefer addressing the issue afterwards if we know a way to address it.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

65

std::pair

82

nit: typo

85

This seems like it will crash if not a projected permutation.
Plz add a clear assert at the top-level of the func if this can never happen or consider failing.