[stream] Don't group timepoint awaits in different blocks (#11823)
We sort various await ops over the same timepoint according to their
order in the same block, but the list was pulling in await ops in other
blocks too. This caused crash.
In general, unifying ops across the blocks is triky given we need to
consider dominance and referened values in each op. Prefer to be
conservative for now.
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
index 3b75830..d67ec79 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
@@ -2584,6 +2584,7 @@
// there. We rely on other canonicalizers to sink things such that
// (hopefully) we get them directly accessible here.
if (use.getOwner() == op) continue;
+ if (op->getBlock() != use.getOwner()->getBlock()) continue;
if (dominanceInfo.dominates(use.getOwner(), op)) continue;
auto awaitOp = dyn_cast<TimepointAwaitOp>(use.getOwner());
if (!awaitOp ||
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/test/timepoint_folding.mlir b/compiler/src/iree/compiler/Dialect/Stream/IR/test/timepoint_folding.mlir
index 146ae41..f8edd40 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/test/timepoint_folding.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/test/timepoint_folding.mlir
@@ -210,6 +210,37 @@
// -----
+// Tests that the pattern doesn't kick in when the same timepoint are waited in
+// different blocks.
+
+func.func private @materializeResource() -> !stream.resource<*>
+
+// CHECK-LABEL: @DontGroupAwaitsByTimepointAcrossBlocks
+func.func @DontGroupAwaitsByTimepointAcrossBlocks(
+ %arg0: !stream.timepoint,
+ %arg1: !stream.resource<*>,
+ %arg2: i1
+) -> !stream.resource<*> {
+ %c100 = arith.constant 100 : index
+ %c101 = arith.constant 101 : index
+ %0 = stream.timepoint.await %arg0 => %arg1 : !stream.resource<*>{%c100}
+ // CHECK: cf.cond_br
+ cf.cond_br %arg2, ^bb0, ^bb1
+// CHECK: ^bb
+^bb0:
+ // CHECK: stream.timepoint.await %arg0 => %arg1
+ return %0 : !stream.resource<*>
+// CHECK: ^bb
+^bb1:
+ // CHECK: %[[R:.+]] = call @materializeResource
+ %r = call @materializeResource() : () -> !stream.resource<*>
+ // CHECK: stream.timepoint.await %arg0 => %[[R]]
+ %1 = stream.timepoint.await %arg0 => %r : !stream.resource<*>{%c101}
+ return %1 : !stream.resource<*>
+}
+
+// -----
+
// CHECK-LABEL: @FoldDuplicateAwaitResources
func.func @FoldDuplicateAwaitResources(
%arg0: !stream.timepoint,
