[Codegen][GPU] Enable swizzling for scaled matmuls by Muzammiluddin-Syed-ECE · Pull Request #23175 · iree-org/iree

Muzammiluddin-Syed-ECE · 2026-01-16T21:30:48Z

This is the fourth of a series of PRs that together implement support in
IREE for XOR swizzling through the SwizzleHintOp.

There are four PRs that need to be merged:

Allow rank > 1 swizzle hint op operands and add a pass to flatten
swizzle hint allocs.
Add patterns which can fold reshapes and extract_slice ops into
empty ops through swizzle hint ops.
Add swizzle hint attribute to be set in lowering_config and
consumed in GPUPromoteMatmulOperandsPass.
Update LLVMGPUSelectLoweringStrategy Pass to set xor swizzles for
MXFP4 GEMMs.

This is PR 4, which does three things:

Expresses the row width as a function of CacheLineSizeInBits and the element type of the chosen intrinsic's operands.
Adds swizzle attribute to promotion type.
Adds a test for the swizzle attribute which should have been added in PR 3.

We see an average 4.8% geomean improvement over top of main in the mxfp4 gemms tested. See full data here.

shape	Feature_Throughput (TFLOps)	Top_of_main_Throughput(TFLOps)	improvement_ratio	improvement_percent
1000_1024_8192_512	96.044057	86.408117	1.111517	11.151661
16300_1024_8192_512	1282.544470	1219.842766	1.051401	5.140146
500_16384_26624_1664	713.748860	723.841063	0.986057	-1.394257
53200_53248_8192_512	1194.983741	1107.453688	1.079037	7.903721
8100_16384_8192_512	1388.739043	1301.227179	1.067253	6.725333
8192_512_256_16	305.040291	301.612872	1.011364	1.136364
8192_53248_8192_512	1734.982869	1728.772840	1.003592	0.359216
...	...	...	...	...
OVERALL_GEOMEAN			1.049751	4.975053

kuhar

What's the speedup after we land this?

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

Muzammiluddin-Syed-ECE · 2026-01-16T22:04:23Z

What's the speedup after we land this?

Roughly a Geomean speedup of %4.4 across shapes that previously had bank conflicts. Those that had no bank conflicts to begin with were unchanged.

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

This is part of a series of PR's implementing support for XOR swizzles in IREE. We require the LDS bank count to figure out XOR swizzle parameters. See PR: #23175 --------- Signed-off-by: Muzammiluddin Syed <[email protected]>

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp

Muzammiluddin-Syed-ECE · 2026-01-29T05:27:04Z

What's the speedup after we land this?

Roughly a Geomean speedup of %4.4 across shapes that previously had bank conflicts. Those that had no bank conflicts to begin with were unchanged.

Here's a gist showing the perf against most recent top of main: data.

kuhar · 2026-01-29T14:48:24Z

What's the speedup after we land this?

Roughly a Geomean speedup of %4.4 across shapes that previously had bank conflicts. Those that had no bank conflicts to begin with were unchanged.

Here's a gist showing the perf against most recent top of main: data.

I don't know how to read this data -- there are no units and it's unclear if positive change means improvement or regression.

