Fixing tied op indexing in the copy-on-write pass. (#7643)

iree-org · Nov 12, 2021 · c9ff825 · c9ff825
1 parent fabc494
commit c9ff825
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/iree/compiler/Dialect/Stream/Transforms/MaterializeCopyOnWrite.cpp b/iree/compiler/Dialect/Stream/Transforms/MaterializeCopyOnWrite.cpp
@@ -99,12 +99,11 @@ static bool materializeTiedOpCOW(IREE::Util::TiedOpInterface tiedOp) {
 
   // Clones each operand that is tied to a result and it may be required.
   OpBuilder builder(tiedOp);
-  unsigned tiedOperandsOffset = tiedOp.getTiedOperandsIndexAndLength().first;
   auto tiedOperandIndices = tiedOp.getTiedResultOperandIndices();
   for (unsigned i = 0; i < tiedOperandIndices.size(); ++i) {
     int64_t operandIdx = tiedOperandIndices[i];
     if (operandIdx == IREE::Util::TiedOpInterface::kUntiedIndex) continue;
-    auto &operand = tiedOp->getOpOperand(tiedOperandsOffset + operandIdx);
+    auto &operand = tiedOp->getOpOperand(operandIdx);
     didChange =
         materializeOperandCOW(tiedOp.getLoc(), operand, affinity, builder) ||
         didChange;

diff --git a/iree/compiler/Dialect/Stream/Transforms/test/materialize_copy_on_write.mlir b/iree/compiler/Dialect/Stream/Transforms/test/materialize_copy_on_write.mlir
@@ -68,6 +68,37 @@ func @multiUseTiedOperand(%size: index) -> (!stream.resource<*>, !stream.resourc
 
 // -----
 
+// Tests tied dispatches with a data dependency.
+// %splat0 is mutated by @dispatch0 and a clone gets inserted to preserve its
+// original contents for use by @dispatch1.
+
+// CHECK-LABEL: @tiedDispatches
+func private @tiedDispatches() {
+  %c0_i32 = arith.constant 0 : i32
+  %c1_i32 = arith.constant 1 : i32
+  %c1 = arith.constant 1 : index
+  %c16 = arith.constant 16 : index
+  %c40 = arith.constant 40 : index
+
+  // CHECK: %[[SPLAT0:.+]] = stream.async.splat %c0_i32
+  %splat0 = stream.async.splat %c0_i32 : i32 -> !stream.resource<*>{%c40}
+  // CHECK: %[[SPLAT1:.+]] = stream.async.splat %c1_i32
+  %splat1 = stream.async.splat %c1_i32 : i32 -> !stream.resource<*>{%c16}
+
+  // CHECK: %[[CLONE0:.+]] = stream.async.clone %[[SPLAT0]]
+  // CHECK: %[[DISPATCH0:.+]] = stream.async.dispatch @ex::@dispatch0[%c1, %c1, %c1](%[[CLONE0]], %[[SPLAT1]])
+  // CHECK-SAME: (!stream.resource<*>{%c40}, !stream.resource<*>{%c16}) -> %[[CLONE0]]{%c40}
+  %dispatch0 = stream.async.dispatch @ex::@dispatch0[%c1, %c1, %c1](%splat0, %splat1) : (!stream.resource<*>{%c40}, !stream.resource<*>{%c16}) -> %splat0{%c40}
+
+  // CHECK: %[[DISPATCH1:.+]] = stream.async.dispatch @ex::@dispatch1[%c1, %c1, %c1](%[[DISPATCH0]], %[[SPLAT0]])
+  // CHECK-SAME: (!stream.resource<*>{%c40}, !stream.resource<*>{%c40}) -> %[[DISPATCH0]]{%c40}
+  %dispatch1 = stream.async.dispatch @ex::@dispatch1[%c1, %c1, %c1](%dispatch0, %splat0) : (!stream.resource<*>{%c40}, !stream.resource<*>{%c40}) -> %dispatch0{%c40}
+
+  return
+}
+
+// -----
+
 // Tests that block args (like function args) are copied until copy elision can
 // take care of them later.