Adding verification after EncodeHostTensors/EncodeDeviceTensors passes. (#16255)
This is a lighter check than the existing verification of async lowering
that only checks that there are no tensor ops remaining.
Fixes #16231.
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
index d6b4e5a..f741745 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
@@ -125,6 +125,10 @@
addCleanupPatterns(passManager);
+ // Everything must now be in stream.async.* form but we don't yet have
+ // lifetime assigned.
+ passManager.addPass(IREE::Stream::createVerifyLoweringToAsyncResourcesPass());
+
// Materialize copy-on-write behavior with explicit stream.async.* ops.
// This will insert a lot of copies, so follow it up with a pass that elides
// ones that aren't needed. This is easier to verify than if there was one
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
index f7d760b..fabe47f 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
@@ -492,9 +492,14 @@
let summary = "Verifies that input dialects are converted to stream.tensor.* ops.";
}
+def VerifyLoweringToAsyncResourcesPass :
+ Pass<"iree-stream-verify-lowering-to-async-resources", "mlir::ModuleOp"> {
+ let summary = "Verifies that all stream.tensor.* ops and types are fully lowered to stream.async.* resource ops.";
+}
+
def VerifyLoweringToAsyncPass :
Pass<"iree-stream-verify-lowering-to-async", "mlir::ModuleOp"> {
- let summary = "Verifies that all stream.tensor.* ops and types are fully lowered to stream.async.* ops.";
+ let summary = "Verifies that all stream.tensor.* ops and types are fully lowered to stream.async.* ops and all resources have an assigned lifetime.";
}
def VerifyAsyncAccessRangesPass :
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/VerifyLowerings.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/VerifyLowerings.cpp
index b515a97..2860f48 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/VerifyLowerings.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/VerifyLowerings.cpp
@@ -24,6 +24,7 @@
#define GEN_PASS_DEF_VERIFYINPUTPASS
#define GEN_PASS_DEF_VERIFYLOWERINGTOTENSORSPASS
+#define GEN_PASS_DEF_VERIFYLOWERINGTOASYNCRESOURCESPASS
#define GEN_PASS_DEF_VERIFYLOWERINGTOASYNCPASS
#define GEN_PASS_DEF_VERIFYLOWERINGTOCMDPASS
#include "iree/compiler/Dialect/Stream/Transforms/Passes.h.inc"
@@ -297,7 +298,28 @@
};
//===----------------------------------------------------------------------===//
-// --iree-stream-verify-lowering-to-tensors
+// --iree-stream-verify-lowering-to-async-resources
+//===----------------------------------------------------------------------===//
+
+struct VerifyLoweringToAsyncResourcesPass
+ : public IREE::Stream::impl::VerifyLoweringToAsyncResourcesPassBase<
+ VerifyLoweringToAsyncResourcesPass> {
+ void runOnOperation() override {
+ // We cannot have stream.cmd.* ops mixed with stream.tensor/async.* ops
+ // as they use different memory models. We need to allow them through,
+ // though, to allow for compiler re-entrancy.
+ Verifier verifier;
+ setupDefaultOpLegality(verifier);
+ markTensorInputsIllegal(verifier);
+ markStreamTensorOpsIllegal(verifier);
+ if (failed(verifier.run(getOperation()))) {
+ return signalPassFailure();
+ }
+ }
+};
+
+//===----------------------------------------------------------------------===//
+// --iree-stream-verify-lowering-to-async
//===----------------------------------------------------------------------===//
struct VerifyLoweringToAsyncPass
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/refine_usage.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/refine_usage.mlir
index b1cb56a..7276441 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/refine_usage.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/refine_usage.mlir
@@ -318,14 +318,14 @@
%c1 = arith.constant 1 : index
%c4 = arith.constant 4 : index
-// CHECK: %[[C0:.+]] = arith.constant 0 : index
-// CHECK: %[[C1:.+]] = arith.constant 1 : index
-// CHECK: %[[C4:.+]] = arith.constant 4 : index
-// CHECK: %[[DISP0:.+]] = stream.async.dispatch @dispatch0(%arg1[%[[C0]] to %[[ARG0]] for %[[ARG0]]]) : (!stream.resource<external>{%[[C4]]}) -> !stream.resource<transient>{%[[C4]]}
+ // CHECK: %[[C0:.+]] = arith.constant 0 : index
+ // CHECK: %[[C1:.+]] = arith.constant 1 : index
+ // CHECK: %[[C4:.+]] = arith.constant 4 : index
+ // CHECK: %[[DISP0:.+]] = stream.async.dispatch @dispatch0(%arg1[%[[C0]] to %[[ARG0]] for %[[ARG0]]]) : (!stream.resource<external>{%[[C4]]}) -> !stream.resource<transient>{%[[C4]]}
%dispatch6 = stream.async.dispatch @dispatch0(%arg1[%c0 to %arg0 for %arg0]) : (!stream.resource<*>{%c4}) -> !stream.resource<*>{%c4}
+
// CHECK: %[[FOR:.+]] = scf.for %[[ARG2:.+]] = %[[C0]] to %[[ARG0]] step %[[C1]] iter_args(%[[ARG3:.+]] = %[[DISP0]]) -> (!stream.resource<transient>) {
%for = scf.for %i = %c0 to %arg0 step %c1 iter_args(%arg3 = %dispatch6) -> (!stream.resource<*>) {
-
// CHECK: %[[DISP1:.+]] = stream.async.dispatch @dispatch1(%[[ARG3]][%[[C0]] to %[[ARG0]] for %[[ARG0]]]) : (!stream.resource<transient>{%[[C4]]}) -> !stream.resource<transient>{%[[C4]]}
// CHECK: %[[DISP2:.+]] = stream.async.dispatch @dispatch2(%[[DISP1]][%[[C0]] to %[[ARG0]] for %[[ARG0]]]) : (!stream.resource<transient>{%[[C4]]}) -> !stream.resource<transient>{%[[C4]]}
// CHECK: %[[DISP3:.+]] = stream.async.dispatch @dispatch3(%[[DISP2]][%[[C0]] to %[[ARG0]] for %[[ARG0]]]) : (!stream.resource<transient>{%[[C4]]}) -> !stream.resource<transient>{%[[C4]]}
@@ -335,6 +335,7 @@
// CHECK: scf.yield %[[DISP3]] : !stream.resource<transient>
scf.yield %dispatch3 : !stream.resource<*>
}
+
// CHECK: %[[DISP4:.+]] = stream.async.dispatch @dispatch4(%[[FOR]][%[[C0]] to %[[ARG0]] for %[[ARG0]]]) : (!stream.resource<transient>{%[[C4]]}) -> !stream.resource<external>{%[[C4]]}
%dispatch5 = stream.async.dispatch @dispatch4(%for[%c0 to %arg0 for %arg0]) : (!stream.resource<*>{%c4}) -> !stream.resource<*>{%c4}
%transfer = stream.async.transfer %dispatch5 : !stream.resource<*>{%arg0} -> !stream.resource<external>{%arg0}
