[GPU][Codegen] Expand iteration space based on new `expand_dims` attribute #22342

efric · 2025-10-17T00:41:16Z

This patch introduces iteration space expansion for reductions in the VectorDistribute path.

Specifically, we:

Add a new attribute, expand_dims, for reductions.
Introduce a new pass, GPUExpandDimensions, which uses expand_dims to expand the iteration space of relevant dimensions.
Refactor common functionality shared between GPUExpandDimensions and BlockDynamicDimensions into reusable utilities.
Refactor encoding helpers from EncodingAttrs.cpp into reusable utilities.

This change also enables chain FMA in matvec codegen as we iterate along the K reduction dimension.

Performance Summary

IREE benchmark module

Only expansion: ~4% improvement
Expansion + chain FMA: ~11% improvement

rocprof

Only expansion: ~13% worse
Expansion + chain FMA: ~9% better

Register usage

10% reduction (60 → 54 registers for matvec dispatches)

Instruction latency (post-reduction loop epilogue)

3.5% improvement (340 → 328 total mean latency)

Notes

As a follow-up, we can explore applying iteration space expansion to the reduction in attention
Right now, we only expand one dimension into two although the implementation supports expansion to N dimensions.
Please note this PR changes the reduction order, some expect some minor changes to the numerics
This is does not improve performance by itself/can cause regression without chain FMA [LLVMGPU][Codegen] Emit packed chain FMA from select multi_reductions and contracts #21855

Traces for matvec dispatches are attached for all variations (original, only expansion, and expansion + chain FMA).

115_expansion_and_chain.tar.gz
115_nothing.tar.gz
115_only_expansion.tar.gz

Fixes: #22153

ci-extra: test_torch

Signed-off-by: Eric Feng <[email protected]>

Max191

Overall looks good now, just some nits and one comment about the configuration logic. Nice work, and thanks for addressing all my comments so far!

In case I'm OOO by the time you address these changes, I'll remove my request for changes now. It LGTM after these final comments, so anyone else can give the final approval after this round of comments.

Max191 · 2025-12-10T19:44:59Z

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

+#include "iree/compiler/Codegen/Common/Transforms.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
+#include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVectorExtras.h"
+#include "llvm/Support/DebugLog.h"
+#include "llvm/Support/LogicalResult.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"


nit: Please clean up the includes. I think a lot of them are probably not necessary.

I think we need these remaining ones

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

Max191 · 2025-12-10T20:22:31Z

compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_expand_dimensions.mlir

