ScheduleExecution enhancements for timeline-aware scheduling and SCF. (#22483)
Paritioning now handles values used in SCF (and other) region ops
properly and ScheduleExecution now handles timeline-aware ops (those ops
that operate on a timeline and can materialize timepoints even if they
don't naturally have them).
A test op is added to avoid relying on the HAL ops that implement the
interface (I'll probably try to move more tests to test ops/attrs in the
future to keep HAL out of Stream).
Part of #16168 PR sequence (4/6).
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Analysis/BUILD.bazel
index 3cbb5b5..c26b567 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/BUILD.bazel
@@ -35,6 +35,7 @@
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Analysis",
"@llvm-project//mlir:ArithDialect",
+ "@llvm-project//mlir:ControlFlowInterfaces",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FunctionInterfaces",
"@llvm-project//mlir:IR",
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Analysis/CMakeLists.txt
index c2dd74c..b229ec2 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/CMakeLists.txt
@@ -28,6 +28,7 @@
LLVMSupport
MLIRAnalysis
MLIRArithDialect
+ MLIRControlFlowInterfaces
MLIRFuncDialect
MLIRFunctionInterfaces
MLIRIR
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.cpp
index 52a2e66..0bfcf5c 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.cpp
@@ -10,11 +10,56 @@
#include "llvm/Support/Debug.h"
#include "mlir/IR/AsmState.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
#define DEBUG_TYPE "iree-stream-partitioning"
namespace mlir::iree_compiler::IREE::Stream {
+//===----------------------------------------------------------------------===//
+// Utilities
+//===----------------------------------------------------------------------===//
+
+// Collects all values consumed by an operation, including those used in nested
+// regions (e.g. scf.for bodies). This is a conservative analysis that walks all
+// regions to find consumed values.
+void collectConsumedValues(Operation *rootOp,
+ SetVector<Value> &consumedValues) {
+ SmallVector<Operation *> worklist;
+ DenseSet<Operation *> visitedOps;
+ worklist.push_back(rootOp);
+ while (!worklist.empty()) {
+ Operation *op = worklist.pop_back_val();
+
+ // Skip if already visited.
+ if (!visitedOps.insert(op).second) {
+ continue;
+ }
+
+ // Collect direct operands.
+ for (auto operand : op->getOperands()) {
+ consumedValues.insert(operand);
+ }
+
+ // Conservatively walk all regions to find consumed values.
+ // We can't use RegionBranchOpInterface::getEntrySuccessorRegions because
+ // it requires compile-time constant operands to determine reachability,
+ // but we need to handle runtime values. For hazard detection we need to be
+ // conservative and assume all regions may execute.
+ for (auto ®ion : op->getRegions()) {
+ for (auto &block : region) {
+ for (auto &nestedOp : block) {
+ worklist.push_back(&nestedOp);
+ }
+ }
+ }
+ }
+}
+
+//===----------------------------------------------------------------------===//
+// Partition data structures
+//===----------------------------------------------------------------------===//
+
#ifndef NDEBUG
void dumpPartition(Partition &partition, AsmState &asmState) {
@@ -93,6 +138,40 @@
}
}
+ // Check for circular dependencies: input ops using partition outputs.
+ // This can happen if an operation with nested regions (like scf.for) is
+ // incorrectly placed as a partition input when its nested regions actually
+ // consume partition outputs.
+ for (auto in : ins) {
+ // Only check ops, not bare values.
+ auto definingOp = in.getDefiningOp();
+ if (!definingOp)
+ continue;
+
+ // Collect all values used by this input op (including nested regions).
+ SetVector<Value> inputConsumedValues;
+ collectConsumedValues(definingOp, inputConsumedValues);
+
+ // Check if any consumed value comes from partition outputs.
+ //
+ // NOTE: We only check outputs, not all defValues. When operations are
+ // duplicated across partitions (e.g., clones with preferCloneToConsumers),
+ // the same operation may define values in multiple partitions. An input
+ // operation might consume a value that is defined in THIS partition's
+ // defValues, but if that value is not exported (not in outputs), the input
+ // operation must be using a copy from another partition that also defines
+ // it. Checking defValues would create false positives for such cases.
+ for (auto consumedValue : inputConsumedValues) {
+ if (outs.contains(consumedValue)) {
+ return mlir::emitError(loc)
+ << "circular dependency: input operation uses partition output "
+ << consumedValue
+ << " - this indicates the operation should be inside the "
+ "partition";
+ }
+ }
+ }
+
return success();
}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.h b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.h
index dbf2ba7..923a5a4 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.h
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.h
@@ -60,6 +60,15 @@
};
//===----------------------------------------------------------------------===//
+// Utilities
+//===----------------------------------------------------------------------===//
+
+// Collects all values consumed by an operation, including those used in nested
+// regions (e.g. scf.for bodies). Uses a worklist-based approach to avoid stack
+// overflow on deeply nested regions.
+void collectConsumedValues(Operation *rootOp, SetVector<Value> &consumedValues);
+
+//===----------------------------------------------------------------------===//
// Stream partitioning algorithms
//===----------------------------------------------------------------------===//
//
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 e4baa0c..0ae3d0d 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp
@@ -42,6 +42,10 @@
llvm::BitVector membership;
// Which partitions transitively depend on this operation.
llvm::BitVector hazards;
+ // Hazards specifically from nested region captures creating circular
+ // dependencies. These must always be respected, even for operations with
+ // preferCloneToConsumers.
+ llvm::BitVector nestedRegionHazards;
};
struct PartitionBuilder {
@@ -64,9 +68,78 @@
opInfo.hazards.reset(ordinal);
ops.insert(op);
hazards |= opInfo.hazards;
+ hazards |= opInfo.nestedRegionHazards;
}
};
+// Find the effective OpInfo for a user operation.
+// Returns the OpInfo* to use for hazard tracking, or nullptr if the user
+// should be skipped (either because hazards were already set or not found).
+static OpInfo *getEffectiveUserInfo(Operation *user, Block *block,
+ DenseMap<Operation *, OpInfo> &opInfos,
+ OpInfo &opInfo) {
+ // Check if user is in the same block we're partitioning.
+ if (user->getBlock() == block) {
+ // Direct lookup - user in same block.
+ LLVM_DEBUG(llvm::dbgs() << " User in same block\n");
+ auto userInfoIt = opInfos.find(user);
+ if (userInfoIt != opInfos.end()) {
+ return &userInfoIt->second;
+ }
+ return nullptr;
+ }
+
+ LLVM_DEBUG(llvm::dbgs() << " User in different block (nested region)\n");
+
+ // User is in a nested region - find the containing operation in our block.
+ // This handles cases like scf.for where operations inside the loop body
+ // use values from the parent block. We treat the containing operation
+ // (scf.for) as a proxy for all its nested users.
+ Operation *containing = user->getParentOp();
+ while (containing && containing->getBlock() != block) {
+ containing = containing->getParentOp();
+ }
+
+ if (!containing) {
+ return nullptr;
+ }
+
+ // Use the containing op's info as proxy for nested users.
+ auto containingInfoIt = opInfos.find(containing);
+ if (containingInfoIt != opInfos.end()) {
+ // Only use containing op's info if it's streamable.
+ // Non-streamable ops (scf.for, scf.if) need special circular dependency
+ // tracking via nestedRegionHazards to prevent grouping captured values
+ // with operations that consume the containing op's results.
+ if (isa<IREE::Stream::StreamableOpInterface>(containing)) {
+ return &containingInfoIt->second;
+ }
+ // Fall through to set nestedRegionHazards for non-streamable containing
+ // ops.
+ }
+
+ // The containing operation is not streamable (e.g., scf.for).
+ // Values captured into it create circular dependency hazards with any
+ // partitions that use the containing operation's results.
+ for (auto result : containing->getResults()) {
+ for (auto resultUser : result.getUsers()) {
+ if (resultUser->getBlock() == block) {
+ auto resultUserInfoIt = opInfos.find(resultUser);
+ if (resultUserInfoIt != opInfos.end()) {
+ opInfo.nestedRegionHazards |= resultUserInfoIt->second.membership;
+ opInfo.nestedRegionHazards |= resultUserInfoIt->second.hazards;
+ LLVM_DEBUG({
+ llvm::dbgs() << "Setting nestedRegionHazards for captured value, "
+ << "hazard count: "
+ << opInfo.nestedRegionHazards.count() << "\n";
+ });
+ }
+ }
+ }
+ }
+ return nullptr;
+}
+
// This is terrible. See Stream/Analysis/Partition.h for a description of what
// a real implementation would do. We want cost modeling for tie breakers when
// an op could be in multiple partitions, cloning for ops that are not worth
@@ -131,8 +204,31 @@
opHazardsInCandidatePartition |= opInfos[user].hazards;
}
}
+ // For preferCloneToConsumers ops, we DON'T include nestedRegionHazards
+ // because cloning into multiple partitions solves the circular
+ // dependency. The op will be duplicated in each consumer partition, so
+ // there's no shared state causing circular dependencies.
} else {
- opHazards = &opInfo.hazards;
+ // For non-clonable ops, nestedRegionHazards must be respected to prevent
+ // circular dependencies from nested region captures. Unlike normal
+ // hazards which propagate through use-def chains, nestedRegionHazards may
+ // not be fully captured in opInfo.hazards when the containing operation
+ // (e.g., scf.for) is not streamable and returns nullptr from
+ // getEffectiveUserInfo.
+ opHazardsInCandidatePartition = opInfo.hazards;
+ opHazardsInCandidatePartition |= opInfo.nestedRegionHazards;
+ opHazards = &opHazardsInCandidatePartition;
+ }
+
+ // Check nestedRegionHazards separately first. Unlike normal hazards
+ // (which indicate dependencies), nestedRegionHazards indicate partitions
+ // the op CANNOT join due to circular dependency constraints.
+ for (auto nestedHazardOrdinal : opInfo.nestedRegionHazards.set_bits()) {
+ if (nestedHazardOrdinal == partitionOrdinal) {
+ // This op cannot be in this partition due to nested region capture
+ // circular dependency constraints.
+ return false;
+ }
}
for (auto opHazardOrdinal : opHazards->set_bits()) {
@@ -155,8 +251,36 @@
auto asmState = getRootAsmState(block);
- llvm::DenseMap<Operation *, llvm::SmallVector<Operation *>> syncOps;
+ // Build a map from values to operations that consume them in nested regions.
+ // This allows us to detect when an operation with nested regions (like
+ // scf.for) consumes values produced by other operations that may create
+ // hazards even though there's no direct use-def edge.
+ DenseMap<Value, SmallVector<Operation *>> nestedConsumers;
+ for (auto &op : *block) {
+ // Only check operations with regions that are not streamable.
+ // Streamable ops are handled normally but non-streamable ops with regions
+ // (like scf.for, scf.if, scf.while) need special handling.
+ if (op.getNumRegions() == 0 ||
+ isa<IREE::Stream::StreamableOpInterface>(op)) {
+ continue;
+ }
+ // Collect all values consumed by this operation (including in nested
+ // regions).
+ SetVector<Value> consumedValues;
+ collectConsumedValues(&op, consumedValues);
+
+ // For each consumed value record this op as a nested consumer.
+ for (auto value : consumedValues) {
+ // Only track values defined in the same block (not block arguments).
+ auto definingOp = value.getDefiningOp();
+ if (definingOp && definingOp->getBlock() == block) {
+ nestedConsumers[value].push_back(&op);
+ }
+ }
+ }
+
+ llvm::DenseMap<Operation *, llvm::SmallVector<Operation *>> syncOps;
for (auto &op : llvm::reverse(*block)) {
LLVM_DEBUG({
@@ -223,12 +347,16 @@
// We prune the set based on whether the users are part of a transitive
// dependency chain down the use-def chain to a partition.
llvm::BitVector consumers(builders.size(), /*t=*/false);
+
+ // Process direct users (common case).
for (auto user : op.getUsers()) {
- auto userInfoIt = opInfos.find(user);
- if (userInfoIt == opInfos.end()) {
+ OpInfo *effectiveUserInfo =
+ getEffectiveUserInfo(user, block, opInfos, opInfo);
+ if (!effectiveUserInfo) {
continue;
}
- auto &userInfo = userInfoIt->second;
+
+ auto &userInfo = *effectiveUserInfo;
LLVM_DEBUG({
llvm::dbgs() << "Testing user:\n";
user->print(llvm::dbgs(), *asmState);
@@ -245,14 +373,47 @@
opInfo.hazards |= userInfo.hazards;
}
+ // Process indirect users (operations that consume this op's results in
+ // their nested regions e.g. scf.for using a value in its body).
+ for (auto result : op.getResults()) {
+ auto nestedConsumersIt = nestedConsumers.find(result);
+ if (nestedConsumersIt != nestedConsumers.end()) {
+ for (auto *nestedUser : nestedConsumersIt->second) {
+ auto userInfoIt = opInfos.find(nestedUser);
+ if (userInfoIt == opInfos.end()) {
+ continue;
+ }
+ auto &userInfo = userInfoIt->second;
+ LLVM_DEBUG({
+ llvm::dbgs() << "Testing nested region user:\n";
+ nestedUser->print(llvm::dbgs(), *asmState);
+ llvm::dbgs() << "\n";
+ for (auto membershipOrdinal : userInfo.membership.set_bits()) {
+ llvm::dbgs() << " member of partition " << membershipOrdinal
+ << "\n";
+ }
+ for (auto hazardOrdinal : userInfo.hazards.set_bits()) {
+ llvm::dbgs() << " hazard w/ partition " << hazardOrdinal << "\n";
+ }
+ });
+ // For nested consumers we only propagate hazards and not membership.
+ // The nested consumer itself is not streamable and won't be in a
+ // partition but it creates a dependency barrier.
+ opInfo.hazards |= userInfo.membership;
+ opInfo.hazards |= userInfo.hazards;
+ }
+ }
+ }
+
for (auto syncOp : syncOps[&op]) {
for (auto user : syncOp->getUsers()) {
- auto userInfoIt = opInfos.find(user);
- if (userInfoIt == opInfos.end()) {
+ OpInfo *effectiveUserInfo =
+ getEffectiveUserInfo(user, block, opInfos, opInfo);
+ if (!effectiveUserInfo) {
continue;
}
- auto &userInfo = userInfoIt->second;
+ auto &userInfo = *effectiveUserInfo;
LLVM_DEBUG({
llvm::dbgs() << "Testing sync user:\n";
user->print(llvm::dbgs(), *asmState);
@@ -382,8 +543,13 @@
LLVM_DEBUG(llvm::dbgs() << "Moving to first candidate partition "
<< firstCandidateOrdinal << " (continue)\n");
// If we are a clonable op (like splat) clone us into every partition.
+ // We also clone if there are nested region hazards, indicating consumers
+ // in nested regions that will need their own copy.
// Otherwise we just pick the first we find (probably a bad heuristic).
- if (consumers.count() > 1 && streamableOp.preferCloneToConsumers()) {
+ bool shouldClone =
+ streamableOp.preferCloneToConsumers() &&
+ (consumers.count() > 1 || opInfo.nestedRegionHazards.any());
+ if (shouldClone && consumers.any()) {
for (auto consumerOrdinal : consumers.set_bits()) {
LLVM_DEBUG(llvm::dbgs() << "Cloning into consumer partition "
<< consumerOrdinal << "\n");
@@ -391,7 +557,29 @@
consumerBuilder->insert(&op, opInfo);
consumerBuilder->clonedOps.insert(&op);
}
+ } else if (shouldClone && !consumers.any()) {
+ // We want to clone but have no streamable consumers. Check if the op
+ // has any uses in the same block - if so, it needs to be in a partition.
+ // If all uses are in nested regions, leave it for re-materialization.
+ bool hasUsesInSameBlock = false;
+ for (auto result : op.getResults()) {
+ for (auto user : result.getUsers()) {
+ if (user->getBlock() == block) {
+ hasUsesInSameBlock = true;
+ break;
+ }
+ }
+ if (hasUsesInSameBlock) {
+ break;
+ }
+ }
+ if (hasUsesInSameBlock) {
+ // Has uses in same block, insert into first candidate partition.
+ builder->insert(&op, opInfo);
+ }
+ // Else: all uses are in nested regions, leave for re-materialization.
} else {
+ // Not a clonable op, insert into first candidate partition.
builder->insert(&op, opInfo);
}
}
@@ -410,10 +598,17 @@
SetVector<Value> producedValues;
SetVector<Value> escapingValues;
for (auto *op : llvm::reverse(builder->ops)) {
+ // Collect direct operands only. Do NOT recursively collect values from
+ // nested regions of operations that define our operands - that would
+ // incorrectly add values consumed inside control flow ops (scf.for, etc)
+ // to partition inputs, creating spurious circular dependencies.
+ // The circular dependency check in Partition::verify() will correctly
+ // detect actual circular dependencies using collectConsumedValues.
bool didCloneEscape = false;
for (auto operand : op->getOperands()) {
consumedValues.insert(operand);
}
+
for (auto result : op->getResults()) {
producedValues.insert(result);
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
index fc2bcab..a27ae0b 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
@@ -4492,6 +4492,31 @@
return getResourceOperandsMutable();
}
+//===----------------------------------------------------------------------===//
+// stream.test.timeline_aware
+//===----------------------------------------------------------------------===//
+
+SmallVector<Value>
+TestTimelineAwareOp::buildAwaitTimepoints(OpBuilder &builder) {
+ // Build timepoint.import ops for each wait fence.
+ SmallVector<Value> timepoints;
+ for (auto fence : getWaitFenceLikes()) {
+ auto importOp = IREE::Stream::TimepointImportOp::create(
+ builder, getLoc(), builder.getType<IREE::Stream::TimepointType>(),
+ ValueRange{fence}, /*affinity=*/nullptr);
+ timepoints.push_back(importOp.getResultTimepoint());
+ }
+ return timepoints;
+}
+
+Value TestTimelineAwareOp::buildResultTimepoint(OpBuilder &builder) {
+ // Build timepoint.import for signal fence.
+ auto importOp = IREE::Stream::TimepointImportOp::create(
+ builder, getLoc(), builder.getType<IREE::Stream::TimepointType>(),
+ ValueRange{getSignalFenceLike()}, /*affinity=*/nullptr);
+ return importOp.getResultTimepoint();
+}
+
} // namespace mlir::iree_compiler::IREE::Stream
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
index 4c9f326..5edc9aa 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
@@ -4664,4 +4664,66 @@
} // OpGroupMiscellaneousOps
+//===----------------------------------------------------------------------===//
+// Test ops
+//===----------------------------------------------------------------------===//
+
+def OpGroupTestOps : OpDocGroup {
+ let summary = "Test-only ops for Stream dialect testing";
+ let description = "Ops used only in unit tests.";
+}
+
+let opDocGroup = OpGroupTestOps in {
+
+def Stream_TestTimelineAwareOp : Stream_PureOp<"test.timeline_aware", [
+ AttrSizedOperandSegments,
+ DeclareOpInterfaceMethods<Stream_TimelineAwareOp, [
+ "buildAwaitTimepoints",
+ "buildResultTimepoint",
+ ]>,
+]> {
+ let summary = [{Test op implementing TimelineAwareOpInterface.}];
+ let description = [{
+ Test-only operation that implements TimelineAwareOpInterface directly,
+ allowing Stream dialect tests to verify timeline-aware behavior without
+ requiring HAL/etc dialect dependencies.
+
+ Mimics util.call with the HAL coarse-fences model:
+ - Takes wait fences that must complete before execution.
+ - Takes a signal fence to signal when complete.
+ - Returns an arbitrary result.
+ }];
+
+ let arguments = (ins
+ Variadic<AnyType>:$args,
+ Variadic<Stream_TestFence>:$wait_fence_likes,
+ Stream_TestFence:$signal_fence_like
+ );
+ let results = (outs
+ Variadic<AnyType>:$results
+ );
+
+ let assemblyFormat = [{
+ `(` $args `)` `waits` `(` $wait_fence_likes `)` `signals` `(` $signal_fence_like `)`
+ attr-dict `:` functional-type($args, $results)
+ }];
+
+ let extraClassDeclaration = [{
+ bool participatesInTimeline() { return true; }
+
+ SmallVector<Value> getAwaitFences() {
+ // Return all wait fences (excluding signal fence).
+ SmallVector<Value> fences;
+ for (auto fence : getWaitFenceLikes()) {
+ fences.push_back(fence);
+ }
+ return fences;
+ }
+
+ Value getSignalFence() { return getSignalFenceLike(); }
+ }];
+}
+
+} // OpGroupTestOps
+
#endif // IREE_DIALECT_STREAM_OPS
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.td
index 5668a45..1abb9d8 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.td
@@ -49,6 +49,24 @@
let hasCustomAssemblyFormat = 1;
}
+def Stream_TestFence : Stream_TypeDef<"TestFence", []> {
+ let mnemonic = "test.fence";
+
+ let summary = [{Test-only fence type for timeline-aware op testing.}];
+ let description = [{
+ A test-only type that behaves like hal.fence for testing TimelineAwareOpInterface
+ without introducing HAL dependencies in Stream unit tests.
+ }];
+
+ let parameters = (ins);
+
+ let builders = [
+ TypeBuilder<(ins), [{
+ return $_get($_ctxt);
+ }]>,
+ ];
+}
+
//===----------------------------------------------------------------------===//
// Stream resource value types
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp
index ebe8d1b..67e9ab0 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp
@@ -5,7 +5,6 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include "iree/compiler/Dialect/Stream/Analysis/Partitioning.h"
-#include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
@@ -224,14 +223,63 @@
return orderedBlocks;
}
+// Returns true if an operation can be placed into a partition (execute region).
+// Only streamable operations should be partitioned; all other operations stay
+// in the parent block and have their nested regions recursively processed.
+//
+// TODO(benvanik): allow limited scf ops inside the regions when we can ensure
+// they rely on only constant values (in the future we can do indirect stuff).
+static bool canPlaceInPartition(Operation *op) {
+ // Only StreamableOpInterface ops can be placed into partitions.
+ // This includes ops like stream.async.slice, stream.async.clone, etc.
+ // but excludes stream.async.execute (which is the partition container itself)
+ // and control flow ops like scf.for.
+ return isa<IREE::Stream::StreamableOpInterface>(op);
+}
+
+// Tries to find a timeline-aware consumer of |value| that can provide a
+// timepoint for synchronization. This looks through stream.tensor.export ops
+// since timeline-aware ops (like util.call with fences) typically work with
+// HAL types rather than stream resources.
+// Returns the timepoint if found and otherwise returns null.
+static Value tryGetConsumerTimepoint(Value value) {
+ SmallVector<Operation *> usersToCheck;
+ usersToCheck.append(value.getUsers().begin(), value.getUsers().end());
+
+ // Look through tensor exports to find timeline-aware ops.
+ for (auto *user : value.getUsers()) {
+ if (auto exportOp = dyn_cast<IREE::Stream::TensorExportOp>(user)) {
+ usersToCheck.append(exportOp.getResult().getUsers().begin(),
+ exportOp.getResult().getUsers().end());
+ }
+ }
+
+ for (auto *user : usersToCheck) {
+ if (auto awareOp = dyn_cast<IREE::Stream::TimelineAwareOpInterface>(user)) {
+ if (awareOp.participatesInTimeline()) {
+ // Let the timeline-aware op build its result timepoint.
+ OpBuilder awareBuilder(user);
+ awareBuilder.setInsertionPointAfter(user);
+ if (auto timepoint = awareOp.buildResultTimepoint(awareBuilder)) {
+ return timepoint;
+ }
+ }
+ }
+ }
+
+ return nullptr;
+}
+
LogicalResult processRegion(Location loc, MLIRContext *context, Region ®ion,
const PartitioningConfigAttr &configAttr) {
for (auto *block : sortBlocksInDominanceOrder(region)) {
// Compute a set of partitions covering all of the streamable ops in the
// block.
auto partitionSet = partitionStreamableOps(configAttr, block);
- if (partitionSet.empty())
+ if (partitionSet.empty()) {
continue;
+ }
+
if (failed(partitionSet.verify(loc))) {
return failure();
}
@@ -275,22 +323,64 @@
for (auto [oldResult, newResult, newResultSize] : llvm::zip_equal(
partitionBuilder.partition->outs, executeOp.getResults(),
executeOp.getResultSizes())) {
- // Insert one await per result. We could batch them all but that would
- // prematurely tie their lifetimes together. By having unique awaits
- // we allow propagation to move the waits further to where the values
- // are used (including right into other execution regions).
- auto awaitOp = IREE::Stream::TimepointAwaitOp::create(
- builder, executeOp.getLoc(), newResult, newResultSize,
- executeOp.getResultTimepoint());
-
- // Explicitly copy the Value since it is marked as const.
- Value toBeDeleted = oldResult;
-
- toBeDeleted.replaceAllUsesWith(awaitOp.getResults().front());
- deadOps.insert(oldResult.getDefiningOp());
+ // Check if any consumer is timeline-aware and can provide its own
+ // timepoint - if so we reuse that to guarantee we don't ever wait.
+ Value deferredTimepoint = tryGetConsumerTimepoint(oldResult);
+ if (deferredTimepoint) {
+ // Timeline-aware op provided a timepoint - use it.
+ auto awaitOp = IREE::Stream::TimepointAwaitOp::create(
+ builder, executeOp.getLoc(), newResult, newResultSize,
+ deferredTimepoint);
+ deadOps.insert(oldResult.getDefiningOp());
+ Value toReplace = oldResult;
+ toReplace.replaceAllUsesWith(awaitOp.getResults().front());
+ } else {
+ // Normal case - immediate await on the execute's timepoint.
+ auto awaitOp = IREE::Stream::TimepointAwaitOp::create(
+ builder, executeOp.getLoc(), newResult, newResultSize,
+ executeOp.getResultTimepoint());
+ deadOps.insert(oldResult.getDefiningOp());
+ Value toReplace = oldResult;
+ toReplace.replaceAllUsesWith(awaitOp.getResults().front());
+ }
}
}
+ // Before erasing preferCloneToConsumers ops, re-materialize them in nested
+ // regions that reference them.
+ DenseMap<Operation *, Operation *> rematerializedOps;
+ for (auto *deadOp : deadOps) {
+ auto streamableOp = dyn_cast<IREE::Stream::StreamableOpInterface>(deadOp);
+ if (!streamableOp || !streamableOp.preferCloneToConsumers()) {
+ continue;
+ }
+
+ // Find all uses in nested regions and re-materialize.
+ for (auto &use : llvm::make_early_inc_range(deadOp->getUses())) {
+ Operation *user = use.getOwner();
+ // Only look at users in other regions.
+ if (user->getBlock() == deadOp->getBlock()) {
+ continue;
+ }
+
+ // Re-materialize the op at the start of the user's block.
+ Block *nestedBlock = user->getBlock();
+ auto *clonedOp = deadOp->clone();
+ nestedBlock->getOperations().insert(nestedBlock->begin(), clonedOp);
+
+ // Replace this use with the cloned op's result.
+ Value oldValue = use.get();
+ auto resultValue = dyn_cast<OpResult>(oldValue);
+ unsigned resultIndex = resultValue.getResultNumber();
+ use.set(clonedOp->getResult(resultIndex));
+ }
+ }
+
for (auto *deadOp : llvm::reverse(deadOps)) {
+ if (!deadOp->use_empty()) {
+ // Keep ops that are still used (shouldn't happen after
+ // re-materialization).
+ continue;
+ }
deadOp->erase();
}
@@ -304,9 +394,27 @@
});
}
+ // Process nested regions AFTER partitioning the current block.
+ // This allows nested regions to use the remapped values from parent
+ // partitions. We only process operations that were NOT placed into
+ // partitions above (e.g., control flow operations like scf.for).
for (auto *block : sortBlocksInDominanceOrder(region)) {
for (auto &op : *block) {
- if (isa<scf::SCFDialect>(op.getDialect())) {
+ // Skip ops that were placed into partitions (they're being moved).
+ if (canPlaceInPartition(&op)) {
+ continue;
+ }
+
+ // Process any op with regions that supports control flow, except for
+ // stream.async.execute and stream.async.concurrent which are the
+ // partition containers created by this pass.
+ if (isa<RegionBranchOpInterface>(op)) {
+ // Don't recurse into the partition containers we just created.
+ if (isa<IREE::Stream::AsyncExecuteOp>(op) ||
+ isa<IREE::Stream::AsyncConcurrentOp>(op)) {
+ continue;
+ }
+
for (auto &subregion : op.getRegions()) {
if (failed(processRegion(loc, context, subregion, configAttr)))
return failure();
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 4817ff2..a68e003 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
@@ -55,6 +55,8 @@
"schedule_allocation.mlir",
"schedule_concurrency.mlir",
"schedule_execution.mlir",
+ "schedule_execution_scf.mlir",
+ "schedule_execution_timeline_aware.mlir",
"specialize_dispatches.mlir",
"specialize_encodings.mlir",
"sync_initializers.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 658bcc6..83329a5 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
@@ -53,6 +53,8 @@
"schedule_allocation.mlir"
"schedule_concurrency.mlir"
"schedule_execution.mlir"
+ "schedule_execution_scf.mlir"
+ "schedule_execution_timeline_aware.mlir"
"specialize_dispatches.mlir"
"specialize_encodings.mlir"
"sync_initializers.mlir"
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 0220467..f2786a2 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
@@ -514,3 +514,50 @@
}
util.return %sum, %arg1 : !stream.resource<*>, index
}
+
+// -----
+
+// Tests clones that span multiple partitions separated by a load barrier do not
+// create a circular dependency when the clone is used on both sides of the
+// barrier. The clone is duplicated into both partitions.
+
+// CHECK-LABEL: @cloneAcrossLoadBarrier
+util.func public @cloneAcrossLoadBarrier(%arg0: !stream.resource<external>) -> (!stream.resource<transient>, !stream.resource<transient>) {
+ %c0 = arith.constant 0 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %c24 = arith.constant 24 : index
+
+ // First partition: clone + dispatch0 + transfer (before load barrier).
+ // CHECK: %[[RESULTS0:.+]]:2, %[[TP0:.+]] = stream.async.execute
+ // CHECK-SAME: with({{.+}} as %[[ARG0:.+]]: !stream.resource<external>{%c24})
+ // CHECK-SAME: -> (!stream.resource<transient>{%c16}, !stream.resource<staging>{%c16})
+ // CHECK-NEXT: %[[CLONE0:.+]] = stream.async.clone %[[ARG0]]
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c24} -> !stream.resource<external>{%c24}
+ // CHECK-NEXT: %[[DISPATCH0:.+]] = stream.async.dispatch @ex::@dispatch_0(%[[CLONE0]]
+ %dispatch0 = stream.async.dispatch @ex::@dispatch_0(%clone[%c0 to %c24 for %c24]) : (!stream.resource<external>{%c24}) -> !stream.resource<transient>{%c16}
+ // CHECK-NEXT: %[[TRANSFER:.+]] = stream.async.transfer %[[DISPATCH0]]
+ %transfer = stream.async.transfer %dispatch0 : !stream.resource<transient>{%c16} -> !stream.resource<staging>{%c16}
+ // CHECK-NEXT: stream.yield %[[DISPATCH0]], %[[TRANSFER]]
+
+ // CHECK: stream.timepoint.await %[[TP0]] => %[[RESULTS0]]#1
+ // CHECK: %[[LOADED:.+]] = stream.async.load
+ %loaded = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c16} -> i64
+ // Second partition: clone (duplicated) + splat + dispatch1 + dispatch2 (after load barrier).
+ // CHECK: %[[RESULTS1:.+]]:2, %[[TP1:.+]] = stream.async.execute
+ // CHECK-SAME: await(%[[TP0]])
+ // CHECK-SAME: with({{.+}} as %[[ARG1:.+]]: !stream.resource<external>{%c24},
+ // CHECK-SAME: %[[RESULTS0]]#0 as %[[DISPATCH0_IN:.+]]: !stream.resource<transient>{%c16})
+ // CHECK-SAME: -> (!stream.resource<transient>{%c8}, !stream.resource<transient>{%c8})
+ // CHECK-NEXT: %[[CLONE1:.+]] = stream.async.clone %[[ARG1]]
+ // CHECK-NEXT: %[[SPLAT:.+]] = stream.async.splat %[[LOADED]]
+ %filled = stream.async.splat %loaded : i64 -> !stream.resource<transient>{%c24}
+ // CHECK-NEXT: %[[DISPATCH1:.+]] = stream.async.dispatch @ex::@dispatch_1(%[[DISPATCH0_IN]]{{.+}}, %[[SPLAT]]
+ %dispatch1 = stream.async.dispatch @ex::@dispatch_1(%dispatch0[%c0 to %c16 for %c16], %filled[%c0 to %c24 for %c24]) : (!stream.resource<transient>{%c16}, !stream.resource<transient>{%c24}) -> !stream.resource<transient>{%c8}
+ // CHECK-NEXT: %[[DISPATCH2:.+]] = stream.async.dispatch @ex::@dispatch_2(%[[CLONE1]]
+ %dispatch2 = stream.async.dispatch @ex::@dispatch_2(%clone[%c0 to %c24 for %c24]) : (!stream.resource<external>{%c24}) -> !stream.resource<transient>{%c8}
+ // CHECK-NEXT: stream.yield %[[DISPATCH1]], %[[DISPATCH2]]
+
+ // CHECK: stream.timepoint.await %[[TP1]] => %[[RESULTS1]]
+ util.return %dispatch1, %dispatch2 : !stream.resource<transient>, !stream.resource<transient>
+}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution_scf.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution_scf.mlir
new file mode 100644
index 0000000..5a4dbf8
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution_scf.mlir
@@ -0,0 +1,1147 @@
+// RUN: iree-opt --split-input-file --allow-unregistered-dialect --pass-pipeline="builtin.module(util.func(iree-stream-schedule-execution))" %s | FileCheck %s
+
+// Tests basic scf.for loop with dependency on cloned resource (reproducer case).
+// The loop uses a cloned resource from a partition, which creates a dependency
+// that must be respected during partitioning.
+
+// CHECK-LABEL: @scfForWithDependency
+util.func public @scfForWithDependency(%arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c11 = arith.constant 11 : index
+ %c48 = arith.constant 48 : index
+ %c0_i64 = arith.constant 0 : i64
+ %c5_i64 = arith.constant 5 : i64
+
+ // CHECK: scf.for
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c48} -> !stream.resource<external>{%c48}
+
+
+ %result = scf.for %i = %c0 to %c11 step %c1 iter_args(%iter = %c0_i64) -> (i64) {
+ %offset = arith.muli %i, %c4 : index
+ %end = arith.addi %offset, %c4 : index
+ // Loop uses the cloned resource - this should NOT be partitioned into the clone partition.
+ // CHECK: %[[EXEC_RESULT:.+]], %[[EXEC_TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: with(%{{.+}} as %[[CAPTURE:.+]]: !stream.resource<external>{%c48})
+ // CHECK-SAME: -> !stream.resource<staging>{%c4}
+ // CHECK-NEXT: %[[CLONE:.+]] = stream.async.clone %[[CAPTURE]]
+ // CHECK-NEXT: %[[SLICE:.+]] = stream.async.slice %[[CLONE]]
+ // CHECK-NEXT: %[[TRANSFER:.+]] = stream.async.transfer %[[SLICE]]
+ // CHECK-NEXT: stream.yield %[[TRANSFER]]
+ %slice = stream.async.slice %clone[%offset to %end] : !stream.resource<external>{%c48} -> !stream.resource<external>{%c4}
+ %transfer = stream.async.transfer %slice : !stream.resource<external>{%c4} -> !stream.resource<staging>{%c4}
+ // CHECK: stream.timepoint.await %[[EXEC_TIMEPOINT]]
+ %load = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c4} -> i32
+ %ext = arith.extsi %load : i32 to i64
+ %cmp = arith.cmpi eq, %ext, %c5_i64 : i64
+ %add = arith.extui %cmp : i1 to i64
+ %next = arith.addi %iter, %add : i64
+ scf.yield %next : i64
+ }
+
+ // CHECK: stream.async.execute
+ // CHECK-NEXT: stream.async.splat
+ %splat = stream.async.splat %result : i64 -> !stream.resource<external>{%c4}
+ util.return %splat : !stream.resource<external>
+}
+
+// -----
+
+// Tests scf.for loop result being used by a dispatch.
+// The loop produces a value that is consumed by a partition.
+
+stream.async.func private @dispatch(%arg0: i64, %arg1: !stream.resource<*>) -> !stream.resource<*>
+
+// CHECK-LABEL: @scfForProducingDispatchInput
+util.func public @scfForProducingDispatchInput(%arg0: !stream.resource<*>) -> !stream.resource<*> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c10 = arith.constant 10 : index
+ %c128 = arith.constant 128 : index
+ %c0_i64 = arith.constant 0 : i64
+ %c1_i64 = arith.constant 1 : i64
+
+ // CHECK: scf.for
+ %count = scf.for %i = %c0 to %c10 step %c1 iter_args(%iter = %c0_i64) -> (i64) {
+ %next = arith.addi %iter, %c1_i64 : i64
+ scf.yield %next : i64
+ }
+
+ // The dispatch should be in a partition that waits for the loop to complete.
+ // CHECK: stream.async.execute
+ // CHECK-NEXT: stream.async.call @dispatch
+ %result = stream.async.call @dispatch(%count, %arg0[%c0 to %c128 for %c128]) : (i64, !stream.resource<*>{%c128}) -> %arg0{%c128}
+ // CHECK: stream.timepoint.await
+ util.return %result : !stream.resource<*>
+}
+
+// -----
+
+// Tests nested scf.for loops with stream operations.
+// Both inner and outer loops should remain outside partitions.
+
+// CHECK-LABEL: @scfNestedFor
+util.func public @scfNestedFor(%arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c3 = arith.constant 3 : index
+ %c4 = arith.constant 4 : index
+ %c16 = arith.constant 16 : index
+ %c0_i32 = arith.constant 0 : i32
+
+ // CHECK: scf.for
+ %outer_result = scf.for %i = %c0 to %c3 step %c1 iter_args(%outer_iter = %c0_i32) -> (i32) {
+ // CHECK: scf.for
+ %inner_result = scf.for %j = %c0 to %c3 step %c1 iter_args(%inner_iter = %outer_iter) -> (i32) {
+ %offset = arith.muli %j, %c4 : index
+ %end = arith.addi %offset, %c4 : index
+ // CHECK: stream.async.execute
+ // CHECK-NEXT: stream.async.slice
+ // CHECK-NEXT: stream.async.transfer
+ %slice = stream.async.slice %arg0[%offset to %end] : !stream.resource<external>{%c16} -> !stream.resource<external>{%c4}
+ %transfer = stream.async.transfer %slice : !stream.resource<external>{%c4} -> !stream.resource<staging>{%c4}
+ // CHECK: stream.timepoint.await
+ %load = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c4} -> i32
+ %add = arith.addi %inner_iter, %load : i32
+ scf.yield %add : i32
+ }
+ scf.yield %inner_result : i32
+ }
+
+ // CHECK: stream.async.execute
+ // CHECK-NEXT: stream.async.splat
+ %splat = stream.async.splat %outer_result : i32 -> !stream.resource<external>{%c4}
+ util.return %splat : !stream.resource<external>
+}
+
+// -----
+
+// Tests scf.for loop that uses multiple values produced by a single partition
+// before the loop. The key expectations:
+// 1. The two slices are created in a single execute region before the loop.
+// 2. The loop does NOT await the timepoint before entering.
+// 3. Inside the loop body, both slices are used in a single execute region
+// that awaits the partition's timepoint.
+// 4. The two transfers are merged into one execute region.
+
+// CHECK-LABEL: @scfForWithMultipleRegionUses
+util.func public @scfForWithMultipleRegionUses(%arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %c0_i32 = arith.constant 0 : i32
+
+ // CHECK: %[[RESULTS:.+]]:2, %[[TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: with(%{{.+}} as %[[CAPTURE:.+]]: !stream.resource<external>{%c16})
+ // CHECK-SAME: -> (!stream.resource<external>{%c8}, !stream.resource<external>{%c8})
+ // CHECK-NEXT: %[[SLICE0:.+]] = stream.async.slice %[[CAPTURE]][%c0 to %c8]
+ %slice0 = stream.async.slice %arg0[%c0 to %c8] : !stream.resource<external>{%c16} -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: %[[SLICE1:.+]] = stream.async.slice %[[CAPTURE]][%c8 to %c16]
+ %slice1 = stream.async.slice %arg0[%c8 to %c16] : !stream.resource<external>{%c16} -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: stream.yield %[[SLICE0]], %[[SLICE1]]
+
+ // CHECK-NOT: stream.timepoint.await
+ // CHECK: scf.for
+ %result = scf.for %i = %c0 to %c4 step %c1 iter_args(%iter = %c0_i32) -> (i32) {
+ // Loop uses both slices from the partition. Both transfers merged into one execute.
+ // CHECK: stream.async.execute await(%[[TIMEPOINT]])
+ // CHECK-NEXT: stream.async.transfer
+ %transfer0 = stream.async.transfer %slice0 : !stream.resource<external>{%c8} -> !stream.resource<staging>{%c8}
+ // CHECK-NEXT: stream.async.transfer
+ %transfer1 = stream.async.transfer %slice1 : !stream.resource<external>{%c8} -> !stream.resource<staging>{%c8}
+ // CHECK: stream.timepoint.await
+ %load0 = stream.async.load %transfer0[%c0] : !stream.resource<staging>{%c8} -> i32
+
+ // CHECK-NOT: stream.async.execute
+ // CHECK-NOT: stream.timepoint.await
+ %load1 = stream.async.load %transfer1[%c0] : !stream.resource<staging>{%c8} -> i32
+
+ %add0 = arith.addi %iter, %load0 : i32
+ %add1 = arith.addi %add0, %load1 : i32
+ scf.yield %add1 : i32
+ }
+
+ // CHECK: stream.async.execute
+ // CHECK-NEXT: stream.async.splat
+ %splat = stream.async.splat %result : i32 -> !stream.resource<external>{%c4}
+ util.return %splat : !stream.resource<external>
+}
+
+// -----
+
+// Tests scf.for loop with zero iterations.
+// The loop may not execute at all, which affects partitioning.
+
+// CHECK-LABEL: @scfForZeroIterations
+util.func public @scfForZeroIterations(%arg0: !stream.resource<external>, %lb: index, %ub: index) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c16 = arith.constant 16 : index
+ %c0_i32 = arith.constant 0 : i32
+
+ // Clone is NOT materialized outside the loop since the loop may not execute.
+ // CHECK-NOT: stream.async.execute
+ // CHECK-NOT: stream.async.clone
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c16} -> !stream.resource<external>{%c16}
+
+ // CHECK: scf.for
+ %result = scf.for %i = %lb to %ub step %c1 iter_args(%iter = %c0_i32) -> (i32) {
+ // Clone is materialized INSIDE the execute region within the loop body.
+ // CHECK: stream.async.execute
+ // CHECK-NEXT: stream.async.clone
+ // CHECK-NEXT: stream.async.slice
+ // CHECK-NEXT: stream.async.transfer
+ %slice = stream.async.slice %clone[%c0 to %c4] : !stream.resource<external>{%c16} -> !stream.resource<external>{%c4}
+ %transfer = stream.async.transfer %slice : !stream.resource<external>{%c4} -> !stream.resource<staging>{%c4}
+ // CHECK: stream.timepoint.await
+ %load = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c4} -> i32
+ %add = arith.addi %iter, %load : i32
+ scf.yield %add : i32
+ }
+
+ // CHECK: stream.async.execute
+ // CHECK-NEXT: stream.async.splat
+ %splat = stream.async.splat %result : i32 -> !stream.resource<external>{%c4}
+ util.return %splat : !stream.resource<external>
+}
+
+// -----
+
+// Tests scf.if with both then and else branches containing stream operations.
+// Both branches should be partitioned independently.
+
+// CHECK-LABEL: @scfIfThenElse
+util.func public @scfIfThenElse(%cond: i1, %arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+
+ // CHECK: scf.if
+ %result = scf.if %cond -> !stream.resource<external> {
+ // CHECK: %[[RESULTS_0:.+]], %[[RESULT_TIMEPOINT_1:.+]] = stream.async.execute
+ // CHECK-SAME: with(%{{.+}} as %[[ARG2:.+]]: !stream.resource<external>{%c16})
+ // CHECK-SAME: -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: %[[SLICE:.+]] = stream.async.slice %[[ARG2]][%c0 to %c8]
+ %slice = stream.async.slice %arg0[%c0 to %c8] : !stream.resource<external>{%c16} -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: stream.yield %[[SLICE]]
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[RESULT_TIMEPOINT_1]] => %[[RESULTS_0]]
+ // CHECK: scf.yield %[[AWAITED]]
+ scf.yield %slice : !stream.resource<external>
+ } else {
+ // CHECK: %[[RESULTS_0:.+]], %[[RESULT_TIMEPOINT_1:.+]] = stream.async.execute
+ // CHECK-SAME: with(%{{.+}} as %[[ARG2:.+]]: !stream.resource<external>{%c16})
+ // CHECK-SAME: -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: %[[SLICE:.+]] = stream.async.slice %[[ARG2]][%c8 to %c16]
+ %slice = stream.async.slice %arg0[%c8 to %c16] : !stream.resource<external>{%c16} -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: stream.yield %[[SLICE]]
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[RESULT_TIMEPOINT_1]] => %[[RESULTS_0]]
+ // CHECK: scf.yield %[[AWAITED]]
+ scf.yield %slice : !stream.resource<external>
+ }
+
+ // CHECK: %[[RESULTS:.+]], %[[RESULT_TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: with(%{{.+}} as %[[ARG2:.+]]: !stream.resource<external>{%c8})
+ // CHECK-SAME: -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: %[[CLONE:.+]] = stream.async.clone %[[ARG2]]
+ %transfer = stream.async.transfer %result : !stream.resource<external>{%c8} -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: stream.yield %[[CLONE]]
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[RESULT_TIMEPOINT]] => %[[RESULTS]]
+ // CHECK: util.return %[[AWAITED]]
+ util.return %transfer : !stream.resource<external>
+}
+
+// -----
+
+// Tests scf.if with only then branch (no else).
+// The then branch contains stream operations.
+
+// CHECK-LABEL: @scfIfThenOnly
+util.func public @scfIfThenOnly(%cond: i1, %arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+
+ // CHECK: scf.if
+ %result = scf.if %cond -> !stream.resource<external> {
+ // CHECK: %[[RESULTS:.+]], %[[RESULT_TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: with(%{{.+}} as %[[ARG2:.+]]: !stream.resource<external>{%c8})
+ // CHECK-SAME: -> !stream.resource<external>{%c4}
+ // CHECK-NEXT: %[[SLICE:.+]] = stream.async.slice %[[ARG2]][%c0 to %c4]
+ %slice = stream.async.slice %arg0[%c0 to %c4] : !stream.resource<external>{%c8} -> !stream.resource<external>{%c4}
+ // CHECK-NEXT: stream.yield %[[SLICE]]
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[RESULT_TIMEPOINT]] => %[[RESULTS]]
+ // CHECK: scf.yield %[[AWAITED]]
+ scf.yield %slice : !stream.resource<external>
+ } else {
+ scf.yield %arg0 : !stream.resource<external>
+ }
+
+ // CHECK: util.return
+ util.return %result : !stream.resource<external>
+}
+
+// -----
+
+// Tests scf.if using cloned resource from partition.
+// The if operation depends on a partition output.
+
+// CHECK-LABEL: @scfIfWithDependency
+util.func public @scfIfWithDependency(%cond: i1, %arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+
+ // CHECK: %[[RESULTS:.+]], %[[RESULT_TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: with(%{{.+}} as %[[ARG2:.+]]: !stream.resource<external>{%c8})
+ // CHECK-SAME: -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: %[[CLONE_OP:.+]] = stream.async.clone %[[ARG2]]
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c8} -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: stream.yield %[[CLONE_OP]]
+
+ // The else branch uses clone directly (not in a stream op), so await must happen before scf.if.
+ // CHECK: %[[CLONE_READY:.+]] = stream.timepoint.await %[[RESULT_TIMEPOINT]] => %[[RESULTS]]
+ // CHECK: scf.if
+ %result = scf.if %cond -> !stream.resource<external> {
+ // Then branch uses cloned resource in stream op - can reference original clone.
+ // CHECK: %[[RESULTS_0:.+]], %[[RESULT_TIMEPOINT_1:.+]] = stream.async.execute await(%[[RESULT_TIMEPOINT]])
+ // CHECK-SAME: with(%[[RESULTS]] as %[[ARG2:.+]]: !stream.resource<external>{%c8})
+ // CHECK-SAME: -> !stream.resource<external>{%c4}
+ // CHECK-NEXT: %[[SLICE:.+]] = stream.async.slice %[[ARG2]][%c0 to %c4]
+ %slice = stream.async.slice %clone[%c0 to %c4] : !stream.resource<external>{%c8} -> !stream.resource<external>{%c4}
+ // CHECK-NEXT: stream.yield %[[SLICE]]
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[RESULT_TIMEPOINT_1]] => %[[RESULTS_0]]
+ // CHECK: scf.yield %[[AWAITED]]
+ scf.yield %slice : !stream.resource<external>
+ } else {
+ // CHECK: scf.yield %[[CLONE_READY]]
+ scf.yield %clone : !stream.resource<external>
+ }
+
+ // CHECK: util.return
+ util.return %result : !stream.resource<external>
+}
+
+// -----
+
+// Tests scf.if result being used by a dispatch.
+// The if produces a value consumed by a partition.
+
+stream.async.func private @dispatch(%arg0: !stream.resource<*>, %arg1: index) -> %arg0
+
+// CHECK-LABEL: @scfIfProducingDispatchInput
+util.func public @scfIfProducingDispatchInput(%cond: i1, %arg0: !stream.resource<*>) -> !stream.resource<*> {
+ %c0 = arith.constant 0 : index
+ %c128 = arith.constant 128 : index
+ %c256 = arith.constant 256 : index
+
+ %size = scf.if %cond -> index {
+ // To prevent folding to arith.select.
+ "some.sideeffect"() : () -> ()
+ scf.yield %c128 : index
+ } else {
+ scf.yield %c256 : index
+ }
+
+ // The dispatch should be in a partition.
+ // CHECK: %[[RESULTS:.+]], %[[RESULT_TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-NEXT: stream.async.call @dispatch
+ %result = stream.async.call @dispatch(%arg0[%c0 to %size for %size], %size) : (!stream.resource<*>{%size}, index) -> %arg0{%size}
+ // CHECK: stream.timepoint.await %[[RESULT_TIMEPOINT]] => %[[RESULTS]]
+ util.return %result : !stream.resource<*>
+}
+
+// -----
+
+// Tests scf.if nested inside scf.for with stream operations.
+// Both control flow operations should remain outside partitions.
+
+// CHECK-LABEL: @scfNestedIfInFor
+util.func public @scfNestedIfInFor(%arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %c0_i32 = arith.constant 0 : i32
+
+ // CHECK: scf.for
+ %result = scf.for %i = %c0 to %c4 step %c1 iter_args(%iter = %c0_i32) -> (i32) {
+ %c2 = arith.constant 2 : index
+ %rem = arith.remui %i, %c2 : index
+ %cond = arith.cmpi eq, %rem, %c0 : index
+
+ // CHECK: scf.if
+ %value = scf.if %cond -> i32 {
+ // CHECK: %[[RESULTS_0:.+]], %[[RESULT_TIMEPOINT_1:.+]] = stream.async.execute
+ // CHECK-SAME: with(%{{.+}} as %[[ARG3:.+]]: !stream.resource<external>{%c16})
+ // CHECK-SAME: -> !stream.resource<staging>{%c4}
+ // CHECK-NEXT: %[[SLICE:.+]] = stream.async.slice %[[ARG3]][%c0 to %c4]
+ // CHECK-NEXT: %[[TRANSFER:.+]] = stream.async.transfer %[[SLICE]]
+ %slice = stream.async.slice %arg0[%c0 to %c4] : !stream.resource<external>{%c16} -> !stream.resource<external>{%c4}
+ %transfer = stream.async.transfer %slice : !stream.resource<external>{%c4} -> !stream.resource<staging>{%c4}
+ // CHECK-NEXT: stream.yield %[[TRANSFER]]
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[RESULT_TIMEPOINT_1]] => %[[RESULTS_0]]
+ // CHECK: %[[LOAD:.+]] = stream.async.load %[[AWAITED]][%c0]
+ %load = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c4} -> i32
+ // CHECK: scf.yield %[[LOAD]]
+ scf.yield %load : i32
+ } else {
+ // CHECK: %[[RESULTS_0:.+]], %[[RESULT_TIMEPOINT_1:.+]] = stream.async.execute
+ // CHECK-SAME: with(%{{.+}} as %[[ARG3:.+]]: !stream.resource<external>{%c16})
+ // CHECK-SAME: -> !stream.resource<staging>{%c8}
+ // CHECK-NEXT: %[[SLICE:.+]] = stream.async.slice %[[ARG3]][%c8 to %c16]
+ // CHECK-NEXT: %[[TRANSFER:.+]] = stream.async.transfer %[[SLICE]]
+ %slice = stream.async.slice %arg0[%c8 to %c16] : !stream.resource<external>{%c16} -> !stream.resource<external>{%c8}
+ %transfer = stream.async.transfer %slice : !stream.resource<external>{%c8} -> !stream.resource<staging>{%c8}
+ // CHECK-NEXT: stream.yield %[[TRANSFER]]
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[RESULT_TIMEPOINT_1]] => %[[RESULTS_0]]
+ // CHECK: %[[LOAD:.+]] = stream.async.load %[[AWAITED]][%c0]
+ %load = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c8} -> i32
+ // CHECK: scf.yield %[[LOAD]]
+ scf.yield %load : i32
+ }
+
+ %add = arith.addi %iter, %value : i32
+ scf.yield %add : i32
+ }
+
+ // CHECK: %[[RESULTS:.+]], %[[RESULT_TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-NEXT: stream.async.splat
+ %splat = stream.async.splat %result : i32 -> !stream.resource<external>{%c4}
+ // CHECK: stream.timepoint.await %[[RESULT_TIMEPOINT]] => %[[RESULTS]]
+ util.return %splat : !stream.resource<external>
+}
+
+// -----
+
+// Tests basic scf.while loop with stream operations in before/after regions.
+// Both regions should be partitioned independently.
+
+// CHECK-LABEL: @scfWhileBasic
+// CHECK-SAME: (%[[ARG0:.+]]: !stream.resource<external>)
+util.func public @scfWhileBasic(%arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c10 = arith.constant 10 : index
+ %c0_i32 = arith.constant 0 : i32
+ %c100_i32 = arith.constant 100 : i32
+
+ // CHECK: %[[WHILE_RESULTS:.+]]:2 = scf.while
+ // CHECK-SAME: (%[[ARG1:.+]] = %c0_i32, %[[ARG2:.+]] = %c0)
+ // CHECK-SAME: (i32, index) -> (i32, index)
+ %result:2 = scf.while (%iter = %c0_i32, %count = %c0) : (i32, index) -> (i32, index) {
+ %cond = arith.cmpi slt, %count, %c10 : index
+ // CHECK: scf.condition
+ scf.condition(%cond) %iter, %count : i32, index
+ } do {
+ ^bb0(%iter: i32, %count: index):
+ // CHECK: ^bb0(%[[ITER:.+]]: i32, %[[COUNT:.+]]: index):
+ // CHECK: %[[RESULTS:.+]], %[[TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: with(%[[ARG0]] as %[[CAPTURE:.+]]: !stream.resource<external>{%c8})
+ // CHECK-SAME: -> !stream.resource<staging>{%c4}
+ // CHECK-NEXT: %[[SLICE:.+]] = stream.async.slice %[[CAPTURE]][%c0 to %c4]
+ // CHECK-NEXT: %[[TRANSFER:.+]] = stream.async.transfer %[[SLICE]]
+ // CHECK-NEXT: stream.yield %[[TRANSFER]]
+ %slice = stream.async.slice %arg0[%c0 to %c4] : !stream.resource<external>{%c8} -> !stream.resource<external>{%c4}
+ %transfer = stream.async.transfer %slice : !stream.resource<external>{%c4} -> !stream.resource<staging>{%c4}
+ // CHECK: %[[READY:.+]] = stream.timepoint.await %[[TIMEPOINT]] => %[[RESULTS]]
+ %load = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c4} -> i32
+ // CHECK: %[[LOAD:.+]] = stream.async.load %[[READY]][%c0]
+ %add = arith.addi %iter, %load : i32
+ %next_count = arith.addi %count, %c1 : index
+ scf.yield %add, %next_count : i32, index
+ }
+
+ // CHECK: %[[SPLAT_RESULTS:.+]], %[[SPLAT_TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: -> !stream.resource<external>{%c4}
+ // CHECK-NEXT: %[[SPLAT:.+]] = stream.async.splat %[[WHILE_RESULTS]]#0
+ // CHECK-NEXT: stream.yield %[[SPLAT]]
+ %splat = stream.async.splat %result#0 : i32 -> !stream.resource<external>{%c4}
+ // CHECK: %[[FINAL:.+]] = stream.timepoint.await %[[SPLAT_TIMEPOINT]] => %[[SPLAT_RESULTS]]
+ // CHECK: util.return %[[FINAL]]
+ util.return %splat : !stream.resource<external>
+}
+
+// -----
+
+// Tests scf.while loop using cloned resource from partition.
+// The while depends on a partition output.
+
+// CHECK-LABEL: @scfWhileWithDependency
+// CHECK-SAME: (%[[ARG0:.+]]: !stream.resource<external>)
+util.func public @scfWhileWithDependency(%arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c5 = arith.constant 5 : index
+ %c0_i32 = arith.constant 0 : i32
+
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c8} -> !stream.resource<external>{%c8}
+
+ // CHECK: %[[WHILE_RESULTS:.+]]:2 = scf.while
+ // CHECK-SAME: (%[[ARG1:.+]] = %c0_i32, %[[ARG2:.+]] = %c0)
+ // CHECK-SAME: (i32, index) -> (i32, index)
+ %result:2 = scf.while (%iter = %c0_i32, %count = %c0) : (i32, index) -> (i32, index) {
+ %cond = arith.cmpi slt, %count, %c5 : index
+ scf.condition(%cond) %iter, %count : i32, index
+ } do {
+ ^bb0(%iter: i32, %count: index):
+ // CHECK: ^bb0(%[[ITER:.+]]: i32, %[[COUNT:.+]]: index):
+ // While uses the cloned resource.
+ // CHECK: %[[RESULTS:.+]], %[[TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: with(%[[ARG0]] as %[[CAPTURE:.+]]: !stream.resource<external>{%c8})
+ // CHECK-SAME: -> !stream.resource<staging>{%c4}
+ // CHECK-NEXT: %[[CLONE:.+]] = stream.async.clone %[[CAPTURE]]
+ // CHECK-NEXT: %[[SLICE:.+]] = stream.async.slice %[[CLONE]][%c0 to %c4]
+ // CHECK-NEXT: %[[TRANSFER:.+]] = stream.async.transfer %[[SLICE]]
+ // CHECK-NEXT: stream.yield %[[TRANSFER]]
+ %slice = stream.async.slice %clone[%c0 to %c4] : !stream.resource<external>{%c8} -> !stream.resource<external>{%c4}
+ %transfer = stream.async.transfer %slice : !stream.resource<external>{%c4} -> !stream.resource<staging>{%c4}
+ // CHECK: %[[READY:.+]] = stream.timepoint.await %[[TIMEPOINT]] => %[[RESULTS]]
+ %load = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c4} -> i32
+ // CHECK: %[[LOAD:.+]] = stream.async.load %[[READY]][%c0]
+ %add = arith.addi %iter, %load : i32
+ %next_count = arith.addi %count, %c1 : index
+ scf.yield %add, %next_count : i32, index
+ }
+
+ // CHECK: %[[SPLAT_RESULTS:.+]], %[[SPLAT_TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: -> !stream.resource<external>{%c4}
+ // CHECK-NEXT: %[[SPLAT:.+]] = stream.async.splat %[[WHILE_RESULTS]]#0
+ // CHECK-NEXT: stream.yield %[[SPLAT]]
+ %splat = stream.async.splat %result#0 : i32 -> !stream.resource<external>{%c4}
+ // CHECK: %[[FINAL:.+]] = stream.timepoint.await %[[SPLAT_TIMEPOINT]] => %[[SPLAT_RESULTS]]
+ // CHECK: util.return %[[FINAL]]
+ util.return %splat : !stream.resource<external>
+}
+
+// -----
+
+// Tests scf.while result being used by a dispatch.
+// The while produces a value consumed by a partition.
+
+stream.async.func private @dispatch(%arg0: i64, %arg1: !stream.resource<*>) -> !stream.resource<*>
+
+// CHECK-LABEL: @scfWhileProducingDispatchInput
+// CHECK-SAME: (%[[ARG0:.+]]: !stream.resource<*>)
+util.func public @scfWhileProducingDispatchInput(%arg0: !stream.resource<*>) -> !stream.resource<*> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c128 = arith.constant 128 : index
+ %c0_i64 = arith.constant 0 : i64
+ %c1_i64 = arith.constant 1 : i64
+
+ // CHECK: %[[WHILE_RESULTS:.+]]:2 = scf.while
+ // CHECK-SAME: (%[[ARG1:.+]] = %c0_i64, %[[ARG2:.+]] = %c0)
+ // CHECK-SAME: (i64, index) -> (i64, index)
+ %count:2 = scf.while (%iter = %c0_i64, %idx = %c0) : (i64, index) -> (i64, index) {
+ %cond = arith.cmpi slt, %idx, %c10 : index
+ scf.condition(%cond) %iter, %idx : i64, index
+ } do {
+ ^bb0(%iter: i64, %idx: index):
+ %next = arith.addi %iter, %c1_i64 : i64
+ %next_idx = arith.addi %idx, %c1 : index
+ scf.yield %next, %next_idx : i64, index
+ }
+
+ // The dispatch should be in a partition.
+ // CHECK: %[[RESULTS:.+]], %[[TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: with(%[[ARG0]] as %[[CAPTURE:.+]]: !stream.resource<*>{%c128})
+ // CHECK-SAME: -> {{.+}}{%c128}
+ // CHECK-NEXT: %[[CALL:.+]] = stream.async.call @dispatch(%[[WHILE_RESULTS]]#0, %[[CAPTURE]][%c0 to %c128 for %c128])
+ // CHECK-NEXT: stream.yield %[[CALL]]
+ %result = stream.async.call @dispatch(%count#0, %arg0[%c0 to %c128 for %c128]) : (i64, !stream.resource<*>{%c128}) -> %arg0{%c128}
+ // CHECK: %[[FINAL:.+]] = stream.timepoint.await %[[TIMEPOINT]] => %[[RESULTS]]
+ // CHECK: util.return %[[FINAL]]
+ util.return %result : !stream.resource<*>
+}
+
+// -----
+
+// Tests combination of scf.if, scf.for, and scf.while with stream operations.
+// All control flow operations should remain outside partitions.
+
+// CHECK-LABEL: @scfMixedControlFlow
+util.func public @scfMixedControlFlow(%cond: i1, %arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %c0_i32 = arith.constant 0 : i32
+
+ // CHECK-NOT: stream.async.execute
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c16} -> !stream.resource<external>{%c16}
+
+ // CHECK-NOT: stream.timepoint.await
+ // CHECK: scf.if
+ %if_result = scf.if %cond -> i32 {
+ // CHECK: scf.for
+ %for_result = scf.for %i = %c0 to %c2 step %c1 iter_args(%iter = %c0_i32) -> (i32) {
+ // CHECK: stream.async.execute
+ // CHECK-NEXT: stream.async.clone
+ // CHECK-NEXT: stream.async.slice
+ // CHECK-NEXT: stream.async.transfer
+ %slice = stream.async.slice %clone[%c0 to %c4] : !stream.resource<external>{%c16} -> !stream.resource<external>{%c4}
+ %transfer = stream.async.transfer %slice : !stream.resource<external>{%c4} -> !stream.resource<staging>{%c4}
+ // CHECK: stream.timepoint.await
+ %load = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c4} -> i32
+ %add = arith.addi %iter, %load : i32
+ scf.yield %add : i32
+ }
+ scf.yield %for_result : i32
+ } else {
+ // CHECK: scf.while
+ %while_result:2 = scf.while (%iter = %c0_i32, %count = %c0) : (i32, index) -> (i32, index) {
+ %cond_while = arith.cmpi slt, %count, %c2 : index
+ scf.condition(%cond_while) %iter, %count : i32, index
+ } do {
+ ^bb0(%iter: i32, %count: index):
+ // CHECK: stream.async.execute
+ // CHECK-NEXT: stream.async.clone
+ // CHECK-NEXT: stream.async.slice
+ // CHECK-NEXT: stream.async.transfer
+ %slice = stream.async.slice %clone[%c8 to %c16] : !stream.resource<external>{%c16} -> !stream.resource<external>{%c8}
+ %transfer = stream.async.transfer %slice : !stream.resource<external>{%c8} -> !stream.resource<staging>{%c8}
+ // CHECK: stream.timepoint.await
+ %load = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c8} -> i32
+ %add = arith.addi %iter, %load : i32
+ %next_count = arith.addi %count, %c1 : index
+ scf.yield %add, %next_count : i32, index
+ }
+ scf.yield %while_result#0 : i32
+ }
+
+ // CHECK: stream.async.execute
+ // CHECK-NEXT: stream.async.splat
+ %splat = stream.async.splat %if_result : i32 -> !stream.resource<external>{%c4}
+ // CHECK: stream.timepoint.await
+ util.return %splat : !stream.resource<external>
+}
+
+// -----
+
+// Tests control flow operations across multiple partitions.
+// Stream operations before and after control flow should form separate partitions.
+
+stream.async.func private @dispatch0(%arg0: !stream.resource<*>, %arg1: index) -> %arg0
+stream.async.func private @dispatch1(%arg0: !stream.resource<*>, %arg1: index) -> %arg0
+
+// CHECK-LABEL: @scfControlFlowAcrossPartitions
+util.func public @scfControlFlowAcrossPartitions(%cond: i1, %arg0: !stream.resource<*>, %loop_bound: index) -> (!stream.resource<*>, i32) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c128 = arith.constant 128 : index
+ %c0_i32 = arith.constant 0 : i32
+
+ // First partition before control flow.
+ // CHECK: %[[RESULTS0:.+]], %[[TIMEPOINT0:.+]] = stream.async.execute
+ // CHECK-NEXT: stream.async.call @dispatch0
+ %result0 = stream.async.call @dispatch0(%arg0[%c0 to %c128 for %c128], %c128) : (!stream.resource<*>{%c128}, index) -> %arg0{%c128}
+
+ // Control flow in the middle - loop that loads data.
+ // CHECK-NOT: stream.timepoint.await
+ // CHECK: scf.for
+ %loop_result:2 = scf.for %i = %c0 to %loop_bound step %c1 iter_args(%iter = %result0, %sum = %c0_i32) -> (!stream.resource<*>, i32) {
+ // CHECK: stream.async.execute await(%[[TIMEPOINT0]])
+ // CHECK-NEXT: stream.async.transfer
+ %transfer = stream.async.transfer %iter : !stream.resource<*>{%c128} -> !stream.resource<staging>{%c128}
+ // CHECK: stream.timepoint.await
+ %loaded = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c128} -> i32
+ %next_sum = arith.addi %sum, %loaded : i32
+ scf.yield %iter, %next_sum : !stream.resource<*>, i32
+ }
+
+ // Second partition after control flow.
+ // CHECK: stream.async.execute await(%[[TIMEPOINT0]])
+ // CHECK-NEXT: stream.async.call @dispatch1
+ %result1 = stream.async.call @dispatch1(%loop_result#0[%c0 to %c128 for %c128], %c128) : (!stream.resource<*>{%c128}, index) -> %loop_result#0{%c128}
+ // CHECK: stream.timepoint.await
+ util.return %result1, %loop_result#1 : !stream.resource<*>, i32
+}
+
+// -----
+
+// Tests control flow with device affinities.
+// Partitions should respect affinity boundaries with control flow.
+
+util.global private @device_a : !hal.device
+util.global private @device_b : !hal.device
+
+// CHECK-LABEL: @scfControlFlowWithAffinities
+util.func public @scfControlFlowWithAffinities(%cond: i1, %arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c0_i32 = arith.constant 0 : i32
+
+ // Partition on device_a.
+ // CHECK-NOT: stream.async.execute
+ %clone_a = stream.async.clone on(#hal.device.affinity<@device_a>) %arg0 : !stream.resource<external>{%c8} -> !stream.resource<external>{%c8}
+
+ // Control flow on host.
+ // CHECK-NOT: stream.timepoint.await
+ // CHECK: scf.for
+ %result = scf.for %i = %c0 to %c2 step %c1 iter_args(%iter = %c0_i32) -> (i32) {
+ // Partition on device_a for clone inside loop.
+ // CHECK: stream.async.execute on(#hal.device.affinity<@device_a>)
+ // CHECK-NEXT: stream.async.clone
+ // Partition on device_b inside loop.
+ // CHECK: stream.async.execute on(#hal.device.affinity<@device_b>)
+ // CHECK-NEXT: stream.async.slice
+ // CHECK-NEXT: stream.async.transfer
+ %slice = stream.async.slice on(#hal.device.affinity<@device_b>) %clone_a[%c0 to %c4] : !stream.resource<external>{%c8} -> !stream.resource<external>{%c4}
+ %transfer = stream.async.transfer %slice : !stream.resource<external>{%c4} from(#hal.device.affinity<@device_b>) -> to(#hal.device.affinity<@device_b>) !stream.resource<staging>{%c4}
+ // CHECK: stream.timepoint.await
+ %load = stream.async.load %transfer[%c0] : !stream.resource<staging>{%c4} -> i32
+ %add = arith.addi %iter, %load : i32
+ scf.yield %add : i32
+ }
+
+ // Final partition on device_a.
+ // CHECK: stream.async.execute on(#hal.device.affinity<@device_a>)
+ // CHECK-NEXT: stream.async.splat
+ %splat = stream.async.splat on(#hal.device.affinity<@device_a>) %result : i32 -> !stream.resource<external>{%c4}
+ // CHECK: stream.timepoint.await
+ util.return %splat : !stream.resource<external>
+}
+
+// -----
+
+// Tests that operations captured in nested regions (scf.for) are not grouped
+// with operations that depend on the nested region's results. This prevents
+// circular dependencies where the partition would use the scf result while the
+// scf captures values from the partition.
+
+// CHECK-LABEL: @scfNestedCapture
+// CHECK-SAME: (%[[ARG0:.+]]: !stream.resource<*>, %[[ARG1:.+]]: index)
+util.func public @scfNestedCapture(%arg0: !stream.resource<*>, %arg1: index) -> !stream.resource<*> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c0_i64 = arith.constant 0 : i64
+
+ // The cloned resource will be duplicated - one copy in the loop's partition,
+ // one copy in the dispatch's partition. No standalone clone partition.
+ %cloned = stream.async.clone %arg0 : !stream.resource<*>{%arg1} -> !stream.resource<*>{%arg1}
+
+ // Loop gets its own materialized clone.
+ // CHECK: %[[LOOP_RESULT:.+]] = scf.for
+ %loop_result = scf.for %i = %c0 to %c10 step %c1 iter_args(%acc = %c0_i64) -> (i64) {
+ // Operations inside loop should be partitioned with cloned op materialized.
+ // CHECK: stream.async.execute
+ // CHECK-SAME: with(%[[ARG0]]
+ // CHECK-NEXT: stream.async.clone
+ // CHECK: stream.async.transfer
+ %slice = stream.async.slice %cloned[%c0 to %arg1] : !stream.resource<*>{%arg1} -> !stream.resource<*>{%arg1}
+ %transfer = stream.async.transfer %slice : !stream.resource<*>{%arg1} -> !stream.resource<staging>{%arg1}
+ // CHECK: stream.timepoint.await
+ // CHECK: stream.async.load
+ %loaded = stream.async.load %transfer[%c0] : !stream.resource<staging>{%arg1} -> i32
+ %loaded_i64 = arith.extsi %loaded : i32 to i64
+ %new_acc = arith.addi %acc, %loaded_i64 : i64
+ scf.yield %new_acc : i64
+ }
+
+ // Dispatch doesn't use %cloned, only %loop_result, so no clone is materialized here.
+ // Before the fix, partitioning would try to group %cloned with this dispatch,
+ // creating a circular dependency (dispatch needs loop result, loop needs cloned).
+ // After fix, %cloned is only materialized where used (in loop partition).
+ // CHECK: stream.async.execute
+ // CHECK-NOT: stream.async.clone
+ // CHECK: stream.async.dispatch
+ %dispatched = stream.async.dispatch @ex::@dispatch[%c1](%loop_result) : (i64) -> !stream.resource<*>{%arg1}
+
+ util.return %dispatched : !stream.resource<*>
+}
+
+// -----
+
+// Tests scf.if with divergent captures - only one branch captures from parent.
+// Partitioning must be conservative and treat the containing scf.if as a hazard.
+
+// CHECK-LABEL: @scfIfDivergentCapture
+// CHECK-SAME: (%[[ARG0:.+]]: !stream.resource<*>, %[[ARG1:.+]]: index, %[[COND:.+]]: i1)
+util.func public @scfIfDivergentCapture(%arg0: !stream.resource<*>, %arg1: index, %cond: i1) -> !stream.resource<*> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c0_i64 = arith.constant 0 : i64
+
+ %cloned = stream.async.clone %arg0 : !stream.resource<*>{%arg1} -> !stream.resource<*>{%arg1}
+
+ // Only true branch captures %cloned - it will get its own materialized clone.
+ // CHECK: %[[IF_RESULT:.+]] = scf.if
+ %if_result = scf.if %cond -> (i64) {
+ // CHECK: stream.async.execute
+ // CHECK-SAME: with(%[[ARG0]]
+ // CHECK-NEXT: stream.async.clone
+ %slice = stream.async.slice %cloned[%c0 to %arg1] : !stream.resource<*>{%arg1} -> !stream.resource<*>{%arg1}
+ %transfer = stream.async.transfer %slice : !stream.resource<*>{%arg1} -> !stream.resource<staging>{%arg1}
+ %loaded = stream.async.load %transfer[%c0] : !stream.resource<staging>{%arg1} -> i32
+ %result = arith.extsi %loaded : i32 to i64
+ scf.yield %result : i64
+ } else {
+ // False branch doesn't capture anything.
+ scf.yield %c0_i64 : i64
+ }
+
+ // Dispatch doesn't use %cloned, so no clone is materialized here.
+ // CHECK: stream.async.execute
+ // CHECK-NOT: stream.async.clone
+ // CHECK: stream.async.dispatch
+ %dispatched = stream.async.dispatch @ex::@dispatch[%c1](%if_result) : (i64) -> !stream.resource<*>{%arg1}
+
+ util.return %dispatched : !stream.resource<*>
+}
+
+// -----
+
+// Tests nested scf.for loops where inner loop captures from grandparent scope.
+// The containing operation walk should find the outer scf.for correctly.
+
+// CHECK-LABEL: @scfNestedLoopCapture
+// CHECK-SAME: (%[[ARG0:.+]]: !stream.resource<*>, %[[ARG1:.+]]: index)
+util.func public @scfNestedLoopCapture(%arg0: !stream.resource<*>, %arg1: index) -> !stream.resource<*> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c5 = arith.constant 5 : index
+ %c0_i64 = arith.constant 0 : i64
+
+ %cloned = stream.async.clone %arg0 : !stream.resource<*>{%arg1} -> !stream.resource<*>{%arg1}
+
+ // Outer loop.
+ // Inner loop gets its own materialized clone.
+ // CHECK: %[[OUTER_RESULT:.+]] = scf.for
+ %outer_result = scf.for %i = %c0 to %c5 step %c1 iter_args(%outer_acc = %c0_i64) -> (i64) {
+ // Inner loop captures %cloned from grandparent block.
+ // It gets its own materialized clone.
+ // CHECK: scf.for
+ %inner_result = scf.for %j = %c0 to %c5 step %c1 iter_args(%inner_acc = %c0_i64) -> (i64) {
+ // CHECK: stream.async.execute
+ // CHECK-SAME: with(%[[ARG0]]
+ // CHECK-NEXT: stream.async.clone
+ %slice = stream.async.slice %cloned[%c0 to %arg1] : !stream.resource<*>{%arg1} -> !stream.resource<*>{%arg1}
+ %transfer = stream.async.transfer %slice : !stream.resource<*>{%arg1} -> !stream.resource<staging>{%arg1}
+ %loaded = stream.async.load %transfer[%c0] : !stream.resource<staging>{%arg1} -> i32
+ %loaded_i64 = arith.extsi %loaded : i32 to i64
+ %new_inner = arith.addi %inner_acc, %loaded_i64 : i64
+ scf.yield %new_inner : i64
+ }
+ %new_outer = arith.addi %outer_acc, %inner_result : i64
+ scf.yield %new_outer : i64
+ }
+
+ // Dispatch doesn't use %cloned, so no clone is materialized here.
+ // CHECK: stream.async.execute
+ // CHECK-NOT: stream.async.clone
+ // CHECK: stream.async.dispatch
+ %dispatched = stream.async.dispatch @ex::@dispatch[%c1](%outer_result) : (i64) -> !stream.resource<*>{%arg1}
+
+ util.return %dispatched : !stream.resource<*>
+}
+
+// -----
+
+// Tests multiple scf operations at the same level capturing the same resource.
+// All should mark the resource as hazardous for grouping with their consumers.
+
+// CHECK-LABEL: @scfMultipleCaptures
+// CHECK-SAME: (%[[ARG0:.+]]: !stream.resource<*>, %[[ARG1:.+]]: index)
+util.func public @scfMultipleCaptures(%arg0: !stream.resource<*>, %arg1: index) -> !stream.resource<*> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c5 = arith.constant 5 : index
+ %c0_i64 = arith.constant 0 : i64
+
+ %cloned = stream.async.clone %arg0 : !stream.resource<*>{%arg1} -> !stream.resource<*>{%arg1}
+
+ // First loop gets its own materialized clone.
+ // CHECK: %[[RESULT1:.+]] = scf.for
+ %result1 = scf.for %i = %c0 to %c5 step %c1 iter_args(%acc = %c0_i64) -> (i64) {
+ // CHECK: stream.async.execute
+ // CHECK-SAME: with(%[[ARG0]]
+ // CHECK-NEXT: stream.async.clone
+ %slice = stream.async.slice %cloned[%c0 to %arg1] : !stream.resource<*>{%arg1} -> !stream.resource<*>{%arg1}
+ %transfer = stream.async.transfer %slice : !stream.resource<*>{%arg1} -> !stream.resource<staging>{%arg1}
+ %loaded = stream.async.load %transfer[%c0] : !stream.resource<staging>{%arg1} -> i32
+ %loaded_i64 = arith.extsi %loaded : i32 to i64
+ %new_acc = arith.addi %acc, %loaded_i64 : i64
+ scf.yield %new_acc : i64
+ }
+
+ // Second loop also gets its own materialized clone.
+ // CHECK: %[[RESULT2:.+]] = scf.for
+ %result2 = scf.for %i = %c0 to %c5 step %c1 iter_args(%acc = %c0_i64) -> (i64) {
+ // CHECK: stream.async.execute
+ // CHECK-SAME: with(%[[ARG0]]
+ // CHECK-NEXT: stream.async.clone
+ %slice = stream.async.slice %cloned[%c0 to %arg1] : !stream.resource<*>{%arg1} -> !stream.resource<*>{%arg1}
+ %transfer = stream.async.transfer %slice : !stream.resource<*>{%arg1} -> !stream.resource<staging>{%arg1}
+ %loaded = stream.async.load %transfer[%c0] : !stream.resource<staging>{%arg1} -> i32
+ %loaded_i64 = arith.extsi %loaded : i32 to i64
+ %new_acc = arith.addi %acc, %loaded_i64 : i64
+ scf.yield %new_acc : i64
+ }
+
+ // Dispatch uses both loop results. Should be in separate partition from %cloned.
+ // CHECK: stream.async.execute
+ // CHECK-NOT: stream.async.clone
+ // CHECK: stream.async.dispatch
+ %sum = arith.addi %result1, %result2 : i64
+ %dispatched = stream.async.dispatch @ex::@dispatch[%c1](%sum) : (i64) -> !stream.resource<*>{%arg1}
+
+ util.return %dispatched : !stream.resource<*>
+}
+
+// -----
+
+// Tests that nested scf operations correctly handle iter_args pattern.
+// Iter_args should allow values to flow through the loop without creating hazards.
+
+// CHECK-LABEL: @scfIterArgsPattern
+util.func public @scfIterArgsPattern(%arg0: !stream.resource<*>, %arg1: index) -> !stream.resource<*> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c5 = arith.constant 5 : index
+
+ // Loop with iter_args - resource flows through the loop correctly.
+ // CHECK: scf.for
+ // CHECK-SAME: iter_args
+ %loop_result = scf.for %i = %c0 to %c5 step %c1 iter_args(%iter_resource = %arg0) -> !stream.resource<*> {
+ // Each iteration partitions independently.
+ // CHECK: stream.async.execute
+ // CHECK: stream.async.dispatch
+ %updated = stream.async.dispatch @ex::@dispatch[%c1](%iter_resource[%c0 to %arg1 for %arg1]) : (!stream.resource<*>{%arg1}) -> !stream.resource<*>{%arg1}
+ // CHECK: stream.timepoint.await
+ scf.yield %updated : !stream.resource<*>
+ }
+
+ util.return %loop_result : !stream.resource<*>
+}
+
+// -----
+
+// Tests mixed iter_args (good) and captures (hazard) in same loop.
+// This is a common real-world pattern where a loop updates one resource via iter_args
+// while reading from another resource captured from parent scope.
+
+// CHECK-LABEL: @scfMixedIterArgsAndCaptures
+util.func public @scfMixedIterArgsAndCaptures(%arg0: !stream.resource<*>, %arg1: !stream.resource<*>, %size: index) -> (!stream.resource<*>, i64) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c5 = arith.constant 5 : index
+ %c0_i64 = arith.constant 0 : i64
+
+ // This resource will be materialized in the loop.
+ %read_only = stream.async.clone %arg1 : !stream.resource<*>{%size} -> !stream.resource<*>{%size}
+
+ // Loop: %arg0 flows through iter_args (updated each iteration),
+ // %read_only materialized in loop (read each iteration),
+ // %accumulator tracks scalar across iterations.
+ // CHECK: %[[LOOP_RESULTS:.+]]:2 = scf.for
+ // CHECK-SAME: iter_args(
+ %loop_resource, %loop_sum = scf.for %i = %c0 to %c5 step %c1
+ iter_args(%iter_resource = %arg0, %iter_sum = %c0_i64) -> (!stream.resource<*>, i64) {
+
+ // Inside loop: both iter_arg resource and materialized clone are used.
+ // They must be in the same execute region since the dispatch needs the
+ // updated iter_arg, and the transfer needs the materialized clone.
+ // CHECK: stream.async.execute
+ // CHECK: stream.async.clone
+ // CHECK: stream.async.dispatch
+ // CHECK: stream.async.transfer
+ %updated = stream.async.dispatch @ex::@dispatch[%c1](%iter_resource[%c0 to %size for %size])
+ : (!stream.resource<*>{%size}) -> !stream.resource<*>{%size}
+
+ %slice = stream.async.slice %read_only[%c0 to %size] : !stream.resource<*>{%size} -> !stream.resource<*>{%size}
+ %transfer = stream.async.transfer %slice : !stream.resource<*>{%size} -> !stream.resource<staging>{%size}
+ %loaded = stream.async.load %transfer[%c0] : !stream.resource<staging>{%size} -> i32
+ %loaded_i64 = arith.extsi %loaded : i32 to i64
+ %new_sum = arith.addi %iter_sum, %loaded_i64 : i64
+
+ scf.yield %updated, %new_sum : !stream.resource<*>, i64
+ }
+
+ // The key test: the final dispatch uses loop results but NOT the captured
+ // %read_only. This proves that %read_only is correctly tracked as a hazard
+ // and not incorrectly merged with operations that only use the loop's outputs.
+ // CHECK: stream.async.execute
+ // CHECK-NOT: stream.async.clone
+ // CHECK: stream.async.dispatch
+ %final = stream.async.dispatch @ex::@dispatch[%c1](%loop_sum, %loop_resource[%c0 to %size for %size])
+ : (i64, !stream.resource<*>{%size}) -> !stream.resource<*>{%size}
+
+ util.return %final, %loop_sum : !stream.resource<*>, i64
+}
+
+// -----
+
+// Tests splat with optimization_barrier in same block + use inside scf.if.
+// This reproduces the bug from if.mlir: splat with barrier use AND use in nested
+// region was not being placed in any partition, causing verification failure.
+// The barrier forces the splat into a partition before the scf.if.
+
+// CHECK-LABEL: @splatWithBarrierAndScfIfUse
+util.func public @splatWithBarrierAndScfIfUse(%cond: i1, %arg0: !stream.resource<external>) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c10_i32 = arith.constant 10 : i32
+ %c4 = arith.constant 4 : index
+
+ // CHECK: %[[RESULTS:.+]], %[[TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: with() -> !stream.resource<external>{%c4}
+ // CHECK-NEXT: %[[SPLAT:.+]] = stream.async.splat %c10_i32
+ // CHECK-NEXT: stream.yield %[[SPLAT]]
+ %splat = stream.async.splat %c10_i32 : i32 -> !stream.resource<external>{%c4}
+
+ // Same-block non-streamable user forces await before scf.if.
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[TIMEPOINT]]
+ %barrier = util.optimization_barrier %splat : !stream.resource<external>
+
+ // CHECK: scf.if
+ %result = scf.if %cond -> !stream.resource<external> {
+ // Nested region can use both the awaited value and reference original.
+ // CHECK: stream.async.execute
+ // CHECK: stream.async.dispatch
+ %dispatch = stream.async.dispatch @dispatch::@entry(%barrier[%c0 to %c4 for %c4], %splat[%c0 to %c4 for %c4]) : (!stream.resource<external>{%c4}, !stream.resource<external>{%c4}) -> !stream.resource<external>{%c4}
+ scf.yield %dispatch : !stream.resource<external>
+ } else {
+ scf.yield %barrier : !stream.resource<external>
+ }
+
+ util.return %result : !stream.resource<external>
+}
+
+stream.executable private @dispatch {
+ stream.executable.export public @entry
+ builtin.module {
+ func.func @entry(%arg0: !stream.binding, %arg1: !stream.binding, %arg2: !stream.binding) {
+ return
+ }
+ }
+}
+
+// -----
+
+// Tests clone with optimization_barrier in same block + use inside scf.for.
+// Similar to test above but with clone + scf.for instead of splat + scf.if.
+
+// CHECK-LABEL: @cloneWithBarrierAndScfForUse
+util.func public @cloneWithBarrierAndScfForUse(%arg0: !stream.resource<external>, %ub: index) -> i64 {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c0_i64 = arith.constant 0 : i64
+
+ // CHECK: %[[RESULTS:.+]], %[[TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: with(%{{.+}} as %[[ARG:.+]]: !stream.resource<external>{%c8})
+ // CHECK-SAME: -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: %[[CLONE:.+]] = stream.async.clone %[[ARG]]
+ // CHECK-NEXT: stream.yield %[[CLONE]]
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c8} -> !stream.resource<external>{%c8}
+
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[TIMEPOINT]]
+ %barrier = util.optimization_barrier %clone : !stream.resource<external>
+
+ // CHECK: scf.for
+ %result = scf.for %i = %c0 to %ub step %c1 iter_args(%iter = %c0_i64) -> (i64) {
+ // CHECK: stream.async.execute
+ // CHECK: stream.async.slice
+ // CHECK: stream.async.transfer
+ %results_0, %result_timepoint_1 = stream.async.execute with(%barrier as %arg1: !stream.resource<external>{%c8}) -> !stream.resource<staging>{%c4} {
+ %3 = stream.async.slice %arg1[%c0 to %c4] : !stream.resource<external>{%c8} -> !stream.resource<external>{%c4}
+ %4 = stream.async.transfer %3 : !stream.resource<external>{%c4} -> !stream.resource<staging>{%c4}
+ stream.yield %4 : !stream.resource<staging>{%c4}
+ } => !stream.timepoint
+ %1 = stream.timepoint.await %result_timepoint_1 => %results_0 : !stream.resource<staging>{%c4}
+ %2 = stream.async.load %1[%c0] : !stream.resource<staging>{%c4} -> i32
+ %ext = arith.extsi %2 : i32 to i64
+ %next = arith.addi %iter, %ext : i64
+ scf.yield %next : i64
+ }
+
+ util.return %result : i64
+}
+
+// -----
+
+// Tests splat with multiple optimization_barrier uses in same block.
+// Ensures that multiple same-block non-streamable users work correctly.
+
+// CHECK-LABEL: @splatWithMultipleBarriers
+util.func public @splatWithMultipleBarriers() -> (!stream.resource<external>, !stream.resource<external>) {
+ %c10_i32 = arith.constant 10 : i32
+ %c4 = arith.constant 4 : index
+
+ // CHECK: %[[RESULTS:.+]], %[[TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-NEXT: %[[SPLAT:.+]] = stream.async.splat
+ // CHECK-NEXT: stream.yield %[[SPLAT]]
+ %splat = stream.async.splat %c10_i32 : i32 -> !stream.resource<external>{%c4}
+
+ // Both barriers use the same awaited value.
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[TIMEPOINT]]
+ // CHECK: util.optimization_barrier %[[AWAITED]]
+ %barrier1 = util.optimization_barrier %splat : !stream.resource<external>
+ // CHECK: util.optimization_barrier %[[AWAITED]]
+ %barrier2 = util.optimization_barrier %splat : !stream.resource<external>
+
+ util.return %barrier1, %barrier2 : !stream.resource<external>, !stream.resource<external>
+}
+
+// -----
+
+// Tests clone with optimization_barrier + use in deeply nested regions.
+// Ensures barrier causes partitioning even with deep nesting (scf.for > scf.if).
+
+// CHECK-LABEL: @clonableBarrierDeeplyNested
+util.func public @clonableBarrierDeeplyNested(%arg0: !stream.resource<external>, %cond: i1, %ub: index) -> !stream.resource<external> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+
+ // CHECK: %[[RESULTS:.+]], %[[TIMEPOINT:.+]] = stream.async.execute
+ // CHECK-SAME: with(%{{.+}} as %[[ARG:.+]]: !stream.resource<external>{%c8})
+ // CHECK-SAME: -> !stream.resource<external>{%c8}
+ // CHECK-NEXT: %[[CLONE:.+]] = stream.async.clone %[[ARG]]
+ // CHECK-NEXT: stream.yield %[[CLONE]]
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c8} -> !stream.resource<external>{%c8}
+
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[TIMEPOINT]]
+ %barrier = util.optimization_barrier %clone : !stream.resource<external>
+
+ // CHECK: scf.for
+ %result = scf.for %i = %c0 to %ub step %c1 iter_args(%iter = %arg0) -> !stream.resource<external> {
+ // CHECK: scf.if
+ %next = scf.if %cond -> !stream.resource<external> {
+ // Deeply nested use of barrier.
+ // CHECK: stream.async.execute
+ // CHECK: stream.async.slice
+ %results_0, %result_timepoint_1 = stream.async.execute with(%barrier as %arg1: !stream.resource<external>{%c8}) -> !stream.resource<external>{%c4} {
+ %2 = stream.async.slice %arg1[%c0 to %c4] : !stream.resource<external>{%c8} -> !stream.resource<external>{%c4}
+ stream.yield %2 : !stream.resource<external>{%c4}
+ } => !stream.timepoint
+ %1 = stream.timepoint.await %result_timepoint_1 => %results_0 : !stream.resource<external>{%c4}
+ scf.yield %1 : !stream.resource<external>
+ } else {
+ scf.yield %iter : !stream.resource<external>
+ }
+ scf.yield %next : !stream.resource<external>
+ }
+
+ util.return %result : !stream.resource<external>
+}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution_timeline_aware.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution_timeline_aware.mlir
new file mode 100644
index 0000000..1933bde
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_execution_timeline_aware.mlir
@@ -0,0 +1,118 @@
+// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(util.func(iree-stream-schedule-execution))" %s | FileCheck %s
+
+// Tests that timeline-aware ops don't get immediate awaits.
+// The TimelineAwareOpInterface allows ops to participate in timeline scheduling.
+
+// CHECK-LABEL: @timelineAwareCall
+util.func public @timelineAwareCall(%resource: !stream.resource<external>, %signal_fence: !stream.test.fence) -> !stream.resource<external> {
+ %c16384 = arith.constant 16384 : index
+
+ // Create async execution that produces a timepoint.
+ // CHECK: %[[RESULT:.+]], %[[TIMEPOINT:.+]] = stream.async.execute
+ %result, %timepoint = stream.async.execute with(%resource as %arg0: !stream.resource<external>{%c16384}) -> !stream.resource<external>{%c16384} {
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c16384} -> !stream.resource<external>{%c16384}
+ stream.yield %clone : !stream.resource<external>{%c16384}
+ } => !stream.timepoint
+
+ // Export timepoint to fence for timeline-aware op.
+ // CHECK-NOT: stream.timepoint.await %[[TIMEPOINT]]
+ // CHECK: %[[WAIT_FENCE:.+]] = stream.timepoint.export %[[TIMEPOINT]]
+ %wait_fence = stream.timepoint.export %timepoint => (!stream.test.fence)
+
+ // Call timeline-aware op - should NOT have await before it.
+ // CHECK: stream.test.timeline_aware(%[[RESULT]]) waits(%[[WAIT_FENCE]]) signals(%{{.+}})
+ %computed = stream.test.timeline_aware(%result) waits(%wait_fence) signals(%signal_fence) : (!stream.resource<external>) -> !stream.resource<external>
+
+ // Import signal fence back to timepoint.
+ // CHECK: %[[IMPORTED_TP:.+]] = stream.timepoint.import
+ %imported_tp = stream.timepoint.import %signal_fence : (!stream.test.fence) => !stream.timepoint
+
+ // Subsequent async.execute should await imported timepoint.
+ // CHECK: %[[FINAL:.+]], %{{.+}} = stream.async.execute await(%[[IMPORTED_TP]])
+ %awaited = stream.timepoint.await %imported_tp => %computed : !stream.resource<external>{%c16384}
+ %final, %final_tp = stream.async.execute with(%awaited as %arg0: !stream.resource<external>{%c16384}) -> !stream.resource<external>{%c16384} {
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c16384} -> !stream.resource<external>{%c16384}
+ stream.yield %clone : !stream.resource<external>{%c16384}
+ } => !stream.timepoint
+
+ // CHECK: util.return %[[FINAL]]
+ util.return %final : !stream.resource<external>
+}
+
+// -----
+
+// Tests timeline awareness between multiple sequential timeline-aware calls.
+
+// CHECK-LABEL: @sequentialTimelineAwareCalls
+util.func public @sequentialTimelineAwareCalls(%resource: !stream.resource<external>, %fence_a: !stream.test.fence, %fence_b: !stream.test.fence) -> !stream.resource<external> {
+ %c16384 = arith.constant 16384 : index
+
+ // First async region.
+ // CHECK: %[[R1:.+]], %[[T1:.+]] = stream.async.execute
+ %r1, %t1 = stream.async.execute with(%resource as %arg0: !stream.resource<external>{%c16384}) -> !stream.resource<external>{%c16384} {
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c16384} -> !stream.resource<external>{%c16384}
+ stream.yield %clone : !stream.resource<external>{%c16384}
+ } => !stream.timepoint
+
+ // First timeline-aware call - should NOT have await before it.
+ // CHECK-NOT: stream.timepoint.await %[[T1]]
+ // CHECK: %[[WAIT_A:.+]] = stream.timepoint.export %[[T1]]
+ %wait_a = stream.timepoint.export %t1 => (!stream.test.fence)
+ // CHECK: %[[RESULT_A:.+]] = stream.test.timeline_aware(%[[R1]]) waits(%[[WAIT_A]]) signals(%{{.+}})
+ %result_a = stream.test.timeline_aware(%r1) waits(%wait_a) signals(%fence_a) : (!stream.resource<external>) -> !stream.resource<external>
+
+ // Import result from first call.
+ // CHECK: %[[TP_A:.+]] = stream.timepoint.import
+ %tp_a = stream.timepoint.import %fence_a : (!stream.test.fence) => !stream.timepoint
+
+ // Second async region - pass optimizes await into async.execute.
+ // CHECK: %[[R2:.+]], %[[T2:.+]] = stream.async.execute await(%[[TP_A]]) => with(%[[RESULT_A]] as %{{.+}}: !stream.resource<external>{%c16384})
+ %awaited_a = stream.timepoint.await %tp_a => %result_a : !stream.resource<external>{%c16384}
+ %r2, %t2 = stream.async.execute with(%awaited_a as %arg0: !stream.resource<external>{%c16384}) -> !stream.resource<external>{%c16384} {
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c16384} -> !stream.resource<external>{%c16384}
+ stream.yield %clone : !stream.resource<external>{%c16384}
+ } => !stream.timepoint
+
+ // Second timeline-aware call - again, no await before.
+ // CHECK-NOT: stream.timepoint.await %[[T2]]
+ // CHECK: %[[WAIT_B:.+]] = stream.timepoint.export %[[T2]]
+ %wait_b = stream.timepoint.export %t2 => (!stream.test.fence)
+ // CHECK: %[[RESULT_B:.+]] = stream.test.timeline_aware(%[[R2]]) waits(%[[WAIT_B]]) signals(%{{.+}})
+ %result_b = stream.test.timeline_aware(%r2) waits(%wait_b) signals(%fence_b) : (!stream.resource<external>) -> !stream.resource<external>
+
+ // Import and wait for second call.
+ // CHECK: %[[TP_B:.+]] = stream.timepoint.import
+ %tp_b = stream.timepoint.import %fence_b : (!stream.test.fence) => !stream.timepoint
+ // CHECK: %[[AWAITED_B:.+]] = stream.timepoint.await %[[TP_B]] => %[[RESULT_B]]
+ %awaited_b = stream.timepoint.await %tp_b => %result_b : !stream.resource<external>{%c16384}
+
+ // CHECK: util.return %[[AWAITED_B]]
+ util.return %awaited_b : !stream.resource<external>
+}
+
+// -----
+
+// Tests that non-timeline-aware ops still get normal awaits.
+
+util.func private @normal_func(!stream.resource<external>) -> !stream.resource<external>
+
+// CHECK-LABEL: @nonTimelineAwareCall
+util.func public @nonTimelineAwareCall(%resource: !stream.resource<external>) -> !stream.resource<external> {
+ %c16384 = arith.constant 16384 : index
+
+ // Async work.
+ // CHECK: %[[RESULT:.+]], %[[TIMEPOINT:.+]] = stream.async.execute
+ %result, %timepoint = stream.async.execute with(%resource as %arg0: !stream.resource<external>{%c16384}) -> !stream.resource<external>{%c16384} {
+ %clone = stream.async.clone %arg0 : !stream.resource<external>{%c16384} -> !stream.resource<external>{%c16384}
+ stream.yield %clone : !stream.resource<external>{%c16384}
+ } => !stream.timepoint
+
+ // Normal call without timeline awareness - SHOULD have await.
+ // CHECK: %[[AWAITED:.+]] = stream.timepoint.await %[[TIMEPOINT]] => %[[RESULT]]
+ %awaited = stream.timepoint.await %timepoint => %result : !stream.resource<external>{%c16384}
+
+ // CHECK: util.call @normal_func(%[[AWAITED]])
+ %called = util.call @normal_func(%awaited) : (!stream.resource<external>) -> !stream.resource<external>
+
+ util.return %called : !stream.resource<external>
+}
