search

iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.84k stars 611 forks source link

Reshape fusion with multiple results #18514

Open Max191 opened 2 months ago

Max191 commented 2 months ago

The current reshape propagation patterns upstream bail on multi-result operations, but the implementation seems to support such cases as far as I can tell: https://github.com/llvm/llvm-project/blob/3733528e521b7ee6af3950c65c3ff421c8fd0af6/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp#L1253-L1258

This gist is an example of where propagation through a multi result op is needed: https://gist.github.com/Max191/d95fad4a3a47d77c78dcb157f4f2eac7

My plan is to add some sort of aggressive fusion option to the patterns to allow this fusion, but I'm not sure if it is necessary. In codegen, we probably want to try to propagate through these ops, since we want to push any reshapes as far as we can to the edges of the function. In dispatch creation, the same logic would make sense to me, but I am unsure how it will play with the fusion logic that follows it. My question is, is it worth adding an option for this at all, or should I just remove the restriction?

MaheshRavishankar commented 2 months ago

Just remove the restriction... its a bit artificial. If there are issues that come up then we will know.

MaheshRavishankar commented 2 months ago

cc @IanWood1

Max191 commented 2 months ago

Posting the PR here for reference https://github.com/llvm/llvm-project/pull/108576