[Dispatch] Don't fuse bit truncate -> extend ops

Signed-off-by: Ian Wood <[email protected]>
iree-org · Dec 6, 2024 · 9e22402 · 9e22402
1 parent 9f8aad8
commit 9e22402
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/compiler/src/iree/compiler/DispatchCreation/ElementwiseOpFusion.cpp b/compiler/src/iree/compiler/DispatchCreation/ElementwiseOpFusion.cpp
@@ -136,6 +136,7 @@ void ElementwiseOpFusionPass::runOnOperation() {
       };
 
   RewritePatternSet linalgFusionPatterns(context);
+  linalgFusionPatterns.insert<GatherFusionPattern>(context);
   linalg::populateElementwiseOpsFusionPatterns(linalgFusionPatterns,
                                                fuseElementwiseOpsControlFn);
 
@@ -158,7 +159,6 @@ void ElementwiseOpFusionPass::runOnOperation() {
   RewritePatternSet linalgExtFusionPatterns(context);
   IREE::LinalgExt::populateFuseLinalgExtOpsWithTransposes(
       linalgExtFusionPatterns, foldTransposeControlFn);
-  linalgExtFusionPatterns.insert<GatherFusionPattern>(context);
   if (failed(applyPatternsAndFoldGreedily(
           getOperation(), std::move(linalgExtFusionPatterns), rewriteConfig))) {
     getOperation()->emitOpError(

diff --git a/compiler/src/iree/compiler/DispatchCreation/FusionUtils.cpp b/compiler/src/iree/compiler/DispatchCreation/FusionUtils.cpp
@@ -65,7 +65,9 @@ bool areFusableAsElementwiseOps(MLIRContext *context, OpOperand *fusedOperand,
   // (except for bit-extend ops). If the consumer has only one use, then this
   // fusion is fine since cloning wont result in redundant computation of the
   // producer. (Also note that the producer is always an elementwise operation).
-  if (IREE::LinalgExt::isBitExtendOp(consumerOp) && !consumerOp->hasOneUse()) {
+  if (IREE::LinalgExt::isBitExtendOp(consumerOp) &&
+      (!consumerOp->hasOneUse() ||
+       IREE::LinalgExt::isBitTruncateOp(producerOp))) {
     return false;
   }