[mlir][DispatchCreation] Avoid cloning of index_cast operations if they are operands to the dispatch. (#18374)

Cloning an `index_cast` op which is already a operand to the dispatch is
unnecessary. This op will need to be on the host anyway cause it
represents the shape of the input/output tensors.

Fixes #18364

---------

Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>
Co-authored-by: Nirvedh Meshram <nirvedh@gmail.com>

compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp
index 175f822..8ca036d 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.cpp
@@ -761,6 +761,15 @@
 static bool hasUnfusableUseInDispatch(Value v, Operation *dispatchOp) {
   for (OpOperand &use : v.getUses()) {
     Operation *user = use.getOwner();
+
+    // Do not fuse `index_cast` operations if it is already an operand of the
+    // owner.
+    if (auto indexCastOp = v.getDefiningOp<arith::IndexCastOp>()) {
+      if (user == dispatchOp) {
+        return true;
+      }
+    }
+
     Operation *ownerWorkgroupsOp =
         user->getParentOfType<IREE::Flow::DispatchWorkgroupsOp>();
     Operation *ownerRegionOp =

compiler/src/iree/compiler/DispatchCreation/test/clone_producers_into_dispatch_regions.mlir

diff --git a/compiler/src/iree/compiler/DispatchCreation/test/clone_producers_into_dispatch_regions.mlir b/compiler/src/iree/compiler/DispatchCreation/test/clone_producers_into_dispatch_regions.mlir
index 3c65e7a..713ae74 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/clone_producers_into_dispatch_regions.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/clone_producers_into_dispatch_regions.mlir
@@ -326,3 +326,19 @@
 //  CHECK-SAME:         ins(%[[DEQUANT]] :
 //       CHECK:     flow.return %[[REDUCE]]
 //       CHECK:   return %[[RETURN]]
+
+// -----
+
+// Do no clone index cast operations when they are operands to the dispatch
+util.func public @dont_clone_index_cast(%arg0 : i64) {
+  %0 = arith.index_cast %arg0 : i64 to index
+  %1 = flow.dispatch.region[] -> (tensor<?xf32>{%0}) {
+    %2 = tensor.empty(%0) : tensor<?xf32>
+    flow.return %2 : tensor<?xf32>
+  }
+  util.return
+}
+// CHECK-LABEL: func public @dont_clone_index_cast
+//       CHECK:   arith.index_cast
+//       CHECK:   flow.dispatch.region
+//   CHECK-NOT:   arith.index_cast