[Dispatch] Disable scatter fusion with producers (#19565)

Backends don't currently support scatter fusion and will silently compile incorrect code. This should be turned off in order to prevent backends from generating incorrect results. I don't think any users are running into this currently, but its best to keep it off for now. Similar to #19535 but for both `indices` and `updates`. --------- Signed-off-by: Ian Wood <[email protected]>
iree-org · Dec 27, 2024 · a43d893 · a43d893
1 parent d746a57
commit a43d893
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp b/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp
@@ -651,7 +651,8 @@ isFusableWithProducer(OpOperand &operand,
   }
 
   // Don't fuse attention with it's producer
-  if (isa<IREE::LinalgExt::AttentionOp>(consumer)) {
+  // TODO: Enable scatter fusion when supported by backends.
+  if (isa<IREE::LinalgExt::AttentionOp, IREE::LinalgExt::ScatterOp>(consumer)) {
     return false;
   }
 

diff --git a/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_ext_fusion.mlir b/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_ext_fusion.mlir
@@ -37,9 +37,9 @@ util.func public @linalgext_scatter_dispatch() -> tensor<8192x16x8x128xf32> {
 }
 
 // CHECK-LABEL:     util.func public @linalgext_scatter_dispatch
+//   CHECK-DAG:       %[[INDICES:.+]] = flow.dispatch.region
+//   CHECK-DAG:       %[[UPDATE:.+]] = flow.dispatch.region
 //       CHECK:       %[[RESULT:.+]] = flow.dispatch.region
-//       CHECK:         %[[INDICES:.+]] = linalg.generic
-//       CHECK:         %[[UPDATE:.+]] = linalg.generic
 //       CHECK:         %[[SCATTER_RESULT:.+]] = iree_linalg_ext.scatter
 //  CHECK-SAME:           ins(%[[UPDATE]], %[[INDICES]] : tensor<4x1x16x8x128xf32>, tensor<4x1xi32>)
 //       CHECK:         flow.return %[[SCATTER_RESULT]]