Properly marking stream alloca ops as non-pure and expanding folding. (#13863)
Previously only the AsyncCloneOp was folding allocas/splats but this is
useful at the TensorCloneOp level too for making IR more readable.
Progress on #13545 (makes the IR easier to read during stream
transformation).
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
index bad55a8..e6e1f17 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
@@ -313,6 +313,38 @@
}
};
+// Clones ops that prefer to be cloned directly.
+// This prevents us from splatting out a value and then cloning that (keeping
+// the memory live/etc) instead of just splatting it again on-demand.
+//
+// Example:
+// %0 = stream.async.splat %c123_i32
+// %1 = stream.async.clone %0
+// ->
+// %1 = stream.async.splat %c123_i32
+template <typename Op>
+struct PropagateClonableOps : public OpRewritePattern<Op> {
+ using OpRewritePattern<Op>::OpRewritePattern;
+ LogicalResult matchAndRewrite(Op cloneOp,
+ PatternRewriter &rewriter) const override {
+ if (cloneOp.use_empty()) return failure();
+ auto sourceOp =
+ cloneOp.getSource()
+ .template getDefiningOp<IREE::Stream::StreamableOpInterface>();
+ if (!sourceOp || !sourceOp.preferCloneToConsumers()) return failure();
+ for (auto &use :
+ llvm::make_early_inc_range(cloneOp.getResult().getUses())) {
+ rewriter.setInsertionPoint(use.getOwner());
+ auto clonedOp = rewriter.clone(*sourceOp);
+ use.set(clonedOp->getResult(0));
+ }
+ if (cloneOp.use_empty()) {
+ rewriter.eraseOp(cloneOp);
+ }
+ return success();
+ }
+};
+
// Materialize copy-on-write (🐄) ops where required for |rootValue|.
// Only valid in tensor/async ops - don't use with stream.cmd.*.
static bool materializeCOW(Location loc, Value rootValue, OpBuilder &builder) {
@@ -1121,11 +1153,11 @@
void TensorCloneOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- // TODO(benvanik): splat -> clone duplicates splat.
// TODO(benvanik): some way to reduce deep clone->clone->clone chains.
// TODO(benvanik): clone + slice => slice.
// TODO(benvanik): if both operand and result are used once then elide.
// (if not tied block/fn arguments)
+ results.insert<PropagateClonableOps<TensorCloneOp>>(context);
results.insert<ElideUnneededTensorClones>(context);
}
@@ -1326,44 +1358,10 @@
return {};
}
-namespace {
-
-// Clones ops that prefer to be cloned directly.
-// This prevents us from splatting out a value and then cloning that (keeping
-// the memory live/etc) instead of just splatting it again on-demand.
-//
-// Example:
-// %0 = stream.async.splat %c123_i32
-// %1 = stream.async.clone %0
-// ->
-// %1 = stream.async.splat %c123_i32
-struct PropagateClonableOps : public OpRewritePattern<AsyncCloneOp> {
- using OpRewritePattern::OpRewritePattern;
- LogicalResult matchAndRewrite(AsyncCloneOp cloneOp,
- PatternRewriter &rewriter) const override {
- if (cloneOp.use_empty()) return failure();
- auto sourceOp = cloneOp.getSource()
- .getDefiningOp<IREE::Stream::StreamableOpInterface>();
- if (!sourceOp || !sourceOp.preferCloneToConsumers()) return failure();
- for (auto &use :
- llvm::make_early_inc_range(cloneOp.getResult().getUses())) {
- rewriter.setInsertionPoint(use.getOwner());
- auto clonedOp = rewriter.clone(*sourceOp);
- use.set(clonedOp->getResult(0));
- }
- if (cloneOp.use_empty()) {
- rewriter.eraseOp(cloneOp);
- }
- return success();
- }
-};
-
-} // namespace
-
void AsyncCloneOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
// TODO(benvanik): some way to reduce deep clone->clone->clone chains.
- results.insert<PropagateClonableOps>(context);
+ results.insert<PropagateClonableOps<AsyncCloneOp>>(context);
results.insert<ElideUnusedOp<AsyncCloneOp>>(context);
}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
index e7699cf..7a730c4 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
@@ -862,6 +862,10 @@
return success();
}
+bool TensorEmptyOp::isMetadata() { return true; }
+
+bool TensorEmptyOp::preferCloneToConsumers() { return true; }
+
//===----------------------------------------------------------------------===//
// stream.tensor.constant
//===----------------------------------------------------------------------===//
@@ -893,6 +897,8 @@
return success();
}
+bool TensorSplatOp::preferCloneToConsumers() { return true; }
+
//===----------------------------------------------------------------------===//
// stream.tensor.clone
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
index 716b697..d6e3111 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
@@ -26,12 +26,13 @@
// Generic resource ops
//===----------------------------------------------------------------------===//
-def Stream_ResourceAllocOp : Stream_PureOp<"resource.alloc", [
+def Stream_ResourceAllocOp : Stream_Op<"resource.alloc", [
DeclareOpInterfaceMethods<Stream_AffinityOp, [
"getAffinity",
"setAffinity",
]>,
Util_SizeAwareOp,
+ AlwaysSpeculatable,
MemoryEffects<[MemAlloc]>,
]> {
let summary = [{allocates a persistent value with undefined contents}];
@@ -79,13 +80,14 @@
let hasCanonicalizer = 1;
}
-def Stream_ResourceAllocaOp : Stream_PureOp<"resource.alloca", [
+def Stream_ResourceAllocaOp : Stream_Op<"resource.alloca", [
DeclareOpInterfaceMethods<Stream_AffinityOp, [
"getAffinity",
"setAffinity",
]>,
Stream_TimelineOp,
Util_SizeAwareOp,
+ AlwaysSpeculatable,
MemoryEffects<[MemAlloc]>,
]> {
let summary = [{allocates a transient value with undefined contents}];
@@ -747,7 +749,10 @@
def Stream_TensorEmptyOp : Stream_PureOp<"tensor.empty", [
Stream_AffinityOp,
- Stream_StreamableOp,
+ DeclareOpInterfaceMethods<Stream_StreamableOp, [
+ "isMetadata",
+ "preferCloneToConsumers",
+ ]>,
Stream_TensorPhaseOp,
Util_ShapeAwareOp,
Util_SizeAwareOp,
@@ -836,7 +841,9 @@
def Stream_TensorSplatOp : Stream_PureOp<"tensor.splat", [
Stream_AffinityOp,
- Stream_StreamableOp,
+ DeclareOpInterfaceMethods<Stream_StreamableOp, [
+ "preferCloneToConsumers",
+ ]>,
Stream_TensorPhaseOp,
Util_ShapeAwareOp,
Util_SizeAwareOp,
@@ -1331,7 +1338,7 @@
// Resource transfer ops
//===----------------------------------------------------------------------===//
-def Stream_AsyncAllocaOp : Stream_PureOp<"async.alloca", [
+def Stream_AsyncAllocaOp : Stream_Op<"async.alloca", [
DeclareOpInterfaceMethods<Stream_AffinityOp, [
"getAffinity",
"setAffinity",
@@ -1342,6 +1349,7 @@
"preferCloneToConsumers",
]>,
Util_SizeAwareOp,
+ AlwaysSpeculatable,
MemoryEffects<[MemAlloc]>,
]> {
let summary = [{allocates a transient value with undefined contents}];
