diff --git a/compiler/src/iree/compiler/DispatchCreation/ElementwiseOpFusion.cpp b/compiler/src/iree/compiler/DispatchCreation/ElementwiseOpFusion.cpp index 3c1a783ecba3..46eb0f65719c 100644 --- 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(context); linalg::populateElementwiseOpsFusionPatterns(linalgFusionPatterns, fuseElementwiseOpsControlFn); @@ -158,7 +159,6 @@ void ElementwiseOpFusionPass::runOnOperation() { RewritePatternSet linalgExtFusionPatterns(context); IREE::LinalgExt::populateFuseLinalgExtOpsWithTransposes( linalgExtFusionPatterns, foldTransposeControlFn); - linalgExtFusionPatterns.insert(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 index c428091f6cf8..3579e0ab0a2b 100644 --- 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; }