Do you have data about bank conflicts? (Doesn't have to be exhaustive, just sample a dozen random shapes and see if the bank conflict counters align with the observed delta).

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

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp

Muzammiluddin-Syed-ECE · 2026-01-29T19:01:00Z

I don't know how to read this data -- there are no units and it's unclear if positive change means improvement or regression.

Do you have data about bank conflicts? (Doesn't have to be exhaustive, just sample a dozen random shapes and see if the bank conflict counters align with the observed delta).

Sorry I shouldve added to the file that the units were throughput so TFLOPs (higher better).

Also here is some data about bank conflicts on top of main.
https://gist.github.com/Muzammiluddin-Syed-ECE/88e8d52c4a9ecdf9f28c824c3388c937 .

Signed-off-by: Muzammiluddin Syed <[email protected]>

krzysz00

Nitpicks, but I think this makes sense overall as a step for making this feature happen

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp

Signed-off-by: Muzammiluddin Syed <[email protected]>

kuhar

Looks good overall. One more thing to consider is replacing all these pair<int64_t, int64_t> with a custom struct so that each integer is named

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp

Signed-off-by: Muzammiluddin Syed <[email protected]>

This is part of a series of PR's implementing support for XOR swizzles in IREE. We require the LDS bank count to figure out XOR swizzle parameters. See PR: #23175 --------- Signed-off-by: Muzammiluddin Syed <[email protected]>

Signed-off-by: Muzammiluddin Syed <[email protected]>

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h

Signed-off-by: Muzammiluddin Syed <[email protected]>

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h

Signed-off-by: Muzammiluddin Syed <[email protected]>

Muzammiluddin-Syed-ECE · 2026-02-09T01:03:11Z

the clang-tidy warnings are false positives (getXorShuffleAttr is used in ConfigUtils.cpp) or not related to this PR (getNSize and getKSize).

…n CAPI (#23442) After #23175, we now generate swizzle hint ops for scaled gemms whose parameters are set during `lowering_config` selection. These parameters (`rowElems` and `accessElems`) can be chosen via the tuner too (although in the future we intend for there to be an analytically derived solution to this). This PR exposes two functions, `getXorShuffleBounds` and `isXorShuffleValid`, to allow the tuner to constrain its search space for applicable XOR swizzles. Assisted by: composer-1 --------- Signed-off-by: Muzammiluddin Syed <[email protected]>

…rg#23273) This is part of a series of PR's implementing support for XOR swizzles in IREE. We require the LDS bank count to figure out XOR swizzle parameters. See PR: iree-org#23175 --------- Signed-off-by: Muzammiluddin Syed <[email protected]>

This is the fourth of a series of PRs that together implement support in IREE for XOR swizzling through the SwizzleHintOp. There are four PRs that need to be merged: 1) Allow rank > 1 swizzle hint op operands and add a pass to flatten swizzle hint allocs. 2) Add patterns which can fold reshapes and `extract_slice` ops into empty ops through swizzle hint ops. 3) Add swizzle hint attribute to be set in `lowering_config` and consumed in `GPUPromoteMatmulOperandsPass`. 4) Update `LLVMGPUSelectLoweringStrategy` Pass to set xor swizzles for MXFP4 GEMMs. This is PR 4, which does three things: - Expresses the row width as a function of CacheLineSizeInBits and the element type of the chosen intrinsic's operands. - Adds swizzle attribute to promotion type. - Adds a test for the swizzle attribute which should have been added in PR 3. We see an average 4.8% geomean improvement over top of main in the mxfp4 gemms tested. See [full data here](https://gist.github.com/Muzammiluddin-Syed-ECE/71c517206d89018f8706d661c94294b6/). |shape|Feature_Throughput (TFLOps)|Top_of_main_Throughput(TFLOps)|improvement_ratio|improvement_percent| |-----|---------------------------|------------------------------|-----------------|-------------------| |1000_1024_8192_512|96.044057|86.408117|1.111517|11.151661| |16300_1024_8192_512|1282.544470|1219.842766|1.051401|5.140146| |500_16384_26624_1664|713.748860|723.841063|0.986057|-1.394257| |53200_53248_8192_512|1194.983741|1107.453688|1.079037|7.903721| |8100_16384_8192_512|1388.739043|1301.227179|1.067253|6.725333| |8192_512_256_16|305.040291|301.612872|1.011364|1.136364| |8192_53248_8192_512|1734.982869|1728.772840|1.003592|0.359216| |...|...|...|...|...| |OVERALL_GEOMEAN|||1.049751|4.975053| --------- Signed-off-by: Muzammiluddin Syed <[email protected]>

