[Stream] Model execution affinity for cross-device transfers (#24748)
Fixes #21361.
A `stream.async.transfer` currently carries source and target
affinities, but those describe where its input and result reside, not
necessarily where the transfer executes.
This distinction matters when a value produced on device B is
transferred into a result pinned to device A. The problems this can
cause are described in detail in the issue linked above.
This change adds an optional execution affinity to
`stream.async.transfer`, printed as `on(...)`, alongside the existing
`from(...)` and `to(...)` affinities:
```mlir
%result = stream.async.transfer
on(#hal.device.affinity<@device_a>)
%source
from(#hal.device.affinity<@device_b>)
-> to(#hal.device.affinity<@device_a>)
!stream.resource<external>
```
A new PlaceTransferExecutionsPass runs before execution scheduling. It uses affinity analysis to identify transfers whose result has a unique pinned affinity and places those transfers on the pinned device when the default producer-side placement is incompatible.
ScheduleExecutionPass then handles the transfer through its normal affinity partitioning. Transfers are only grouped with their producer when their execution affinities are compatible.
The existing staging-transfer placement rules are preserved as the default when no explicit execution affinity is present. Source and target affinities are no longer modified indirectly when setting an operation's execution affinity.
The full Stream pipeline now produces the desired result for the original multi-device case:
- The device-B dispatch result is allocated with optimal<A, B> affinity.
- The final external result remains pinned to device A.
- The transfer into that result executes on device A.
With this https://github.com/iree-org/iree/pull/21589 can also be closed as this implementation should solve the problem.
---------
Signed-off-by: default <ziereis@roofline.ai>diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp
index c858880..12f4a03 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp
@@ -321,7 +321,16 @@
auto producer = op.getOperand(0).getDefiningOp();
auto streamable =
dyn_cast_if_present<IREE::Stream::StreamableOpInterface>(producer);
- if (streamable) {
+ bool hasCompatibleAffinity = true;
+ if (auto transferOp = dyn_cast<IREE::Stream::AsyncTransferOp>(op)) {
+ auto producerAffinityOp =
+ dyn_cast_if_present<IREE::Stream::AffinityOpInterface>(producer);
+ hasCompatibleAffinity = !producerAffinityOp ||
+ IREE::Stream::AffinityAttr::canExecuteTogether(
+ transferOp.getAffinityAttr(),
+ producerAffinityOp.getAffinityAttr());
+ }
+ if (streamable && hasCompatibleAffinity) {
if (!syncOps.contains(producer)) {
syncOps[producer] = llvm::SmallVector<Operation *>();
}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
index 42f7041..62a53e6 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
@@ -2163,7 +2163,9 @@
OpFoldResult AsyncTransferOp::fold(FoldAdaptor operands) {
if (auto sourceTransferOp = getSource().getDefiningOp<AsyncTransferOp>()) {
- if (sourceTransferOp.getSource().getType() == getResult().getType() &&
+ if (!getExecutionAffinityAttr() &&
+ !sourceTransferOp.getExecutionAffinityAttr() &&
+ sourceTransferOp.getSource().getType() == getResult().getType() &&
sourceTransferOp.getSourceAffinity() == getTargetAffinity()) {
return sourceTransferOp.getSource();
}
@@ -2234,7 +2236,8 @@
transferOp, transferOp.getResult().getType(),
originTransferOp.getSource(), originTransferOp.getSourceSize(),
transferOp.getResultSize(), originTransferOp.getSourceAffinityAttr(),
- transferOp.getResultAffinityAttr());
+ transferOp.getResultAffinityAttr(),
+ transferOp.getExecutionAffinityAttr());
return success();
}
};
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
index 716cc37..dd6b745 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
@@ -2972,6 +2972,13 @@
}
IREE::Stream::AffinityAttr AsyncTransferOp::getAffinityAttr() {
+ if (auto executionAffinityAttr = getExecutionAffinityAttr()) {
+ return executionAffinityAttr;
+ }
+ return getDefaultExecutionAffinityAttr();
+}
+
+IREE::Stream::AffinityAttr AsyncTransferOp::getDefaultExecutionAffinityAttr() {
auto sourceType = cast<IREE::Stream::ResourceType>(getSource().getType());
auto resultType = cast<IREE::Stream::ResourceType>(getResult().getType());
if (sourceType.getLifetime() == IREE::Stream::Lifetime::Staging &&
@@ -2993,37 +3000,10 @@
}
void AsyncTransferOp::setAffinityAttr(IREE::Stream::AffinityAttr value) {
- auto sourceType = cast<IREE::Stream::ResourceType>(getSource().getType());
- auto resultType = cast<IREE::Stream::ResourceType>(getResult().getType());
- if (sourceType.getLifetime() == IREE::Stream::Lifetime::Staging &&
- resultType.getLifetime() == IREE::Stream::Lifetime::Staging) {
- // TODO(multi-device): figure out how to model staging->staging transfers.
- if (value) {
- setSourceAffinityAttr(value);
- } else {
- removeSourceAffinityAttr();
- }
- } else if (sourceType.getLifetime() == IREE::Stream::Lifetime::Staging) {
- // If source is staging then the op should execute on the consumer.
- if (value) {
- setTargetAffinityAttr(value);
- } else {
- removeTargetAffinityAttr();
- }
- } else if (resultType.getLifetime() == IREE::Stream::Lifetime::Staging) {
- // If result is staging then the op should execute on the producer.
- if (value) {
- setSourceAffinityAttr(value);
- } else {
- removeSourceAffinityAttr();
- }
+ if (value) {
+ setExecutionAffinityAttr(value);
} else {
- // Default to result affinity.
- if (value) {
- setTargetAffinityAttr(value);
- } else {
- removeTargetAffinityAttr();
- }
+ removeExecutionAffinityAttr();
}
}
@@ -3034,6 +3014,24 @@
void AsyncTransferOp::removeAffinityAttrs() {
removeSourceAffinityAttr();
removeTargetAffinityAttr();
+ removeExecutionAffinityAttr();
+}
+
+AsyncTransferOp AsyncTransferOp::create(OpBuilder &builder, Location location,
+ Type type, Value source,
+ Value sourceSize, Value resultSize,
+ AffinityAttr sourceAffinity,
+ AffinityAttr targetAffinity) {
+ return create(builder, location, type, source, sourceSize, resultSize,
+ sourceAffinity, targetAffinity, nullptr);
+}
+
+void AsyncTransferOp::build(OpBuilder &builder, OperationState &state,
+ Type type, Value source, Value sourceSize,
+ Value resultSize, AffinityAttr sourceAffinity,
+ AffinityAttr targetAffinity) {
+ build(builder, state, type, source, sourceSize, resultSize, sourceAffinity,
+ targetAffinity, nullptr);
}
void AsyncTransferOp::getAsyncAccessRanges(
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
index f74ecb4..4d1fb78 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
@@ -2524,6 +2524,10 @@
to a `local` lifetime) or different affinities. This is roughly equivalent
to a cast but may have special semantics when later lowered to one or more
devices with discrete memory spaces or pools.
+
+ The optional execution affinity specifies where the transfer must execute
+ independently of the source and target storage affinities. When omitted an
+ execution affinity is derived from the resource lifetimes and endpoints.
}];
let arguments = (ins
@@ -2534,7 +2538,8 @@
Stream_Size:$source_size,
Stream_Size:$result_size,
OptionalAttr<Stream_AffinityAttr>:$source_affinity,
- OptionalAttr<Stream_AffinityAttr>:$target_affinity
+ OptionalAttr<Stream_AffinityAttr>:$target_affinity,
+ OptionalAttr<Stream_AffinityAttr>:$execution_affinity
);
let results = (outs
AnyTypeOf<[
@@ -2544,6 +2549,7 @@
);
let assemblyFormat = [{
+ (`on` `(` $execution_affinity^ `)`)?
$source `:` type($source)
`` `{` $source_size `}`
(`from` `(` $source_affinity^ `)`)?
@@ -2554,6 +2560,15 @@
}];
let extraClassDeclaration = [{
+ static AsyncTransferOp create(OpBuilder &builder, Location location,
+ Type type, Value source, Value sourceSize,
+ Value resultSize, AffinityAttr sourceAffinity,
+ AffinityAttr targetAffinity);
+ static void build(OpBuilder &builder, OperationState &state, Type type,
+ Value source, Value sourceSize, Value resultSize,
+ AffinityAttr sourceAffinity,
+ AffinityAttr targetAffinity);
+ AffinityAttr getDefaultExecutionAffinityAttr();
Value getOperandSize(unsigned idx) { return getSourceSize(); }
Value getResultSize(unsigned idx) { return getResultSize(); }
}];
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_folding.mlir b/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_folding.mlir
index 49a3a32..3cd1613 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_folding.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_folding.mlir
@@ -450,6 +450,22 @@
// -----
+util.global private @device_a : !hal.device
+util.global private @device_b : !hal.device
+
+// CHECK-LABEL: @PreserveTransferExecutionAffinity
+// CHECK-SAME: (%[[SOURCE:.+]]: !stream.resource<transient>, %[[SIZE:.+]]: index)
+util.func private @PreserveTransferExecutionAffinity(%source: !stream.resource<transient>, %size: index) -> !stream.resource<external> {
+ %transfer0 = stream.async.transfer %source : !stream.resource<transient>{%size} from(#hal.device.affinity<@device_b>) -> to(#hal.device.affinity<@device_a>) !stream.resource<staging>{%size}
+ // CHECK: stream.async.transfer on(#hal.device.affinity<@device_a>) %[[SOURCE]]
+ // CHECK-SAME: from(#hal.device.affinity<@device_b>)
+ // CHECK-SAME: to(#hal.device.affinity<@device_a>)
+ %transfer1 = stream.async.transfer on(#hal.device.affinity<@device_a>) %transfer0 : !stream.resource<staging>{%size} from(#hal.device.affinity<@device_a>) -> to(#hal.device.affinity<@device_a>) !stream.resource<external>{%size}
+ util.return %transfer1 : !stream.resource<external>
+}
+
+// -----
+
// CHECK-LABEL: @FoldAsyncCastSameType
// CHECK-SAME: (%[[SOURCE:.+]]: !stream.resource<external>, %[[SIZE:.+]]: index)
util.func private @FoldAsyncCastSameType(%source: !stream.resource<external>, %size: index) -> !stream.resource<external> {
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_ops.mlir b/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_ops.mlir
index 2bc0ae5..79fa894 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_ops.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_ops.mlir
@@ -182,6 +182,35 @@
// -----
+util.global private @device_a : !hal.device
+util.global private @device_b : !hal.device
+
+// CHECK-LABEL: @asyncTransferExecutionAffinity
+util.func private @asyncTransferExecutionAffinity(%arg0: !stream.resource<transient>, %arg1: index) -> !stream.resource<external> {
+ // CHECK: = stream.async.transfer on(#hal.device.affinity<@device_a>)
+ // CHECK-SAME: from(#hal.device.affinity<@device_b>)
+ // CHECK-SAME: to(#hal.device.affinity<@device_a>)
+ %0 = stream.async.transfer on(#hal.device.affinity<@device_a>) %arg0 : !stream.resource<transient>{%arg1} from(#hal.device.affinity<@device_b>) -> to(#hal.device.affinity<@device_a>) !stream.resource<external>{%arg1}
+ util.return %0 : !stream.resource<external>
+}
+
+// -----
+
+util.global private @device_a : !hal.device
+util.global private @device_b : !hal.device
+util.global private @device_c : !hal.device
+
+// CHECK-LABEL: @asyncTransferThirdDeviceExecutionAffinity
+util.func private @asyncTransferThirdDeviceExecutionAffinity(%arg0: !stream.resource<transient>, %arg1: index) -> !stream.resource<external> {
+ // CHECK: = stream.async.transfer on(#hal.device.affinity<@device_c>)
+ // CHECK-SAME: from(#hal.device.affinity<@device_a>)
+ // CHECK-SAME: to(#hal.device.affinity<@device_b>)
+ %0 = stream.async.transfer on(#hal.device.affinity<@device_c>) %arg0 : !stream.resource<transient>{%arg1} from(#hal.device.affinity<@device_a>) -> to(#hal.device.affinity<@device_b>) !stream.resource<external>{%arg1}
+ util.return %0 : !stream.resource<external>
+}
+
+// -----
+
// CHECK-LABEL: @asyncCast
util.func private @asyncCast(%arg0: !stream.resource<external>, %arg1: index) -> !stream.resource<*> {
// CHECK: = stream.async.cast %arg0 : !stream.resource<external>{%arg1} -> !stream.resource<*>{%arg1}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
index 5d17e31..e831e05 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
@@ -44,6 +44,7 @@
"PackDispatchOperands.cpp",
"Passes.cpp",
"Passes.h.inc",
+ "PlaceTransferExecutions.cpp",
"PropagateTimepoints.cpp",
"RefineUsage.cpp",
"ReuseAllocations.cpp",
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
index 406cac6..48abee9 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
@@ -41,6 +41,7 @@
"PackDispatchOperands.cpp"
"Passes.cpp"
"Passes.h.inc"
+ "PlaceTransferExecutions.cpp"
"PropagateTimepoints.cpp"
"RefineUsage.cpp"
"ReuseAllocations.cpp"
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
index 6174334..3a9c196 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
@@ -288,6 +288,8 @@
// Stream formation and scheduling
//----------------------------------------------------------------------------
+ passManager.addPass(IREE::Stream::createPlaceTransferExecutionsPass());
+
FunctionLikeNest(passManager)
// Combine async work into execution regions.
.addPass(IREE::Stream::createScheduleExecutionPass)
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
index a93b8d3..1e15b90 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
@@ -299,6 +299,19 @@
];
}
+def PlaceTransferExecutionsPass :
+ Pass<"iree-stream-place-transfer-executions", "mlir::ModuleOp"> {
+ let summary = "Materializes execution affinities for pinned transfers.";
+ let description = [{
+ Uses whole-program affinity analysis to place transfers producing pinned
+ resources on the device that owns the result when producer-side execution
+ would require the pinned allocation to be accessible from another device.
+ }];
+ let dependentDialects = [
+ "IREE::Stream::StreamDialect",
+ ];
+}
+
//===----------------------------------------------------------------------===//
// Stream formation and scheduling
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/PlaceTransferExecutions.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/PlaceTransferExecutions.cpp
new file mode 100644
index 0000000..c70b5f8
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/PlaceTransferExecutions.cpp
@@ -0,0 +1,88 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/Stream/Analysis/Affinity.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+
+#define DEBUG_TYPE "iree-stream-place-transfer-executions"
+
+namespace mlir::iree_compiler::IREE::Stream {
+
+#define GEN_PASS_DEF_PLACETRANSFEREXECUTIONSPASS
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h.inc"
+
+namespace {
+
+struct PlaceTransferExecutionsPass
+ : IREE::Stream::impl::PlaceTransferExecutionsPassBase<
+ PlaceTransferExecutionsPass> {
+ void runOnOperation() override {
+ AffinityAnalysis affinityAnalysis(getOperation());
+ if (failed(affinityAnalysis.run())) {
+ LLVM_DEBUG(llvm::dbgs() << "affinity analysis failed\n");
+ return signalPassFailure();
+ }
+
+ getOperation().walk([&](IREE::Stream::AsyncTransferOp transferOp) {
+ if (transferOp.getExecutionAffinityAttr()) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "skipping transfer with explicit execution affinity: "
+ << transferOp << "\n");
+ return;
+ }
+
+ SmallVector<IREE::Stream::AffinityAttr> pinnedAffinities;
+ if (!affinityAnalysis.tryLookupPinnedAffinities(transferOp.getResult(),
+ pinnedAffinities)) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "skipping transfer with unknown pinned affinities: "
+ << transferOp << "\n");
+ return;
+ }
+ if (pinnedAffinities.empty() || !llvm::all_equal(pinnedAffinities)) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "skipping transfer without a unique pinned affinity: "
+ << transferOp << "\n");
+ return;
+ }
+
+ auto pinnedAffinityAttr = pinnedAffinities.front();
+ auto targetAffinityAttr = transferOp.getTargetAffinityAttr();
+ if (!targetAffinityAttr ||
+ !IREE::Stream::AffinityAttr::canExecuteTogether(pinnedAffinityAttr,
+ targetAffinityAttr)) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "skipping transfer whose pinned affinity is incompatible "
+ "with its target: "
+ << transferOp << "\n");
+ return;
+ }
+
+ auto derivedExecutionAffinityAttr =
+ transferOp.getDefaultExecutionAffinityAttr();
+ if (IREE::Stream::AffinityAttr::canExecuteTogether(
+ derivedExecutionAffinityAttr, pinnedAffinityAttr)) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "skipping transfer already executing on its pinned "
+ "affinity: "
+ << transferOp << "\n");
+ return;
+ }
+
+ LLVM_DEBUG(llvm::dbgs() << "placing transfer on " << pinnedAffinityAttr
+ << ": " << transferOp << "\n");
+ transferOp.setExecutionAffinityAttr(pinnedAffinityAttr);
+ });
+ }
+};
+
+} // namespace
+} // namespace mlir::iree_compiler::IREE::Stream
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp
index 26c1273..60e9827 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp
@@ -168,6 +168,9 @@
// Note that some ops may have affinities that are more specific and we
// want to preserve those as long as possible.
if (auto transferOp = dyn_cast<IREE::Stream::AsyncTransferOp>(clonedOp)) {
+ if (transferOp.getExecutionAffinityAttr() == partition->affinity) {
+ transferOp.removeExecutionAffinityAttr();
+ }
if (transferOp.getSourceAffinityAttr() == partition->affinity) {
transferOp.setSourceAffinityAttr(nullptr);
}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
index 15b617e..6035df8 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
@@ -58,6 +58,7 @@
"pack_constants.mlir",
"pack_constants_parameters.mlir",
"pack_dispatch_operands.mlir",
+ "place_transfer_executions.mlir",
"propagate_subviews.mlir",
"propagate_timepoints.mlir",
"propagate_timepoints_scf.mlir",
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
index e29a222..9db4899 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
@@ -55,6 +55,7 @@
"pack_constants.mlir"
"pack_constants_parameters.mlir"
"pack_dispatch_operands.mlir"
+ "place_transfer_executions.mlir"
"propagate_subviews.mlir"
"propagate_timepoints.mlir"
"propagate_timepoints_scf.mlir"
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/e2e/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/e2e/BUILD.bazel
index da1f995..83ca371 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/e2e/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/e2e/BUILD.bazel
@@ -24,6 +24,7 @@
[
"async_copies.mlir",
"async_parameters.mlir",
+ "transfer_execution_placement.mlir",
],
include = ["*.mlir"],
),
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/e2e/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/e2e/CMakeLists.txt
index bace965..9d116b2 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/e2e/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/e2e/CMakeLists.txt
@@ -16,6 +16,7 @@
SRCS
"async_copies.mlir"
"async_parameters.mlir"
+ "transfer_execution_placement.mlir"
TOOLS
FileCheck
iree-opt
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/e2e/transfer_execution_placement.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/e2e/transfer_execution_placement.mlir
new file mode 100644
index 0000000..589a3c7
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/e2e/transfer_execution_placement.mlir
@@ -0,0 +1,47 @@
+// RUN: iree-opt --iree-stream-transformation-pipeline %s | FileCheck %s
+
+// End-to-end regression for https://github.com/iree-org/iree/issues/21361.
+// The final result is pinned to device A but produced by a dispatch on device
+// B. The transfer into the external result must execute on A so that the final
+// allocation remains pinned to A. Only the transient feeding that transfer
+// needs to be accessible from both devices.
+
+module {
+ util.global private @device_a : !hal.device
+ util.global private @device_b : !hal.device
+
+ // CHECK-LABEL: util.func public @multi_device_mul
+ util.func public @multi_device_mul(%input: !hal.buffer_view) -> !hal.buffer_view {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+
+ %import = stream.tensor.import on(#hal.device.affinity<@device_a>) %input : !hal.buffer_view -> tensor<4xf32> in !stream.resource<external>{%c16}
+ %produced_a = stream.async.dispatch on(#hal.device.affinity<@device_a>) @executable::@dispatch_a(%import[%c0 to %c16 for %c16]) : (!stream.resource<external>{%c16}) -> !stream.resource<transient>{%c16}
+ %transferred_b = stream.async.transfer %produced_a : !stream.resource<transient>{%c16} from(#hal.device.affinity<@device_a>) -> to(#hal.device.affinity<@device_b>) !stream.resource<transient>{%c16}
+
+ // The dispatch result produced on B and read by the transfer on A must be
+ // accessible from both devices.
+ // CHECK: %[[B_RESULT:.+]], %[[B_RESULT_ALLOCATED:.+]] = stream.resource.alloca uninitialized on(#hal.device.optimal<[#hal.device.affinity<@device_a>, #hal.device.affinity<@device_b>]>) await(%[[A_DONE:.+]])
+ // CHECK: %[[B_DONE:.+]] = stream.cmd.execute on(#hal.device.affinity<@device_b>)
+ // CHECK-SAME: %[[B_RESULT]] as %[[B_RESULT_CAPTURE:[a-zA-Z0-9_]+]]: !stream.resource<transient>
+ // CHECK: stream.cmd.dispatch @executable::@dispatch_b
+ // CHECK: wo %[[B_RESULT_CAPTURE]]
+ %produced_b = stream.async.dispatch on(#hal.device.affinity<@device_b>) @executable::@dispatch_b(%transferred_b[%c0 to %c16 for %c16]) : (!stream.resource<transient>{%c16}) -> !stream.resource<transient>{%c16}
+
+ // The final external allocation remains pinned to A. The transfer is a
+ // separate execution on A that makes the B-written intermediate visible
+ // and copies it into the pinned allocation.
+ // CHECK: %[[FINAL:.+]], %[[FINAL_ALLOCATED:.+]] = stream.resource.alloca uninitialized on(#hal.device.affinity<@device_a>){{.*}} => !stream.resource<external>
+ // CHECK: %[[TRANSFER_DONE:.+]] = stream.cmd.execute on(#hal.device.affinity<@device_a>)
+ // CHECK-SAME: %[[B_RESULT]] as %[[TRANSFER_SOURCE:[a-zA-Z0-9_]+]]: !stream.resource<transient>
+ // CHECK-SAME: %[[FINAL]] as %[[TRANSFER_TARGET:[a-zA-Z0-9_]+]]: !stream.resource<external>
+ // CHECK: stream.cmd.invalidate from(#hal.device.affinity<@device_b>) %[[TRANSFER_SOURCE]]
+ // CHECK-NEXT: stream.cmd.copy %[[TRANSFER_SOURCE]]{{.*}}, %[[TRANSFER_TARGET]]
+ %result = stream.async.transfer %produced_b : !stream.resource<transient>{%c16} from(#hal.device.affinity<@device_b>) -> to(#hal.device.affinity<@device_a>) !stream.resource<external>{%c16}
+
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await {{.*}} => %[[FINAL]]
+ // CHECK: stream.tensor.export on(#hal.device.affinity<@device_a>) %[[AWAITED]]
+ %export = stream.tensor.export on(#hal.device.affinity<@device_a>) %result : tensor<4xf32> in !stream.resource<external>{%c16} -> !hal.buffer_view
+ util.return %export : !hal.buffer_view
+ }
+}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/place_transfer_executions.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/place_transfer_executions.mlir
new file mode 100644
index 0000000..cf9f139
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/place_transfer_executions.mlir
@@ -0,0 +1,46 @@
+// RUN: iree-opt --split-input-file --iree-stream-place-transfer-executions %s | FileCheck %s
+
+module {
+ util.global private @device_a : !hal.device
+ util.global private @device_b : !hal.device
+
+ // CHECK-LABEL: util.func public @pinned_transfer
+ // CHECK: stream.async.transfer on(#hal.device.affinity<@device_a>)
+ // CHECK-SAME: from(#hal.device.affinity<@device_b>)
+ // CHECK-SAME: to(#hal.device.affinity<@device_a>)
+ util.func @pinned_transfer(%source: !stream.resource<transient>, %size: index) -> !hal.buffer_view {
+ %result = stream.async.transfer %source : !stream.resource<transient>{%size} from(#hal.device.affinity<@device_b>) -> to(#hal.device.affinity<@device_a>) !stream.resource<external>{%size}
+ %export = stream.tensor.export on(#hal.device.affinity<@device_a>) %result : tensor<?xi8>{%size} in !stream.resource<external>{%size} -> !hal.buffer_view
+ util.return %export : !hal.buffer_view
+ }
+}
+
+// -----
+
+module {
+ util.global private @device_a : !hal.device
+ util.global private @device_b : !hal.device
+
+ // CHECK-LABEL: util.func public @unpinned_transfer
+ // CHECK-NOT: stream.async.transfer on(
+ // CHECK: stream.async.transfer
+ util.func @unpinned_transfer(%source: !stream.resource<transient>, %size: index) -> !stream.resource<external> {
+ %result = stream.async.transfer %source : !stream.resource<transient>{%size} from(#hal.device.affinity<@device_b>) -> to(#hal.device.affinity<@device_a>) !stream.resource<external>{%size}
+ util.return %result : !stream.resource<external>
+ }
+}
+
+// -----
+
+module {
+ util.global private @device_a : !hal.device
+
+ // CHECK-LABEL: util.func public @compatible_transfer
+ // CHECK-NOT: stream.async.transfer on(
+ // CHECK: stream.async.transfer
+ util.func @compatible_transfer(%source: !stream.resource<transient>, %size: index) -> !hal.buffer_view {
+ %result = stream.async.transfer %source : !stream.resource<transient>{%size} from(#hal.device.affinity<@device_a>) -> to(#hal.device.affinity<@device_a>) !stream.resource<external>{%size}
+ %export = stream.tensor.export on(#hal.device.affinity<@device_a>) %result : tensor<?xi8>{%size} in !stream.resource<external>{%size} -> !hal.buffer_view
+ util.return %export : !hal.buffer_view
+ }
+}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution.mlir
index 2a94c34..79a9175 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution.mlir
@@ -628,3 +628,28 @@
// CHECK-NEXT: util.return %[[READY]]
util.return %dispatch : !stream.resource<transient>
}
+
+// -----
+
+// Tests that an explicitly placed transfer is scheduled on its execution
+// affinity instead of being bundled with its producer.
+
+// CHECK-LABEL: @placedTransfer
+util.func public @placedTransfer(%arg0: !stream.resource<external>, %size: index) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+
+ // CHECK: %[[PRODUCED:.+]], %[[PRODUCED_TIMEPOINT:.+]] = stream.async.execute on(#hal.device.affinity<@device_b>)
+ // CHECK: %[[DISPATCH:.+]] = stream.async.dispatch
+ // CHECK: stream.yield %[[DISPATCH]]
+ %produced = stream.async.dispatch on(#hal.device.affinity<@device_b>) @ex::@dispatch_b[%c1](%arg0[%c0 to %size for %size]) : (!stream.resource<external>{%size}) -> !stream.resource<transient>{%size}
+
+ // CHECK: %[[TRANSFERRED:.+]], %[[TRANSFER_TIMEPOINT:.+]] = stream.async.execute on(#hal.device.affinity<@device_a>)
+ // CHECK-SAME: await(%[[PRODUCED_TIMEPOINT]])
+ // CHECK: %[[TRANSFER:.+]] = stream.async.transfer
+ // CHECK-NEXT: stream.yield %[[TRANSFER]]
+ %transferred = stream.async.transfer on(#hal.device.affinity<@device_a>) %produced : !stream.resource<transient>{%size} from(#hal.device.affinity<@device_b>) -> to(#hal.device.affinity<@device_a>) !stream.resource<external>{%size}
+
+ // CHECK: stream.timepoint.await %[[TRANSFER_TIMEPOINT]] => %[[TRANSFERRED]]
+ util.return %transferred : !stream.resource<external>
+}
