iree icon indicating copy to clipboard operation
iree copied to clipboard

[LinalgExt] Added DecomposeAggregatedOp Pass

Open hhkit opened this issue 1 month ago • 5 comments

Implementing this pass came up as a result of discussions in #22561, while we were implementing ExpReduction's implementation of AggregatedOpInterface. AggregatedOpInterface exposes decomposeOp which is called as part of a Transform dialect transformation (transform.structured.decompose, transform.iree.decompose_aggregate_op). However, we also call AggregatedOpInterface::decomposeOp as a step in many pipelines.

Therefore, in this PR we:

  • Add DecomposeAggregatedOpPass, which calls AggregatedOpInterface::decomposeOp and replaces original uses with the decomposed ops.
  • Added tests based on tests in the Transform dialect equivalent

For discussion:

  • Should we replace DecomposeAttention with this implementation? I noticed not every AggregatedOp fits into this neatly, as some decompositions (eg. decompose im2col) have custom options to take into account.
  • I have asked in Discord and grepped the upstream MLIR code base, because it can't be that something so common doesn't have an upstream equivalent (eg. a pass that invokes a Transform operation?) Happy to close this if anyone knows if there's one!

hhkit avatar Nov 27 '25 13:11 hhkit

Regarding some of the discussion points around when/how we use decomposition:

The main reasons for ops having their own decomposition passes are that we want to either (1) control when in the pipeline the op is decomposed, or (2) perform some additional preprocessing/cleanup/analysis before the decomposition that is specific to the op being decomposed.

For (1), your PR already addresses this with the filtering. For (2), it will be tricky to handle this with your PR, but I think that is okay. It is useful at the moment to have some special case decompositions because it combines the transformation for the specific op into a single pass, which reduces the number of passes we need to call in the pipeline and keeps the op-specific transformation more organized. If you'd like to discuss more on how to address (2), I am happy to, but I think it is okay to leave some special cases, because some op decompositions are more involved than others. The DecomposeAggregatedOp pass is still useful for simpler decompositions, regardless.

I don't have action items or requests here. I mostly just wanted to give some context.

Max191 avatar Dec 01 '25 15:12 Max191

If you'd like to discuss more on how to address (2), I am happy to, but I think it is okay to leave some special cases, because some op decompositions are more involved than others. The DecomposeAggregatedOp pass is still useful for simpler decompositions, regardless.

@Max191 I'm happy with the justification. I suppose the main concern is that there's now two ways to decompose, for example, a winograd. Not sure if there should be a way to opt-out of the pass or if that's overengineering.

If you think it's good to go, would you be able to approve this and run the CI, or should I try to find someone with more context?

hhkit avatar Dec 02 '25 20:12 hhkit

cc @Abhishek-Varma last I remember, you were trying to add such a pass. Did we not land it? Just want to make sure we have no redundancy.

Yes, but it didn't get merged : https://github.com/iree-org/iree/pull/15384 - I recall there was an upstream discussion around this that was a blocker at that point of time.

Abhishek-Varma avatar Dec 10 '25 06:12 Abhishek-Varma

This seems like it is ready to land. Let me know if you want me to force push this (the errors seem unrelated)

MaheshRavishankar avatar Dec 10 '25 19:12 MaheshRavishankar

This seems like it is ready to land. Let me know if you want me to force push this (the errors seem unrelated)

Sure, let's ship it.

hhkit avatar Dec 10 '25 20:12 hhkit