…n CAPI (iree-org#23442) After iree-org#23175, we now generate swizzle hint ops for scaled gemms whose parameters are set during `lowering_config` selection. These parameters (`rowElems` and `accessElems`) can be chosen via the tuner too (although in the future we intend for there to be an analytically derived solution to this). This PR exposes two functions, `getXorShuffleBounds` and `isXorShuffleValid`, to allow the tuner to constrain its search space for applicable XOR swizzles. Assisted by: composer-1 --------- Signed-off-by: Muzammiluddin Syed <[email protected]>

Muzammiluddin-Syed-ECE requested review from Groverkss, Max191, krzysz00, kuhar, nirvedhmeshram and qedawkins as code owners January 16, 2026 21:30

Muzammiluddin-Syed-ECE force-pushed the muzasyed/config branch from b876715 to 40cb00e Compare January 16, 2026 21:33

kuhar reviewed Jan 16, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp Outdated Show resolved Hide resolved

compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp Outdated Show resolved Hide resolved

kuhar reviewed Jan 16, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp Outdated Show resolved Hide resolved

Muzammiluddin-Syed-ECE force-pushed the muzasyed/config branch from 8f86d07 to ae114d9 Compare January 16, 2026 21:35

kuhar reviewed Jan 16, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp Outdated Show resolved Hide resolved

kuhar reviewed Jan 16, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp Outdated Show resolved Hide resolved

kuhar requested changes Jan 16, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp Outdated Show resolved Hide resolved

krzysz00 reviewed Jan 17, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp Outdated Show resolved Hide resolved

compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp Outdated Show resolved Hide resolved

Muzammiluddin-Syed-ECE mentioned this pull request Jan 24, 2026

[GPU] Add workgroupMemoryBankCount parameter to TargetWgpAttr #23273

Merged

Muzammiluddin-Syed-ECE force-pushed the muzasyed/config branch from 3e34e2a to 4bf1aa3 Compare January 24, 2026 13:37

Muzammiluddin-Syed-ECE requested review from MaheshRavishankar, Yu-Zhewen and jtuyls as code owners January 24, 2026 13:37

Muzammiluddin-Syed-ECE force-pushed the muzasyed/config branch from 4bf1aa3 to 50dcd56 Compare January 28, 2026 22:34

Muzammiluddin-Syed-ECE force-pushed the muzasyed/config branch from 50dcd56 to 96b35b3 Compare January 28, 2026 23:07

Muzammiluddin-Syed-ECE commented Jan 28, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp Show resolved Hide resolved

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp Outdated Show resolved Hide resolved

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp Outdated Show resolved Hide resolved

Muzammiluddin-Syed-ECE force-pushed the muzasyed/config branch from 0a6a045 to 7a36eed Compare January 29, 2026 03:22

kuhar reviewed Jan 29, 2026

View reviewed changes

Muzammiluddin-Syed-ECE force-pushed the muzasyed/config branch from 4066376 to 2c8dc7e Compare January 29, 2026 20:16

Muzammiluddin-Syed-ECE force-pushed the muzasyed/config branch from b5ad9e5 to affc0a2 Compare February 6, 2026 18:30

Muzammiluddin-Syed-ECE added 3 commits February 6, 2026 12:33

updating lint disable

d5feff5

Signed-off-by: Muzammiluddin Syed <[email protected]>

removing unrelated change

fef5a1c

Signed-off-by: Muzammiluddin Syed <[email protected]>

build fix

8314808

Signed-off-by: Muzammiluddin Syed <[email protected]>

krzysz00 approved these changes Feb 6, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp Outdated Show resolved Hide resolved

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp Show resolved Hide resolved

resolving nits

565be8b

Signed-off-by: Muzammiluddin Syed <[email protected]>

kuhar reviewed Feb 6, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp Outdated Show resolved Hide resolved

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp Outdated Show resolved Hide resolved

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp Outdated Show resolved Hide resolved

Muzammiluddin-Syed-ECE added 3 commits February 6, 2026 13:40

nit fix

c29bc7d

Signed-off-by: Muzammiluddin Syed <[email protected]>

nit fix

133396f

Signed-off-by: Muzammiluddin Syed <[email protected]>

build fix

361fe86

Signed-off-by: Muzammiluddin Syed <[email protected]>

adding struct for xor shuffle params

8c098e9

Signed-off-by: Muzammiluddin Syed <[email protected]>

Muzammiluddin-Syed-ECE requested a review from kuhar February 7, 2026 04:59

kuhar reviewed Feb 7, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h Outdated Show resolved Hide resolved

kuhar reviewed Feb 7, 2026

View reviewed changes

compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h Outdated Show resolved Hide resolved

addressing pr comments pt 5

56b156b

Signed-off-by: Muzammiluddin Syed <[email protected]>

kuhar reviewed Feb 8, 2026

View reviewed changes

Muzammiluddin-Syed-ECE force-pushed the muzasyed/config branch 2 times, most recently from 6299c66 to 9929999 Compare February 8, 2026 20:29

addressing pr comments pt 6

6219326

Signed-off-by: Muzammiluddin Syed <[email protected]>

Muzammiluddin-Syed-ECE force-pushed the muzasyed/config branch from 9929999 to 6219326 Compare February 9, 2026 00:20

Muzammiluddin-Syed-ECE requested a review from kuhar February 9, 2026 00:54

kuhar approved these changes Feb 9, 2026

View reviewed changes

Muzammiluddin-Syed-ECE merged commit a9fd31b into iree-org:main Feb 9, 2026
57 checks passed

Muzammiluddin-Syed-ECE mentioned this pull request Feb 10, 2026

[Codegen][Tuner] Expose XOR shuffle bounds and validation functions in CAPI #23442

Merged

Conversation

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

Uh oh!

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Muzammiluddin-Syed-ECE commented Jan 29, 2026

Uh oh!

kuhar commented Jan 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Muzammiluddin-Syed-ECE commented Jan 29, 2026

Uh oh!

krzysz00 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

kuhar 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!

Muzammiluddin-Syed-ECE commented Feb 9, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Muzammiluddin-Syed-ECE commented Jan 16, 2026 •

edited

Loading

Muzammiluddin-Syed-ECE commented Jan 16, 2026 •

edited

Loading

kuhar commented Jan 29, 2026 •

edited

Loading