[Global Opt] Don't propagate edge reshapes by IanWood1 · Pull Request #22320 · iree-org/iree

IanWood1 · 2025-10-15T19:25:24Z

Prevent propagating reshapes on the edges of the program in PropagateLinalgTranspose since these reshapes don't block the propagation of transposes. This is similar to what is done in BubbleUpExpandShapes:

iree/compiler/src/iree/compiler/DispatchCreation/BubbleUpExpandShapes.cpp

Lines 441 to 459 in dad4b2d

    
           // Don't reintroduce unit dims via propagating edge unit dim reshapes. 
        
           if (auto expandOp = dyn_cast<tensor::ExpandShapeOp>(consumer); 
        
               expandOp && 
        
               isExpandingUnitDims(expandOp.getReassociationIndices(), 
        
                                   expandOp.getResultType().getShape()) && 
        
               llvm::none_of( 
        
                   consumer->getUsers(), [&](Operation *collapseConsumer) { 
        
                     return isReshapeBlockingFusion(producer, collapseConsumer); 
        
                   })) { 
        
             return false; 
        
           } 
        
           if (auto collapseOp = dyn_cast<tensor::CollapseShapeOp>(producer); 
        
               collapseOp && 
        
               isExpandingUnitDims(collapseOp.getReassociationIndices(), 
        
                                   collapseOp.getSrcType().getShape()) && 
        
               !isReshapeBlockingFusion(producer->getOperand(0).getDefiningOp(), 
        
                                        consumer)) { 
        
             return false; 
        
           }

Closes #22312

Signed-off-by: Ian Wood <ianwood@u.northwestern.edu>

MaheshRavishankar · 2025-10-15T22:42:05Z

compiler/src/iree/compiler/GlobalOptimization/test/propagate_linalg_transpose.mlir


 // -----

