[Stream] Fold transferred stream clones into minimal slices before loads (#24740)
This PR Extends clone folding for `stream.async.load` and
`stream.tensor.load` to
handle the staging transfer introduced before loads.
Both stream.async.load and stream.tensor.load require their source to be
a staging resource, while stream.async.clone and stream.tensor.clone
always produce non-staging resources. As a result, loads from cloned
resources are always preceded by a stream.async.transfer.
This change matches the pattern:
`
clone -> transfer -> load
`
and rewrites it to:
`
slice -> transfer -> load
`
by creating a minimal slice of the original source, transferring only
that slice to staging, and loading from it. This avoids cloning and
transferring the entire resource when only a single value is accessed,
while preserving the required staging semantics.
---------
Signed-off-by: LekkalaSravya3 <lekkala.sravya@multicorewareinc.com>
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
index 3dca743..42f7041 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
@@ -1410,14 +1410,90 @@
//===----------------------------------------------------------------------===//
// stream.tensor.load
//===----------------------------------------------------------------------===//
+namespace {
+
+// Replaces a load from a transferred clone with a load from a
+// transferred single-element slice of the original source.
+struct FoldTensorCloneIntoLoad : OpRewritePattern<TensorLoadOp> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(TensorLoadOp loadOp,
+ PatternRewriter &rewriter) const override {
+ // Match: stream.tensor.load <- stream.async.transfer
+ auto transferOp = loadOp.getSource().getDefiningOp<AsyncTransferOp>();
+ if (!transferOp) {
+ return failure();
+ }
+
+ if (!transferOp.getResult().hasOneUse()) {
+ return failure();
+ }
+
+ // Match: stream.async.transfer <- stream.tensor.clone
+ auto cloneOp = transferOp.getSource().getDefiningOp<TensorCloneOp>();
+ if (!cloneOp) {
+ return failure();
+ }
+
+ // Clones are allowed to change shape and element type ,Only fold when the
+ // clone is a pure copy so slicing the original source is equivalent to
+ // slicing the clone result.
+ if (cloneOp.getSourceEncoding() != cloneOp.getResultEncoding() ||
+ cloneOp.getResultEncoding() != loadOp.getSourceEncoding()) {
+ return failure();
+ }
+ auto sourceType = dyn_cast<RankedTensorType>(loadOp.getSourceEncoding());
+ if (!sourceType) {
+ return failure();
+ }
+ int64_t rank = sourceType.getRank();
+
+ auto loc = loadOp.getLoc();
+
+ // Create a single-element slice (length 1 in every dimension).
+ auto one = arith::ConstantIndexOp::create(rewriter, loc, 1);
+ auto zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ SmallVector<Value> lengths(rank, one);
+ SmallVector<Value> zeroIndices(rank, zero);
+
+ SmallVector<int64_t> unitShape(rank, 1);
+ auto sliceType = RankedTensorType::get(
+ unitShape, sourceType.getElementType(), sourceType.getEncoding());
+
+ Value sliceSize = TensorSizeOfOp::create(
+ rewriter, loc, rewriter.getIndexType(), TypeAttr::get(sliceType),
+ ValueRange{}, cloneOp.getAffinityAttr());
+
+ // Slice the original tensor instead of the cloned tensor.
+ auto sliceOp = TensorSliceOp::create(
+ rewriter, loc, cloneOp.getSource().getType(), cloneOp.getSource(),
+ cloneOp.getSourceEncoding(), cloneOp.getSourceEncodingDims(),
+ cloneOp.getSourceSize(), loadOp.getIndices(), lengths, sliceType,
+ ValueRange{}, sliceSize, cloneOp.getAffinityAttr());
+
+ // Transfer the sliced tensor to staging.
+ auto stagedSliceOp = AsyncTransferOp::create(
+ rewriter, loc, transferOp.getResult().getType(), sliceOp.getResult(),
+ sliceSize, sliceSize, transferOp.getSourceAffinityAttr(),
+ transferOp.getResultAffinityAttr());
+
+ // Load from the staged slice at zero indices.
+ rewriter.replaceOpWithNewOp<TensorLoadOp>(
+ loadOp, loadOp.getResult().getType(), stagedSliceOp.getResult(),
+ sliceType, ValueRange{}, sliceSize, zeroIndices);
+
+ return success();
+ }
+};
+
+} // namespace
void TensorLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
// TODO(benvanik): splat + load -> splat value.
- // TODO(benvanik): clone + ex load -> slice (ranged) + load.
- // TODO(benvanik): slice + ex load -> slice (ranged) + load.
// TODO(benvanik): value->transfer->load -> value->slice->transfer->load?
// TODO(benvanik): combine multiple loads from the same target if contiguous.
+ results.insert<FoldTensorCloneIntoLoad>(context);
}
//===----------------------------------------------------------------------===//
@@ -2250,16 +2326,74 @@
}
};
+// Replaces a load from a transferred clone with a load from a
+// transferred byte-ranged slice of the original source.
+struct FoldAsyncCloneIntoLoad : OpRewritePattern<AsyncLoadOp> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(AsyncLoadOp loadOp,
+ PatternRewriter &rewriter) const override {
+ // Match: stream.async.load <- stream.async.transfer
+ auto transferOp = loadOp.getSource().getDefiningOp<AsyncTransferOp>();
+ if (!transferOp) {
+ return failure();
+ }
+ // Only fold if the transfer is not used by anything else: otherwise the
+ // original clone+transfer will still be required and we'd just be adding
+ // an additional slice+transfer alongside it.
+ if (!transferOp.getResult().hasOneUse()) {
+ return failure();
+ }
+
+ // Match: stream.async.transfer <- stream.async.clone
+ auto cloneOp = transferOp.getSource().getDefiningOp<AsyncCloneOp>();
+ if (!cloneOp) {
+ return failure();
+ }
+
+ auto loc = loadOp.getLoc();
+ int64_t byteSize =
+ IREE::Util::getRoundedElementByteWidth(loadOp.getResult().getType());
+
+ auto byteSizeValue =
+ arith::ConstantIndexOp::create(rewriter, loc, byteSize);
+ auto zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
+
+ // Compute the end of the byte range to slice.
+ auto sourceEnd = rewriter.createOrFold<arith::AddIOp>(
+ loc, loadOp.getSourceOffset(), byteSizeValue);
+
+ // Slice the original resource instead of the cloned resource.
+ auto sliceOp = AsyncSliceOp::create(
+ rewriter, loc, cloneOp.getSource().getType(), cloneOp.getSource(),
+ cloneOp.getSourceSize(), loadOp.getSourceOffset(), sourceEnd,
+ byteSizeValue, cloneOp.getAffinityAttr());
+
+ // Transfer the sliced resource to staging.
+ auto stagedSliceOp = AsyncTransferOp::create(
+ rewriter, loc, transferOp.getResult().getType(), sliceOp.getResult(),
+ byteSizeValue, byteSizeValue, transferOp.getSourceAffinityAttr(),
+ transferOp.getResultAffinityAttr());
+
+ // Load from the staged slice at offset 0.
+ rewriter.replaceOpWithNewOp<AsyncLoadOp>(
+ loadOp, loadOp.getResult().getType(), stagedSliceOp.getResult(),
+ byteSizeValue, zero);
+
+ return success();
+ }
+};
+
} // namespace
void AsyncLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
// TODO(benvanik): splat + load -> splat value.
- // TODO(benvanik): clone + ex load -> slice (ranged) + load.
// TODO(benvanik): slice + ex load -> slice (ranged) + load.
// TODO(benvanik): value->transfer->load -> value->slice->transfer->load?
// TODO(benvanik): combine multiple loads from the same target if contiguous.
results.insert<FoldAsyncLoadBitcast>(context);
+ results.insert<FoldAsyncCloneIntoLoad>(context);
}
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_folding.mlir b/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_folding.mlir
index b8977e4..49a3a32 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_folding.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/test/async_folding.mlir
@@ -501,6 +501,22 @@
// -----
+// CHECK-LABEL: @FoldAsyncCloneIntoLoad
+util.func private @FoldAsyncCloneIntoLoad(%arg0: !stream.resource<external>, %arg1: index, %arg2: index) -> f32 {
+ // CHECK-NOT: stream.async.clone
+ %0 = stream.async.clone %arg0 : !stream.resource<external>{%arg1} -> !stream.resource<*>{%arg1}
+ %1 = stream.async.transfer %0 : !stream.resource<*>{%arg1} -> !stream.resource<staging>{%arg1}
+ // CHECK: %[[END:.+]] = arith.addi %arg2, %c4 : index
+ // CHECK: %[[SLICE:.+]] = stream.async.slice %arg0[%arg2 to %[[END]]] : !stream.resource<external>{%arg1} -> !stream.resource<external>{%c4}
+ // CHECK: %[[STAGED:.+]] = stream.async.transfer %[[SLICE]] : !stream.resource<external>{%c4} -> !stream.resource<staging>{%c4}
+ // CHECK: %[[VALUE:.+]] = stream.async.load %[[STAGED]][%c0] : !stream.resource<staging>{%c4} -> f32
+ %2 = stream.async.load %1[%arg2] : !stream.resource<staging>{%arg1} -> f32
+ // CHECK: util.return %[[VALUE]]
+ util.return %2 : f32
+}
+
+// -----
+
// CHECK-LABEL: @FoldAsyncStoreBitcast
util.func private @FoldAsyncStoreBitcast(%arg0: !stream.resource<staging>, %arg1: index, %arg2: f32) -> !stream.resource<staging> {
%c0 = arith.constant 0 : index
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_folding.mlir b/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_folding.mlir
index cb6088c..1271328 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_folding.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/test/tensor_folding.mlir
@@ -206,6 +206,25 @@
// -----
+// CHECK-LABEL: @FoldTensorCloneIntoLoad
+util.func private @FoldTensorCloneIntoLoad(%arg0: !stream.resource<external>, %arg1: index) -> f32 {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ // CHECK-NOT: stream.tensor.clone
+ %0 = stream.tensor.clone %arg0 : tensor<4x4xf32> in !stream.resource<external>{%arg1} -> tensor<4x4xf32> in !stream.resource<*>{%arg1}
+ %1 = stream.async.transfer %0 : !stream.resource<*>{%arg1} -> !stream.resource<staging>{%arg1}
+ // CHECK: %[[SIZE:.+]] = stream.tensor.sizeof tensor<1x1xf32> : index
+ // CHECK: %[[SLICE:.+]] = stream.tensor.slice %arg0[%c0, %c1 for %c1, %c1] : tensor<4x4xf32> in !stream.resource<external>{%arg1} -> tensor<1x1xf32> in !stream.resource<external>{%[[SIZE]]}
+ // CHECK: %[[STAGED:.+]] = stream.async.transfer %[[SLICE]] : !stream.resource<external>{%[[SIZE]]} -> !stream.resource<staging>{%[[SIZE]]}
+ // CHECK: %[[VALUE:.+]] = stream.tensor.load %[[STAGED]][%c0, %c0] : tensor<1x1xf32> in !stream.resource<staging>{%[[SIZE]]} -> f32
+ %2 = stream.tensor.load %1[%c0, %c1] : tensor<4x4xf32> in !stream.resource<staging>{%arg1} -> f32
+ // CHECK: util.return %[[VALUE]]
+ util.return %2 : f32
+}
+
+
+// -----
+
// CHECK-LABEL: @ElideUnneededTensorClones
util.func private @ElideUnneededTensorClones(%arg0: !stream.resource<*>, %arg1: index) -> f32 {
%c0 = arith.constant 0 : index