+
+// -----
+
+func.func @expand_dynamic_dim(%a: tensor<4x?xf16>, %b: tensor<1x?xf16>) -> tensor<4x1xf32> {


nit: For this test and the test below, can you use something simple like linalg.add? The tests are quite long, and I don't think it needs to be a matvec to test what is needed.

Switched to linalg.add for all except one (I think it is a useful reference since it's the only use real use of this atm)

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

compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ReductionConfigUtils.cpp

I will be OOO, and it looks good to me after the final comments are addressed. Someone else can give the final approval.

Signed-off-by: Eric Feng <[email protected]>

Groverkss · 2025-12-18T17:20:29Z

@efric I added a ci-extra trailer to run test_torch can you check?

Signed-off-by: Eric Feng <[email protected]>

The current binary size for this test is nearing its threshold (459794/460000). Bump to 470000 to allow future changes such as #22342 some room (goes up to 460306) Signed-off-by: Eric Feng <[email protected]>

The current binary size for this test is nearing its threshold (459794/460000). Bump to 470000 to allow future changes such as iree-org#22342 some room (goes up to 460306) Signed-off-by: Eric Feng <[email protected]>

…ibute (iree-org#22342) This patch introduces iteration space expansion for reductions in the VectorDistribute path. Specifically, we: 1. Add a new attribute, `expand_dims`, for reductions. 2. Introduce a new pass, `GPUExpandDimensions`, which uses `expand_dims` to expand the iteration space of relevant dimensions. 3. Refactor common functionality shared between `GPUExpandDimensions` and `BlockDynamicDimensions` into reusable utilities. 4. Refactor encoding helpers from `EncodingAttrs.cpp` into reusable utilities. This change also enables [chain FMA](iree-org#21855) in matvec codegen as we iterate along the K reduction dimension. --- **Performance Summary** **IREE benchmark module** * Only expansion: ~4% improvement * Expansion + chain FMA: ~11% improvement **rocprof** * Only expansion: ~13% worse * Expansion + chain FMA: ~9% better **Register usage** * 10% reduction (60 → 54 registers for matvec dispatches) **Instruction latency (post-reduction loop epilogue)** * 3.5% improvement (340 → 328 total mean latency) --- **Notes** * As a follow-up, we can explore applying iteration space expansion to the reduction in attention * Right now, we only expand one dimension into two although the implementation supports expansion to N dimensions. * Please note this PR changes the reduction order, so expect some minor changes to the numerics * This is does not improve performance by itself/can cause regression without chain FMA iree-org#21855 Traces for matvec dispatches are attached for all variations (original, only expansion, and expansion + chain FMA). [115_expansion_and_chain.tar.gz](https://github.com/user-attachments/files/23268046/115_expansion_and_chain.tar.gz) [115_nothing.tar.gz](https://github.com/user-attachments/files/23268047/115_nothing.tar.gz) [115_only_expansion.tar.gz](https://github.com/user-attachments/files/23268048/115_only_expansion.tar.gz) Fixes: iree-org#22153 ci-extra: test_torch --------- Signed-off-by: Eric Feng <[email protected]> Signed-off-by: Eric Feng <[email protected]>

…ms` attribute (iree-org#22342)" This reverts commit b879ee1.

…ms` attribute (iree-org#22342)" This reverts commit b879ee1. Signed-off-by: Eric Feng <[email protected]>

…ms` attribute (#22342) (#22982) This reverts commit b879ee1. To temporarily fix ONNX models test suite from issues suspected from [[DispatchCreation] Unit dims not properly folded away during GlobalOptimization](#22978) and reported regression. Signed-off-by: Eric Feng <[email protected]>

The current binary size for this test is nearing its threshold (459794/460000). Bump to 470000 to allow future changes such as iree-org#22342 some room (goes up to 460306) Signed-off-by: Eric Feng <[email protected]>

…ibute (iree-org#22342) This patch introduces iteration space expansion for reductions in the VectorDistribute path. Specifically, we: 1. Add a new attribute, `expand_dims`, for reductions. 2. Introduce a new pass, `GPUExpandDimensions`, which uses `expand_dims` to expand the iteration space of relevant dimensions. 3. Refactor common functionality shared between `GPUExpandDimensions` and `BlockDynamicDimensions` into reusable utilities. 4. Refactor encoding helpers from `EncodingAttrs.cpp` into reusable utilities. This change also enables [chain FMA](iree-org#21855) in matvec codegen as we iterate along the K reduction dimension. --- **Performance Summary** **IREE benchmark module** * Only expansion: ~4% improvement * Expansion + chain FMA: ~11% improvement **rocprof** * Only expansion: ~13% worse * Expansion + chain FMA: ~9% better **Register usage** * 10% reduction (60 → 54 registers for matvec dispatches) **Instruction latency (post-reduction loop epilogue)** * 3.5% improvement (340 → 328 total mean latency) --- **Notes** * As a follow-up, we can explore applying iteration space expansion to the reduction in attention * Right now, we only expand one dimension into two although the implementation supports expansion to N dimensions. * Please note this PR changes the reduction order, so expect some minor changes to the numerics * This is does not improve performance by itself/can cause regression without chain FMA iree-org#21855 Traces for matvec dispatches are attached for all variations (original, only expansion, and expansion + chain FMA). [115_expansion_and_chain.tar.gz](https://github.com/user-attachments/files/23268046/115_expansion_and_chain.tar.gz) [115_nothing.tar.gz](https://github.com/user-attachments/files/23268047/115_nothing.tar.gz) [115_only_expansion.tar.gz](https://github.com/user-attachments/files/23268048/115_only_expansion.tar.gz) Fixes: iree-org#22153 ci-extra: test_torch --------- Signed-off-by: Eric Feng <[email protected]> Signed-off-by: Eric Feng <[email protected]>

The current binary size for this test is nearing its threshold (459794/460000). Bump to 470000 to allow future changes such as iree-org#22342 some room (goes up to 460306) Signed-off-by: Eric Feng <[email protected]> Signed-off-by: Bangtian Liu <[email protected]>

…ibute (iree-org#22342) This patch introduces iteration space expansion for reductions in the VectorDistribute path. Specifically, we: 1. Add a new attribute, `expand_dims`, for reductions. 2. Introduce a new pass, `GPUExpandDimensions`, which uses `expand_dims` to expand the iteration space of relevant dimensions. 3. Refactor common functionality shared between `GPUExpandDimensions` and `BlockDynamicDimensions` into reusable utilities. 4. Refactor encoding helpers from `EncodingAttrs.cpp` into reusable utilities. This change also enables [chain FMA](iree-org#21855) in matvec codegen as we iterate along the K reduction dimension. --- **Performance Summary** **IREE benchmark module** * Only expansion: ~4% improvement * Expansion + chain FMA: ~11% improvement **rocprof** * Only expansion: ~13% worse * Expansion + chain FMA: ~9% better **Register usage** * 10% reduction (60 → 54 registers for matvec dispatches) **Instruction latency (post-reduction loop epilogue)** * 3.5% improvement (340 → 328 total mean latency) --- **Notes** * As a follow-up, we can explore applying iteration space expansion to the reduction in attention * Right now, we only expand one dimension into two although the implementation supports expansion to N dimensions. * Please note this PR changes the reduction order, so expect some minor changes to the numerics * This is does not improve performance by itself/can cause regression without chain FMA iree-org#21855 Traces for matvec dispatches are attached for all variations (original, only expansion, and expansion + chain FMA). [115_expansion_and_chain.tar.gz](https://github.com/user-attachments/files/23268046/115_expansion_and_chain.tar.gz) [115_nothing.tar.gz](https://github.com/user-attachments/files/23268047/115_nothing.tar.gz) [115_only_expansion.tar.gz](https://github.com/user-attachments/files/23268048/115_only_expansion.tar.gz) Fixes: iree-org#22153 ci-extra: test_torch --------- Signed-off-by: Eric Feng <[email protected]> Signed-off-by: Eric Feng <[email protected]> Signed-off-by: Bangtian Liu <[email protected]>

efric added 23 commits October 12, 2025 21:15

Initial pass working expand dims

94c0ad0

Signed-off-by: Eric Feng <[email protected]>

nits

6f2e6e0

Signed-off-by: Eric Feng <[email protected]>

switch to listener to propagate lowering config to new linalg.generics

430e8cc

Signed-off-by: Eric Feng <[email protected]>

Merge remote-tracking branch 'origin/main' into it-exp-2

0589d0b

set expand_dims lowering_config

e62e6e6

Signed-off-by: Eric Feng <[email protected]>

nits

e83d058

Signed-off-by: Eric Feng <[email protected]>

refactor out common methods 1st pass

6d2bf2d

Signed-off-by: Eric Feng <[email protected]>

generalize to all linalg ops to fix test + refactor common methods

e79baab

Signed-off-by: Eric Feng <[email protected]>

nits

15385c9

Signed-off-by: Eric Feng <[email protected]>

plumb through existing tests

2cf5a88

Signed-off-by: Eric Feng <[email protected]>

formatting

5a8a2e0

Signed-off-by: Eric Feng <[email protected]>

fix debug wrapper

93c4d99

Signed-off-by: Eric Feng <[email protected]>

new precommit fix

906b5b1

Signed-off-by: Eric Feng <[email protected]>

Merge remote-tracking branch 'origin/main' into it-exp-2

f5f68ca

dont expand dynamic matmuls

e3b793a

Signed-off-by: Eric Feng <[email protected]>

more test reversal

8e18b1e

Signed-off-by: Eric Feng <[email protected]>

nits

9acc711

Signed-off-by: Eric Feng <[email protected]>

more nits

f3015b8

Signed-off-by: Eric Feng <[email protected]>

nit

cc9cdb9

Signed-off-by: Eric Feng <[email protected]>

nits

3f15cdd

Signed-off-by: Eric Feng <[email protected]>

precommit

b81ad92

Signed-off-by: Eric Feng <[email protected]>

nit

ac1091c

Signed-off-by: Eric Feng <[email protected]>

nit

6d88437

Signed-off-by: Eric Feng <[email protected]>

efric changed the title ~~[GPU][Codegen] Expand dimensions based on expand_dims lowering_config~~ [GPU][Codegen] Expand iteration space based on new expand_dims lowering_config Oct 17, 2025

efric added 6 commits October 16, 2025 23:20

refactor to use reassociation indices

454937c

Signed-off-by: Eric Feng <[email protected]>

nit

466b01a

Signed-off-by: Eric Feng <[email protected]>

format

1c84d11

Signed-off-by: Eric Feng <[email protected]>

nits

181996a

Signed-off-by: Eric Feng <[email protected]>

nits comments

c505ef3

Signed-off-by: Eric Feng <[email protected]>

nits, semantics and test

4bba6d1

Signed-off-by: Eric Feng <[email protected]>

Max191 reviewed Dec 10, 2025

View reviewed changes

efric and others added 11 commits December 16, 2025 15:20

refactor to put this into lowering_config

30fede9

Signed-off-by: Eric Feng <[email protected]>

nits

671dd54

Signed-off-by: Eric Feng <[email protected]>

nits

61e54f7

Signed-off-by: Eric Feng <[email protected]>

nits remove extraneous headers

30bddd9

Signed-off-by: Eric Feng <[email protected]>

refactor tests

ced8c45

Signed-off-by: Eric Feng <[email protected]>

nits

451b02d

Signed-off-by: Eric Feng <[email protected]>

test to see if new outputs are now extraneous

51fa3ce

Signed-off-by: Eric Feng <[email protected]>

nit

ea7eef6

Signed-off-by: Eric Feng <[email protected]>

fix typo

873db4f

Signed-off-by: Eric Feng <[email protected]>

add comments

5243493

Signed-off-by: Eric Feng <[email protected]>

nit to test

2e6c2d6

Signed-off-by: Eric Feng <[email protected]>

Groverkss approved these changes Dec 18, 2025

View reviewed changes

trigger ci

8850fc8

Signed-off-by: Eric Feng <[email protected]>

efric force-pushed the static_dimension_expansion branch from edb0ed2 to 8850fc8 Compare December 18, 2025 23:02

efric mentioned this pull request Dec 19, 2025

[CI][Torch Models] Update golden binary size #22951

Merged

efric merged commit b879ee1 into iree-org:main Dec 23, 2025
130 of 142 checks passed

efric mentioned this pull request Dec 24, 2025

[DispatchCreation] Unit dims not properly folded away during GlobalOptimization #22978

Open

efric added a commit to efric/iree that referenced this pull request Dec 24, 2025

Revert "[GPU][Codegen] Expand iteration space based on new `expand_di…

02665b2

…ms` attribute (iree-org#22342)" This reverts commit b879ee1.

efric added a commit to efric/iree that referenced this pull request Dec 24, 2025

Revert "[GPU][Codegen] Expand iteration space based on new `expand_di…

c7a90a1

…ms` attribute (iree-org#22342)" This reverts commit b879ee1. Signed-off-by: Eric Feng <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[GPU][Codegen] Expand iteration space based on new `expand_dims` attribute #22342

[GPU][Codegen] Expand iteration space based on new `expand_dims` attribute #22342

Uh oh!

efric commented Oct 17, 2025 •

edited by Groverkss

Loading

Uh oh!

Max191 left a comment

Uh oh!

Max191 Dec 10, 2025

Uh oh!

efric Dec 17, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Max191 Dec 10, 2025

Uh oh!

efric Dec 17, 2025

Uh oh!

Uh oh!

Uh oh!

Groverkss commented Dec 18, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants


		// -----

		func.func @expand_dynamic_dim(%a: tensor<4x?xf16>, %b: tensor<1x?xf16>) -> tensor<4x1xf32> {

[GPU][Codegen] Expand iteration space based on new expand_dims attribute #22342

[GPU][Codegen] Expand iteration space based on new expand_dims attribute #22342

Uh oh!

Conversation

efric commented Oct 17, 2025 • edited by Groverkss 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!

Max191 Dec 10, 2025

Choose a reason for hiding this comment

Uh oh!

efric Dec 17, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Max191 Dec 10, 2025

Choose a reason for hiding this comment

Uh oh!

efric Dec 17, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Groverkss commented Dec 18, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

[GPU][Codegen] Expand iteration space based on new `expand_dims` attribute #22342

[GPU][Codegen] Expand iteration space based on new `expand_dims` attribute #22342

efric commented Oct 17, 2025 •

edited by Groverkss

Loading