add back fusion in reassociated quantized matmul

nod-ai · Dec 1, 2023 · 40d58f0 · 40d58f0
1 parent 2dc897a
commit 40d58f0
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/compiler/src/iree/compiler/GlobalOptimization/FuseDequantizationMatmul.cpp b/compiler/src/iree/compiler/GlobalOptimization/FuseDequantizationMatmul.cpp
@@ -800,6 +800,14 @@ static LogicalResult reassociateDequantMatmul(RewriterBase &rewriter,
 
   rewriter.replaceOp(matmul, reassociatedDequantization.getResult(0));
 
+  // Fuse dequantization + matmul ops into a single dispatch region
+  SmallVector<Operation *> dequantMatmulOps{quantizedIntegerMatmul,
+                                            reassociatedDequantization};
+  FailureOr<IREE::Flow::DispatchRegionOp> maybeDequantMatmulDispatch =
+      wrapConsecutiveOpsInDispatchRegion(rewriter, dequantMatmulOps);
+  if (failed(maybeDequantMatmulDispatch)) {
+    return failure();
+  }
   return success();
 }