diff --git a/iree/compiler/Dialect/Stream/Transforms/MaterializeCopyOnWrite.cpp b/iree/compiler/Dialect/Stream/Transforms/MaterializeCopyOnWrite.cpp index b5fc0eaf320e..f6392f5e051f 100644 --- 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 index db380786a4d2..afa0584aa82c 100644 --- 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.