-util.func public @sink_non_involution_through_expand_shape(%arg0 : tensor<2x3x4xf32>) -> tensor<1x3x4x2xf32> {


I think we want to keep this behavior... Cause this might allow the transpose to fuse with its consumers...

The point is that the reshapes are on the edge so that there will be no producer/consumers.

I am not sure I follow. The test is changing the behavior that is different than what I assume the end state of the code should be. Could you provide more details here?

The problem is that BubbleUpExpandShapes is not perfect and doesn't currently sink reshapes through reduction operations. The idea is to not touch reshapes that are on the edges of the program since they might get stuck on a reduction operation. This shouldn't harm transpose propagation because these reshapes are on the edges and won't block any transposes from getting propagated throughout the program.

Also, I added some more context on the issue with the mlir before and after this pass #22312

I think the better solution is make a change to BubbleUpExpandShapes, I'm testing if there are any regressions #22341

MaheshRavishankar · 2025-10-15T22:42:22Z

compiler/src/iree/compiler/GlobalOptimization/test/propagate_linalg_transpose.mlir

-// -----
-
-util.func public @bubble_non_involution_through_collapse_shape(%arg0 : tensor<1x2x3x5x7x11xf32>) -> tensor<35x11x6xf32> {
-  %collapsed = tensor.collapse_shape %arg0 [[0, 1, 2], [3, 4], [5]] : tensor<1x2x3x5x7x11xf32> into tensor<6x35x11xf32>


Same here. The transpose could fuse with its producers.

compiler/src/iree/compiler/GlobalOptimization/test/propagate_linalg_transpose.mlir

MaheshRavishankar · 2025-10-15T22:43:26Z

Could you post some more details of the problem (preferably on the issue)

Signed-off-by: Ian Wood <ianwood@u.northwestern.edu>

MaheshRavishankar

Lets land this for now based on what we see as issues with a "more proper way". We would probably need to revisit this though.

Prevent propagating reshapes on the edges of the program in PropagateLinalgTranspose since these reshapes don't block the propagation of transposes. This is similar to what is done in BubbleUpExpandShapes: https://github.com/iree-org/iree/blob/dad4b2d1cab357720a73c05d7bf6e1b946668082/compiler/src/iree/compiler/DispatchCreation/BubbleUpExpandShapes.cpp#L441-L459 Closes iree-org#22312 --------- Signed-off-by: Ian Wood <ianwood@u.northwestern.edu> Signed-off-by: Philipp <philipp.weidel@intel.com>

…nv3d (#2560) Compiler fix: iree-org/iree#22320 Issue: iree-org/iree#22312

…nv3d (ROCm#2560) Compiler fix: iree-org/iree#22320 Issue: iree-org/iree#22312

…nv3d (#2560) Compiler fix: iree-org/iree#22320 Issue: iree-org/iree#22312

Adds flag to allow users to selectively disable the changes from #22320. This is a stopgap for cases where not fusing a transpose into a matmul gives better performance. Signed-off-by: Ian Wood <ianwood@u.northwestern.edu>

…2438) Adds flag to allow users to selectively disable the changes from iree-org#22320. This is a stopgap for cases where not fusing a transpose into a matmul gives better performance. Signed-off-by: Ian Wood <ianwood@u.northwestern.edu>

Compiler fix: iree-org/iree#22320 Issue: iree-org/iree#22312

…2438) Adds flag to allow users to selectively disable the changes from iree-org#22320. This is a stopgap for cases where not fusing a transpose into a matmul gives better performance. Signed-off-by: Ian Wood <ianwood@u.northwestern.edu>

Prevent propagating reshapes on the edges of the program in PropagateLinalgTranspose since these reshapes don't block the propagation of transposes. This is similar to what is done in BubbleUpExpandShapes: https://github.com/iree-org/iree/blob/dad4b2d1cab357720a73c05d7bf6e1b946668082/compiler/src/iree/compiler/DispatchCreation/BubbleUpExpandShapes.cpp#L441-L459 Closes iree-org#22312 --------- Signed-off-by: Ian Wood <ianwood@u.northwestern.edu>

…2438) Adds flag to allow users to selectively disable the changes from iree-org#22320. This is a stopgap for cases where not fusing a transpose into a matmul gives better performance. Signed-off-by: Ian Wood <ianwood@u.northwestern.edu>

[Global Opt] Don't propagate edge reshapes

e3dbbff

Signed-off-by: Ian Wood <ianwood@u.northwestern.edu>

IanWood1 requested a review from hanhanW as a code owner October 15, 2025 19:25

IanWood1 requested a review from MaheshRavishankar October 15, 2025 19:27

MaheshRavishankar requested changes Oct 15, 2025

View reviewed changes

Remove unit from name and add expandshape test

572bd7a

Signed-off-by: Ian Wood <ianwood@u.northwestern.edu>

MaheshRavishankar approved these changes Oct 17, 2025

View reviewed changes

IanWood1 merged commit 189dc86 into iree-org:main Oct 17, 2025
46 checks passed

sjain-stanford mentioned this pull request Oct 20, 2025

[Fusilli] Pick compiler fixes and change dispatch count checks for conv3d nod-ai/amd-shark-ai#2560

Merged

sjain-stanford added a commit to nod-ai/amd-shark-ai that referenced this pull request Oct 23, 2025

[Fusilli] Pick compiler fixes and change dispatch count checks for co…

52e12fe

…nv3d (#2560) Compiler fix: iree-org/iree#22320 Issue: iree-org/iree#22312

AaronStGeorge pushed a commit to AaronStGeorge/rocm-libraries that referenced this pull request Oct 23, 2025

[Fusilli] Pick compiler fixes and change dispatch count checks for co…

53b3415

…nv3d (ROCm#2560) Compiler fix: iree-org/iree#22320 Issue: iree-org/iree#22312

AaronStGeorge pushed a commit to AaronStGeorge/rocm-libraries that referenced this pull request Oct 24, 2025

[Fusilli] Pick compiler fixes and change dispatch count checks for co…

55b156d

…nv3d (ROCm#2560) Compiler fix: iree-org/iree#22320 Issue: iree-org/iree#22312

AaronStGeorge pushed a commit to ROCm/rocm-libraries that referenced this pull request Oct 24, 2025

[Fusilli] Pick compiler fixes and change dispatch count checks for co…

37c118f

…nv3d (#2560) Compiler fix: iree-org/iree#22320 Issue: iree-org/iree#22312

IanWood1 mentioned this pull request Oct 28, 2025

[Global Opt] Add flag to control edge reshape propagation #22438

Merged

AaronStGeorge pushed a commit to AaronStGeorge/fusilli that referenced this pull request Nov 18, 2025

Pick compiler fixes and change dispatch count checks for conv3d (#2560)

aefdea0

Compiler fix: iree-org/iree#22320 Issue: iree-org/iree#22312

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Global Opt] Don't propagate edge reshapes#22320

[Global Opt] Don't propagate edge reshapes#22320
IanWood1 merged 2 commits intoiree-org:mainfrom
IanWood1:tpose_prop_edge_reshapes

IanWood1 commented Oct 15, 2025

Uh oh!

MaheshRavishankar Oct 15, 2025

Uh oh!

IanWood1 Oct 15, 2025

Uh oh!

MaheshRavishankar Oct 16, 2025

Uh oh!

IanWood1 Oct 16, 2025

Uh oh!

IanWood1 Oct 17, 2025

Uh oh!

MaheshRavishankar Oct 15, 2025

Uh oh!

Uh oh!

MaheshRavishankar commented Oct 15, 2025

Uh oh!

MaheshRavishankar left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	// Don't reintroduce unit dims via propagating edge unit dim reshapes.
	if (auto expandOp = dyn_cast<tensor::ExpandShapeOp>(consumer);
	expandOp &&
	isExpandingUnitDims(expandOp.getReassociationIndices(),
	expandOp.getResultType().getShape()) &&
	llvm::none_of(
	consumer->getUsers(), [&](Operation *collapseConsumer) {
	return isReshapeBlockingFusion(producer, collapseConsumer);
	})) {
	return false;
	}
	if (auto collapseOp = dyn_cast<tensor::CollapseShapeOp>(producer);
	collapseOp &&
	isExpandingUnitDims(collapseOp.getReassociationIndices(),
	collapseOp.getSrcType().getShape()) &&
	!isReshapeBlockingFusion(producer->getOperand(0).getDefiningOp(),
	consumer)) {
	return false;
	}


		// -----

		util.func public @sink_non_involution_through_expand_shape(%arg0 : tensor<2x3x4xf32>) -> tensor<1x3x4x2xf32> {

Conversation

IanWood1 commented Oct 15, 2025

Uh oh!

MaheshRavishankar Oct 15, 2025

Choose a reason for hiding this comment

Uh oh!

IanWood1 Oct 15, 2025

Choose a reason for hiding this comment

Uh oh!

MaheshRavishankar Oct 16, 2025

Choose a reason for hiding this comment

Uh oh!

IanWood1 Oct 16, 2025

Choose a reason for hiding this comment

Uh oh!

IanWood1 Oct 17, 2025

Choose a reason for hiding this comment

Uh oh!

MaheshRavishankar Oct 15, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

MaheshRavishankar commented Oct 15, 2025

Uh oh!

MaheshRavishankar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants