[Codegen] Create CombineSourceLayoutTransformation pass for MapGatherOp by Abhishek-Varma · Pull Request #23165 · iree-org/iree

Abhishek-Varma · 2026-01-16T09:09:05Z

-- This commit creates a separate pass iree-codegen-combine-source-layout-transformation
adding support of folding relayout operations into map_gather op.
-- The following ops' folding are being supported through changes in this PR :-
1. linalg::CopyOp
2. linalg::TransposeOp
3. tensor::ExpandShapeOp
4. tensor::CollapseShapeOp
5. tensor::ExtractSliceOp
6. tensor::PadOp
-- Consequently getPaddingValue and createIdentityMapGather member methods were added to MapGatherOp.
-- The implementations for the iree-codegen-combine-source-layout-transformation pass reside in the same file as MapScatter's relayout folding implementations (which has been renamed as iree-codegen-combine-result-layout-transformation).

hanhanW

I did not review all the code yet, but I have a first high-level question: do we want to make both transformation (gather and scatter) be in a single pass? They look orthogonal to me, and it is not clear that which folding is prioritized.

The other high level comment is passing bool-ish argument (i.e., your new enum) to functions. It makes this much more difficult to follow (especially in what is already difficult code). This comment depends on the first mega comment about if we want to split it into two passes or not.

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp

compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td

compiler/src/iree/compiler/Codegen/Common/GPU/GPUCombineLayoutTransformation.cpp

nirvedhmeshram · 2026-01-20T15:10:10Z

-- Unlike map_scatter, we don't allow tensor.pad op in the relayout op chain

I think we want to allow this as the use case mentioned in #23038 has pad op in it?

hanhanW · 2026-01-21T01:12:44Z

Hey @Abhishek-Varma I did not get a chance to review it today. I'll prioritize it tomorrow. You can continue working on other changes and use stacked PR features. (or just send the PR out and ask me/others to review the last commit.)

hanhanW · 2026-01-21T01:31:37Z

I don't follow why a review comment is generated by LLM agents. I hided it and deleted the app.

hanhanW

I haven't reviewed the core logic yet, but can we clean up first? I think some of them are not used at all, and it is adding tech debt and overheads to reviewers.

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformationForMapGather.cpp

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformationForMapGather.h

compiler/src/iree/compiler/Codegen/Common/Passes.td

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformationForMapGather.cpp

...iler/src/iree/compiler/Codegen/Common/test/combine_layout_transformation_for_map_gather.mlir

compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp

compiler/src/iree/compiler/Codegen/Common/Transforms.cpp

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformationForMapGather.cpp

Abhishek-Varma · 2026-01-30T09:14:22Z

Hi @hanhanW.

I've now changed the algo. I'm creating identity map_gather at the result of load_from_buffer and now iteratively fusing in consumer relayout ops instead of the earlier algo of creating an identity map_gather at the end of a relayout op chain. It's more intuitive now.

I've also refactored the code because most of the foldings followed a similar structure - so reviewing it will be easier now!

Locally, I've included this pass at

iree/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp

Lines 544 to 545 in 7aa7be5

    
           funcPassManager.addPass( 
        
               createCombineLayoutTransformationPass(combineLayoutOptions));

and have switched off the temporary workaround flag iree-codegen-test-c-promtion :-

I used the above setup to test the original regression with the repro dispatch and I confirm it resolves the issue :-

real    0m10.902s
user    0m10.810s
sys     0m0.163s

NOTE: I confirmed first by simply switching off the iree-codegen-test-c-promtion flag - the compilation of the repro dispatch was taking > 10 minutes. This was done as a way to "revert" the temporary workaround first and confirm the issue exists. Then I followed the above setup and observed that with the pass introduced through this PR the compilation takes place successfully.

I've even verified the output of the dispatch. Even isolated a few cases like transpose op and verified the e2e correctness too.

I'm adding the output.mlir and input.mlir for reference that the pass can handle complex relayout chains.

So, once this PR merges - I can raise another PR that removes the temporary workaround and adds this pass in the addGPUTileAndFusePassPipeline.

CC: @MaheshRavishankar @nirvedhmeshram @jtuyls @Max191

… Dispatch scope -- This commit adds support of map_gather op in CombineLayoutTransformation in Dispatch scope. -- Unlike map_scatter, we don't allow tensor.pad op in the relayout op chain for map_gather. This is because map_scatter has a known output buffer at the end of the chain, while map_gather only has a known source buffer at the beginning of the relayout ops chain. Signed-off-by: Abhishek Varma <[email protected]>

Signed-off-by: Abhishek Varma <[email protected]>

Max191

Overall, looks good. I left one main comment about the tensor.pad handling, and then mostly nits.

One more high level question: Have you verified numerical correctness with any e2e examples using the CombineLayoutTransformationForMapGather pass yet?

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformationForMapGather.cpp

compiler/src/iree/compiler/Codegen/Common/Passes.td

compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td

hanhanW

Mostly looks good. I was thinking if I should ask refactoring, or defer it when we're evalulating if we want to collapse two passes or not. Moving them to the same file and keep them two separate pass looks good to me.

compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformationForMapGather.cpp

Signed-off-by: Abhishek Varma <[email protected]>

Abhishek-Varma · 2026-02-03T12:55:55Z

One more high level question: Have you verified numerical correctness with any e2e examples using the CombineLayoutTransformationForMapGather pass yet?

Hi @Max191 - Yes, I did in fact check e2e correctness as reported here. I verified the repro's output as well as a standalone transpose op example.

And today I also checked a bunch of complex relayout op chains' e2e correctness too.

Here's what I did :-

Formed a complex chain relayout ops.
Applied the current pass.
Added the output to e2e map_gather test already checked in along with corresponding check.expect* ops with what tensor values are expected.

The tests passed! :D

hanhanW

Mostly looks good, just some nits. I think we over-abstract the "factor" functions.

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp

compiler/src/iree/compiler/Codegen/Common/Passes.td

compiler/src/iree/compiler/Codegen/Common/test/combine_source_layout_transformation.mlir

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp

hanhanW · 2026-02-06T22:44:00Z

The PR title and description need to be updated.

Signed-off-by: Abhishek Varma <[email protected]>

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp

hanhanW

just one final nit. lgtm

Max191

LGTM, nice work! Just a couple nits.

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp

compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp

Signed-off-by: Abhishek Varma <[email protected]>

