This is an archive of the discontinued LLVM Phabricator instance.

[mlir] Add an interface to decompose complex ops
ClosedPublic

Authored by qcolombet on Jul 3 2023, 9:32 AM.

Details

Summary

This patch adds an interface, named AggregatedOpInterface, that decomposes complex operations into simpler ones.

This is WIP because I was wondering how we want to connect this, e.g., regular pass, transform dialect, etc.

To demonstrate the e2e use case, I went with a new the transform dialect op but it would have made sense to reuse the decompose one (albeit it won't work out of the box because that one only support linalg interfaces). I'm guessing we want to express the convolution decomposition in terms of interfaces as well. Anyhow, this is all up for discussion :).

Diff Detail

Event Timeline

qcolombet created this revision.Jul 3 2023, 9:32 AM
Herald added a project: Restricted Project. · View Herald TranscriptJul 3 2023, 9:32 AM
qcolombet requested review of this revision.Jul 3 2023, 9:32 AM
rengolin accepted this revision.Jul 4 2023, 1:46 AM

Thanks! LGTM with a few comment nits.

mlir/include/mlir/Interfaces/AggregatedOpInterface.td
32 ↗(On Diff #536797)

This comment seems at odds with the return type being a small vector of values.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
2412

Nit: "subtract a 'broadcasted' m from x" to make it clear that this is now an "element wise" subtraction on (dim=d), even if the op can just read the same element with a smart affine map.

The broadcast doesn't need to be an op, this would be just to make clear the semantics. (and even if it was, it could still be fused with the following op).

This revision is now accepted and ready to land.Jul 4 2023, 1:46 AM
qcolombet added inline comments.Jul 4 2023, 2:24 AM
mlir/include/mlir/Interfaces/AggregatedOpInterface.td
32 ↗(On Diff #536797)

Good catch.
Forgot to update this when I updated the API.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
2412

👍

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
1166

I am sorry, I already polluted this name with the existing "decompose" op which actually does not in the same way.

We could rename the existing decompose as downscale_to_1d or something similar and you'd have the name for this more apt transformation you are adding.

This is fine for now, we can do the cleanup as a followup, but that should assuage your comment about reuse the decompose one (albeit it won't work out of the box because that one only support linalg interfaces) -> it is a different transform.

chelini added inline comments.Jul 4 2023, 3:38 AM
mlir/include/mlir/Interfaces/AggregatedOpInterface.td
1 ↗(On Diff #536797)

AggregatedOpInterface.td

2 ↗(On Diff #536797)

Why this interface is not part of Linalg?

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
2317

nit: I would create a variable for the context to avoid builder.getContext() multiple times.

2324

nit: /*symbols=*/0

2332

nit: I think we want to have the builder as first arg.

2340

"we should have two maps"

2362

nit: We should have..

2363

I would expect to have identity map for both input operands. Is this not the case?

2382

nit: I think naming is a bit misleading here. Simply, buildDivOp?

2444

nit: I would rename: subtractAndExp to buildSubAndExpOp.

mlir/include/mlir/Interfaces/AggregatedOpInterface.td
2 ↗(On Diff #536797)

+1 to moving the interface in Linalg, this is not a general concept AFAICT

qcolombet added inline comments.Jul 4 2023, 7:56 AM
mlir/include/mlir/Interfaces/AggregatedOpInterface.td
2 ↗(On Diff #536797)

I thought this could generally be useful.
E.g., we have some expand passes (expand-strided-metadata, memref-expand) right now that could fit in this interface. But yeah, let's leave that out.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
2363

No, the second map drops one dimension since the max performed a reduction on dim

2382

Good point.

2444

👍

jpienaar added inline comments.
mlir/include/mlir/Interfaces/AggregatedOpInterface.td
2 ↗(On Diff #536797)

What is not a general concept here? This seems rather general and generally usable.

jpienaar added inline comments.Jul 4 2023, 8:27 AM
mlir/include/mlir/Interfaces/AggregatedOpInterface.td
21 ↗(On Diff #536797)

This seems rather abstract. This seems like querying a rewrite set to see if root pattern matcher exists that matches an op, except the rewrite set is coupled with the op rather than generic. Is that what this is?

qcolombet added inline comments.Jul 4 2023, 8:42 AM
mlir/include/mlir/Interfaces/AggregatedOpInterface.td
21 ↗(On Diff #536797)

Yes, that's exactly that.

qcolombet added inline comments.Jul 4 2023, 9:04 AM
mlir/include/mlir/Interfaces/AggregatedOpInterface.td
21 ↗(On Diff #536797)

@jpienaar
Do you think it would make sense for this interface to return a pattern instead?

qcolombet added inline comments.Jul 4 2023, 9:25 AM
mlir/include/mlir/Interfaces/AggregatedOpInterface.td
21 ↗(On Diff #536797)

Btw maybe a pass with greedy pattern rewriter is fine. I only reproduced what iree was doing here.

What do you all think?

mlir/include/mlir/Interfaces/AggregatedOpInterface.td
2 ↗(On Diff #536797)

Fair enough, if there is value in reusing this interface by all means let's.
If this needs to mature more, Linalg can be a place to shepherd it.

qcolombet updated this revision to Diff 537279.Jul 5 2023, 3:08 AM
qcolombet retitled this revision from [WIP][mlir] Add an interface to decompose complex ops to [mlir] Add an interface to decompose complex ops.
  • Make the AggregatedOpInterface specific to Linalg for now (we may want to drop the interface and just use rewriter patterns.)
  • Fix nits: comments, function names, move the builder first in all the added functions
  • Make the implementation of the transform op cleaner (i.e., support multi result in the returned value).
qcolombet marked 9 inline comments as done.Jul 5 2023, 3:10 AM
qcolombet added inline comments.
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
2382

Forgot to use a free cast here, will fix in a sec.

qcolombet updated this revision to Diff 537284.Jul 5 2023, 3:22 AM
  • Use cast free implementation
qcolombet added inline comments.Jul 5 2023, 3:35 AM
mlir/include/mlir/Interfaces/AggregatedOpInterface.td
21 ↗(On Diff #536797)

@jpienaar I moved this as a linalg only interface for now.
Does this work for you?

I think this needs to mature, like @nicolasvasilache said, but what I'd like to capture with this interface is operations that can be rewritten into simpler ops within the same (set of) dialect(s).

Right now, I found it hard to discover the patterns/passes you need to invoke to get the lowering you need. For instance, going back to my examples from yesterday, when lowering memref one has to know that they need to use memref-expand and expand-strided-metadata before being able to use finalize-memref-to-llvm.

Therefore, I'd like a way to capture the lowering -rewrite- patterns that stay within the same dialect.

jpienaar added inline comments.Jul 5 2023, 3:17 PM
mlir/include/mlir/Interfaces/AggregatedOpInterface.td
21 ↗(On Diff #536797)

This allows for a default expansion to "simpler" ops by some metric. It is rather convenient. I _could_ see it being too restrictive in general (but not by much honestly). But also many fit into this and are simple. Not blocking here. This just seems regular macro expansion not specific to linalg.

qcolombet added inline comments.Jul 18 2023, 2:58 AM
mlir/include/mlir/Interfaces/AggregatedOpInterface.td
21 ↗(On Diff #536797)

Agree on the "not specific to linalg".

I'm going to move forward with the current implementation to let it mature.
Given our conversation, I think we need to think how to make this easy to use and probably reuse the pattern rewriter infrastructure.