…Op (iree-org#23165) -- This commit creates a separate pass `iree-codegen-combine-source-layout-transformation` adding support of folding relayout operations into map_gather op. -- The following ops' folding are being supported through changes in this PR :- 1. `linalg::CopyOp` 2. `linalg::TransposeOp` 3. `tensor::ExpandShapeOp` 4. `tensor::CollapseShapeOp` 5. `tensor::ExtractSliceOp` 6. `tensor::PadOp` -- Consequently `getPaddingValue` and `createIdentityMapGather` member methods were added to `MapGatherOp`. -- The implementations for the `iree-codegen-combine-source-layout-transformation` pass reside in the same file as `MapScatter`'s relayout folding implementations (which has been renamed as `iree-codegen-combine-result-layout-transformation`). Signed-off-by: Abhishek Varma <[email protected]>

Abhishek-Varma requested review from Groverkss, IanWood1, MaheshRavishankar, Max191, hanhanW, kuhar and qedawkins as code owners January 16, 2026 09:09

Abhishek-Varma mentioned this pull request Jan 16, 2026

[LinalgExt][Codegen] iree_linalg_ext.map_gather Design Doc #23038

Open

hanhanW requested changes Jan 19, 2026

View reviewed changes

Abhishek-Varma requested a review from hanhanW January 20, 2026 12:21

This comment was marked as abuse.

Sign in to view

hanhanW requested changes Jan 22, 2026

View reviewed changes

Abhishek-Varma force-pushed the avarma_map_gather_combine_layout branch from 7b5bfa1 to 6821f91 Compare January 29, 2026 06:28

hanhanW requested changes Jan 30, 2026

View reviewed changes

Abhishek-Varma changed the title ~~[Codegen] Add support of map_gather in CombineLayoutTransformation for Dispatch Scope~~ [Codegen] Create CombineLayoutTransformationForMapGather pass Jan 30, 2026

Abhishek-Varma force-pushed the avarma_map_gather_combine_layout branch from 9dc047e to 7e3b728 Compare January 30, 2026 09:08

Abhishek-Varma requested a review from hanhanW January 30, 2026 09:14

Abhishek-Varma added 7 commits February 2, 2026 09:09

Review comment v1.0

1e10028

Signed-off-by: Abhishek Varma <[email protected]>

Add tensor.pad folding into map_gather

eb922e5

Signed-off-by: Abhishek Varma <[email protected]>

Clean up

6f3d3c6

Signed-off-by: Abhishek Varma <[email protected]>

Trivial rebase issue

5baba36

Signed-off-by: Abhishek Varma <[email protected]>

Producer fusion -> consumer fusion

7106d4a

Signed-off-by: Abhishek Varma <[email protected]>

Review comment v2.0

4aa04d7

Signed-off-by: Abhishek Varma <[email protected]>

Max191 reviewed Feb 2, 2026

View reviewed changes

hanhanW reviewed Feb 2, 2026

View reviewed changes

Abhishek-Varma force-pushed the avarma_map_gather_combine_layout branch from 7e3b728 to 35808c2 Compare February 3, 2026 12:40

Abhishek-Varma requested review from krzysz00 and nirvedhmeshram as code owners February 3, 2026 12:40

Review comment Max v1.0

48f1f33

Signed-off-by: Abhishek Varma <[email protected]>

Abhishek-Varma force-pushed the avarma_map_gather_combine_layout branch from 35808c2 to 48f1f33 Compare February 3, 2026 12:42

Abhishek-Varma requested review from Max191 and hanhanW February 3, 2026 12:56

hanhanW reviewed Feb 6, 2026

View reviewed changes

Review comment Hanhan - less abstraction

af8272b

Signed-off-by: Abhishek Varma <[email protected]>

Abhishek-Varma changed the title ~~[Codegen] Create CombineLayoutTransformationForMapGather pass~~ [Codegen] Create CombineSourceLayoutTransformation pass for MapGatherOp Feb 9, 2026

Abhishek-Varma requested a review from hanhanW February 9, 2026 11:37

hanhanW reviewed Feb 9, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp Show resolved Hide resolved

hanhanW approved these changes Feb 9, 2026

View reviewed changes

Max191 approved these changes Feb 9, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp Outdated Show resolved Hide resolved

compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp Outdated Show resolved Hide resolved

Final review comment

7cc5f16

Signed-off-by: Abhishek Varma <[email protected]>

Abhishek-Varma enabled auto-merge (squash) February 10, 2026 13:34

Abhishek-Varma merged commit 1320245 into iree-org:main Feb 10, 2026
57 checks passed

Conversation

Abhishek-Varma commented Jan 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hanhanW left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nirvedhmeshram commented Jan 20, 2026

Uh oh!

hanhanW commented Jan 21, 2026

Uh oh!

This comment was marked as abuse.

Uh oh!

hanhanW commented Jan 21, 2026

Uh oh!

hanhanW left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Abhishek-Varma commented Jan 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Max191 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

hanhanW left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Abhishek-Varma commented Feb 3, 2026

Uh oh!

hanhanW left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

hanhanW commented Feb 6, 2026

Uh oh!

Uh oh!

hanhanW left a comment

Choose a reason for hiding this comment

Uh oh!

Max191 left a comment

Choose a reason for hiding this comment

Abhishek-Varma commented Jan 16, 2026 •

edited

Loading

hanhanW left a comment •

edited

Loading

Abhishek-Varma commented Jan 30, 2026 •

edited

Loading