[Codegen] Support dynamic offsets in collapse_shape fusion to interface stores (#22800)
Previously, `FoldCollapseShapeIntoInterfaceTensorStore` only handled
static offsets when fusing `tensor.collapse_shape` into dispatch tensor
stores. This change extends the pattern to also accept dynamic offsets
when:
1. The offset comes from an `affine.apply` operation
2. The affine expression is provably divisible by the inner dimension
product
This also fixes an issue where offset computation was done at the
subspan insertion point, which could fail when the offset value is
defined in a nested region (e.g., inside a loop). The offset
affine.apply is now created at the store's insertion point where
dominance is guaranteed.
diff --git a/compiler/src/iree/compiler/Codegen/Common/ReshapePatterns.cpp b/compiler/src/iree/compiler/Codegen/Common/ReshapePatterns.cpp
index 552639f..d53084d 100644
--- a/compiler/src/iree/compiler/Codegen/Common/ReshapePatterns.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/ReshapePatterns.cpp
@@ -600,6 +600,16 @@
/// offsets = [%x * 286, 0, 0, 0], sizes = [3, 3, 1, 96]
/// strides = [1, 1, 1, 1] : tensor<3x3x1x96xf32> ->
/// !iree_tensor_ext.dispatch.tensor<writeonly:tensor<9x3x1x96xf32>>
+///
+/// Dynamic offsets are supported if they come from an affine.apply operation
+/// where the result is provably divisible by the inner dimension product. For
+/// example, if the inner dimensions are [4, 128] (product = 512), an offset of
+/// `affine_map<()[s0] -> (s0 * 512)>` can be transformed to `s0`.
+///
+/// The divisibility check handles Add and Mul expressions recursively:
+/// - For Mul: divisible if either operand is divisible (pessimistic)
+/// - For Add: divisible if both operands are divisible
+/// - For constants: divisible if value % innerDimProduct == 0
struct FoldCollapseShapeIntoInterfaceTensorStore
: OpRewritePattern<IREE::TensorExt::DispatchTensorStoreOp> {
using OpRewritePattern<
@@ -648,35 +658,6 @@
continue;
}
- // If the offset is static and a multiple of the inner dimensions'
- // product, we can divide it by the inner product and add it to the first
- // dimension of the expanded store group. If it's not a multiple, we can't
- // be sure that we can decompose in a way that the `0 <= offset &&
- // offset+size <= dim_size` property holds, so we abort the folding. Same
- // for a dynamic offset.
- std::optional<int64_t> maybeStaticOffset = getConstantIntValue(offset);
- if (!maybeStaticOffset.has_value()) {
- return rewriter.notifyMatchFailure(
- storeOp, "has dynamic offset in dimension to be expanded");
- }
- int64_t staticOffset = maybeStaticOffset.value();
- if (staticOffset != 0) {
- int64_t innerDimSize = 1;
- for (auto i : llvm::drop_begin(group)) {
- if (ShapedType::isDynamic(reshapeSrcShape[i])) {
- return rewriter.notifyMatchFailure(
- storeOp, "has a static offset, but the inner dimension product "
- "of the reshape source is dynamic");
- }
- innerDimSize *= reshapeSrcShape[i];
- }
- if (staticOffset % innerDimSize != 0) {
- return rewriter.notifyMatchFailure(
- storeOp, "has a static offset that is not a multiple of the "
- "static inner dimension product");
- }
- }
-
// Check that each reassociation group contains at most a single dynamic
// dimension.
unsigned numDynamicDims =
@@ -688,9 +669,87 @@
collapseShape,
"got multiple dynamic dimensions in group: " + std::to_string(i));
}
- // In case of a single dynamic dimension, check that the subspan size is
- // dynamic as well.
- if (numDynamicDims == 1) {
+
+ // Compute the inner dimension product. This requires all inner dims to
+ // be static for the offset divisibility and subspan size checks below.
+ int64_t innerDimSize = 1;
+ bool hasStaticInnerDims = true;
+ for (int64_t j : llvm::drop_begin(group)) {
+ if (ShapedType::isDynamic(reshapeSrcShape[j])) {
+ hasStaticInnerDims = false;
+ break;
+ }
+ innerDimSize *= reshapeSrcShape[j];
+ }
+
+ // Check that the offset is divisible by the inner dimension product.
+ std::optional<int64_t> maybeStaticOffset = getConstantIntValue(offset);
+ if (maybeStaticOffset.has_value()) {
+ // Static offset case: verify it's a multiple of the inner dim product.
+ int64_t staticOffset = maybeStaticOffset.value();
+ if (staticOffset != 0) {
+ if (!hasStaticInnerDims) {
+ return rewriter.notifyMatchFailure(
+ storeOp, "has non-zero static offset but dynamic inner dims");
+ }
+ if (staticOffset % innerDimSize != 0) {
+ return rewriter.notifyMatchFailure(
+ storeOp, "has a static offset that is not a multiple of the "
+ "inner dimension product");
+ }
+ }
+ } else {
+ // Dynamic offset case: check if it comes from an affine.apply where
+ // the inner dim product is a multiplicand.
+ if (!hasStaticInnerDims) {
+ return rewriter.notifyMatchFailure(
+ storeOp, "has dynamic offset with dynamic inner dimensions");
+ }
+
+ Value offsetValue = cast<Value>(offset);
+ auto applyOp = offsetValue.getDefiningOp<affine::AffineApplyOp>();
+ if (!applyOp) {
+ return rewriter.notifyMatchFailure(
+ storeOp, "has dynamic offset in dimension to be expanded that is "
+ "not from affine.apply");
+ }
+
+ AffineMap map = applyOp.getAffineMap();
+ AffineExpr expr = map.getResult(0);
+
+ // Helper to check if an affine expression is provably divisible by
+ // innerDimSize.
+ std::function<bool(AffineExpr)> isDivisible =
+ [&](AffineExpr e) -> bool {
+ if (auto constExpr = dyn_cast<AffineConstantExpr>(e)) {
+ return constExpr.getValue() % innerDimSize == 0;
+ }
+ if (auto binaryExpr = dyn_cast<AffineBinaryOpExpr>(e)) {
+ if (binaryExpr.getKind() == AffineExprKind::Mul) {
+ // For multiplication, check if either operand makes it divisible.
+ return isDivisible(binaryExpr.getLHS()) ||
+ isDivisible(binaryExpr.getRHS());
+ }
+ if (binaryExpr.getKind() == AffineExprKind::Add) {
+ // For addition, both operands must be divisible.
+ return isDivisible(binaryExpr.getLHS()) &&
+ isDivisible(binaryExpr.getRHS());
+ }
+ }
+ return false;
+ };
+
+ if (!isDivisible(expr)) {
+ return rewriter.notifyMatchFailure(
+ storeOp, "dynamic offset from affine.apply is not provably a "
+ "multiple of inner dimension product");
+ }
+ }
+
+ // In case of a dynamic dimension, check that the subspan size is dynamic
+ // as well. Dynamic -> static is unlikely and eases the implementation
+ // below.
+ if (numDynamicDims >= 1) {
if (ShapedType::isStatic(subspanSize)) {
return rewriter.notifyMatchFailure(
subspanOp, "collapsing and storing with dynamic dimension into a "
@@ -699,11 +758,7 @@
continue;
}
- // Handle static sizes.
- int64_t innerDimSize = 1;
- for (auto i : llvm::drop_begin(group)) {
- innerDimSize *= reshapeSrcShape[i];
- }
+ // Handle static sizes - use innerDimSize computed above.
if (subspanSize % innerDimSize != 0) {
return rewriter.notifyMatchFailure(
storeOp, "subspan type indivisible by expanded shape");
@@ -717,9 +772,17 @@
rewriter.setInsertionPoint(subspanOp);
Location loc = collapseShape.getLoc();
SmallVector<int64_t> expandedSubspanShape;
- SmallVector<OpFoldResult> expandedOffsets;
SmallVector<OpFoldResult> expandedSizes;
SmallVector<Value> newSubspanDynamicDims;
+ // Store offset info for deferred computation. The offset value may be
+ // defined in a nested region that doesn't dominate the subspan op, so we
+ // defer creating affine.apply ops until after moving the insertion point.
+ struct OffsetInfo {
+ OpFoldResult offset;
+ int64_t innerDimSize;
+ bool needsComputation;
+ };
+ SmallVector<OffsetInfo> offsetInfos;
OpFoldResult zero = rewriter.getIndexAttr(0);
size_t dynIndex = 0;
for (auto [subspanSize, group, offset, size] : llvm::zip_equal(
@@ -727,7 +790,7 @@
storeOp.getMixedOffsets(), storeOp.getMixedSizes())) {
if (group.size() == 1) {
expandedSizes.push_back(size);
- expandedOffsets.push_back(offset);
+ offsetInfos.push_back({offset, 1, false});
expandedSubspanShape.push_back(subspanSize);
if (ShapedType::isDynamic(subspanSize)) {
newSubspanDynamicDims.push_back(dynamicDims[dynIndex++]);
@@ -773,15 +836,13 @@
expandedSubspanShape.push_back(subspanSize / innerDimSize);
}
- // Earlier it's guaranteed that the offset is divisible by the inner
- // dimension product and this is the new first offset of the expanded
- // store. Accordingly, the other offsets are set to zero.
- OpFoldResult innerDimSizeAttr = rewriter.getIndexAttr(innerDimSize);
- expandedOffsets.push_back(affine::makeComposedFoldedAffineApply(
- rewriter, loc, div, {offset, innerDimSizeAttr}));
+ // Store offset info for deferred computation. We can't create the
+ // affine.apply here because the offset value may not dominate the
+ // current insertion point.
+ offsetInfos.push_back({offset, innerDimSize, true});
for (int64_t reshapeSrcSize : llvm::drop_begin(innerDimSizes)) {
- expandedOffsets.push_back(zero);
+ offsetInfos.push_back({zero, 1, false});
expandedSubspanShape.push_back(reshapeSrcSize);
if (ShapedType::isDynamic(reshapeSrcSize)) {
OpFoldResult totalSizeAttr = rewriter.getIndexAttr(totalSize);
@@ -820,6 +881,21 @@
dispatchIndexOpFoldResults(expandedSizes, expandedDynamicDims,
expandedStaticDims);
rewriter.setInsertionPoint(storeOp);
+
+ // Now compute the expanded offsets at the store's insertion point where the
+ // offset values are guaranteed to dominate.
+ SmallVector<OpFoldResult> expandedOffsets;
+ for (const OffsetInfo &info : offsetInfos) {
+ if (info.needsComputation) {
+ OpFoldResult innerDimSizeAttr =
+ rewriter.getIndexAttr(info.innerDimSize);
+ expandedOffsets.push_back(affine::makeComposedFoldedAffineApply(
+ rewriter, loc, div, {info.offset, innerDimSizeAttr}));
+ } else {
+ expandedOffsets.push_back(info.offset);
+ }
+ }
+
rewriter.replaceOpWithNewOp<IREE::TensorExt::DispatchTensorStoreOp>(
storeOp, collapseShape.getSrc(), newSubspanOp, expandedDynamicDims,
expandedOffsets, expandedSizes, expandedStrides);
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/fold_reshape_into_interface_tensor.mlir b/compiler/src/iree/compiler/Codegen/Common/test/fold_reshape_into_interface_tensor.mlir
index c03effb..0c5d084 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/fold_reshape_into_interface_tensor.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/fold_reshape_into_interface_tensor.mlir
@@ -24,7 +24,7 @@
// -----
-#pipeline_layout = #hal.pipeline.layout<constants = 2, bindings = [
+#pipeline_layout = #hal.pipeline.layout<constants = 1, bindings = [
#hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">], flags = Indirect>
func.func @fold_expand_into_loads_dynamic() -> tensor<2x?x16x32xf32> {
%c0 = arith.constant 0 : index
@@ -373,7 +373,7 @@
// -----
-#pipeline_layout = #hal.pipeline.layout<constants = 2, bindings = [
+#pipeline_layout = #hal.pipeline.layout<constants = 1, bindings = [
#hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>
func.func @fold_expand_and_collapse(%arg0 : tensor<1x?x1x8xi32>) {
%c128 = arith.constant 128 : index
@@ -399,3 +399,139 @@
// CHECK: iree_tensor_ext.dispatch.tensor.store %{{.+}}, %[[SUBSPAN2]]
// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [1, %[[SHAPE]], 1, 8], strides = [1, 1, 1, 1]
// CHECK-SAME: !iree_tensor_ext.dispatch.tensor<readwrite:tensor<1x?x1x8xi32>>{%[[SHAPE]]}
+
+// -----
+
+// Test dynamic offset from affine.apply where constant is divisible by inner dim product.
+#pipeline_layout = #hal.pipeline.layout<constants = 0, bindings = [
+ #hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>
+func.func @fold_collapse_into_stores_dynamic_offset_divisible(%arg0 : tensor<4x8x4x128xf32>, %idx : index) {
+ %c0 = arith.constant 0 : index
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+ flags("ReadOnly|Indirect") : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ %2 = tensor.collapse_shape %arg0 [[0], [1, 2, 3]] : tensor<4x8x4x128xf32> into tensor<4x4096xf32>
+ // Offset 512 * idx is divisible by inner dim product (4*128=512)
+ %offset = affine.apply affine_map<()[s0] -> (s0 * 512)>()[%idx]
+ iree_tensor_ext.dispatch.tensor.store %2, %1, offsets = [0, %offset], sizes = [4, 4096], strides = [1, 1]
+ : tensor<4x4096xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ return
+}
+// CHECK-LABEL: func @fold_collapse_into_stores_dynamic_offset_divisible(
+// CHECK-SAME: %[[IDX:[A-Za-z0-9]+]]: index
+// CHECK: %[[SUBSPAN:.+]] = hal.interface.binding.subspan
+// CHECK-SAME: !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x40x4x128xf32>>
+// CHECK-NOT: tensor.collapse_shape
+// CHECK: iree_tensor_ext.dispatch.tensor.store %{{.+}}, %[[SUBSPAN]]
+// CHECK-SAME: offsets = [0, %[[IDX]], 0, 0], sizes = [4, 8, 4, 128], strides = [1, 1, 1, 1]
+// CHECK-SAME: !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x40x4x128xf32>>
+
+// -----
+
+// Test that dynamic offset without affine.apply is rejected.
+#pipeline_layout = #hal.pipeline.layout<constants = 0, bindings = [
+ #hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>
+func.func @no_fold_collapse_dynamic_offset_no_affine_apply(%arg0 : tensor<4x8x4x128xf32>, %offset : index) {
+ %c0 = arith.constant 0 : index
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+ flags("ReadOnly|Indirect") : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ %2 = tensor.collapse_shape %arg0 [[0], [1, 2, 3]] : tensor<4x8x4x128xf32> into tensor<4x4096xf32>
+ // Raw dynamic offset without affine.apply should be rejected
+ iree_tensor_ext.dispatch.tensor.store %2, %1, offsets = [0, %offset], sizes = [4, 4096], strides = [1, 1]
+ : tensor<4x4096xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ return
+}
+// CHECK-LABEL: func @no_fold_collapse_dynamic_offset_no_affine_apply
+// CHECK: tensor.collapse_shape
+
+// -----
+
+// Test that dynamic offset not divisible by inner dim is rejected.
+#pipeline_layout = #hal.pipeline.layout<constants = 0, bindings = [
+ #hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>
+func.func @no_fold_collapse_dynamic_offset_not_divisible(%arg0 : tensor<4x8x4x128xf32>, %idx : index) {
+ %c0 = arith.constant 0 : index
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+ flags("ReadOnly|Indirect") : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ %2 = tensor.collapse_shape %arg0 [[0], [1, 2, 3]] : tensor<4x8x4x128xf32> into tensor<4x4096xf32>
+ // Offset 100 * idx is not divisible by inner dim product (512)
+ %offset = affine.apply affine_map<()[s0] -> (s0 * 100)>()[%idx]
+ iree_tensor_ext.dispatch.tensor.store %2, %1, offsets = [0, %offset], sizes = [4, 4096], strides = [1, 1]
+ : tensor<4x4096xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ return
+}
+// CHECK-LABEL: func @no_fold_collapse_dynamic_offset_not_divisible
+// CHECK: tensor.collapse_shape
+
+// -----
+
+// Test dynamic offset from affine.apply with addition where both terms are divisible.
+#pipeline_layout = #hal.pipeline.layout<constants = 0, bindings = [
+ #hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>
+func.func @fold_collapse_dynamic_offset_addition_divisible(%arg0 : tensor<4x8x4x128xf32>, %idx : index) {
+ %c0 = arith.constant 0 : index
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+ flags("ReadOnly|Indirect") : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ %2 = tensor.collapse_shape %arg0 [[0], [1, 2, 3]] : tensor<4x8x4x128xf32> into tensor<4x4096xf32>
+ // Offset (512 * idx + 1024) is divisible by inner dim product (512)
+ %offset = affine.apply affine_map<()[s0] -> (s0 * 512 + 1024)>()[%idx]
+ iree_tensor_ext.dispatch.tensor.store %2, %1, offsets = [0, %offset], sizes = [4, 4096], strides = [1, 1]
+ : tensor<4x4096xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ return
+}
+// CHECK-LABEL: func @fold_collapse_dynamic_offset_addition_divisible(
+// CHECK-SAME: %[[IDX:[A-Za-z0-9]+]]: index
+// CHECK: %[[SUBSPAN:.+]] = hal.interface.binding.subspan
+// CHECK-SAME: !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x40x4x128xf32>>
+// CHECK-NOT: tensor.collapse_shape
+// CHECK: %[[OFFSET:.+]] = affine.apply affine_map<()[s0] -> (s0 + 2)>()[%[[IDX]]]
+// CHECK: iree_tensor_ext.dispatch.tensor.store %{{.+}}, %[[SUBSPAN]]
+// CHECK-SAME: offsets = [0, %[[OFFSET]], 0, 0], sizes = [4, 8, 4, 128], strides = [1, 1, 1, 1]
+
+// -----
+
+// Test dynamic offset from affine.apply with addition where addend is not divisible.
+#pipeline_layout = #hal.pipeline.layout<constants = 0, bindings = [
+ #hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>
+func.func @no_fold_collapse_dynamic_offset_addition_not_divisible(%arg0 : tensor<4x8x4x128xf32>, %idx : index) {
+ %c0 = arith.constant 0 : index
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+ flags("ReadOnly|Indirect") : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ %2 = tensor.collapse_shape %arg0 [[0], [1, 2, 3]] : tensor<4x8x4x128xf32> into tensor<4x4096xf32>
+ // Offset (512 * idx + 100) is not divisible by inner dim product (512) because 100 is not divisible
+ %offset = affine.apply affine_map<()[s0] -> (s0 * 512 + 100)>()[%idx]
+ iree_tensor_ext.dispatch.tensor.store %2, %1, offsets = [0, %offset], sizes = [4, 4096], strides = [1, 1]
+ : tensor<4x4096xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ return
+}
+// CHECK-LABEL: func @no_fold_collapse_dynamic_offset_addition_not_divisible
+// CHECK: tensor.collapse_shape
+
+// -----
+
+// Test that dynamic offset defined inside a nested region works.
+#pipeline_layout = #hal.pipeline.layout<constants = 1, bindings = [
+ #hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>
+func.func @fold_collapse_dynamic_offset_in_nested_region(%arg0 : tensor<4x8x4x128xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+ flags("ReadOnly|Indirect") : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ // The offset is computed inside the forall using the induction variable,
+ // forcing expanded offset computation to go inside the region.
+ scf.forall (%iv) in (%c4) {
+ %2 = tensor.collapse_shape %arg0 [[0], [1, 2, 3]] : tensor<4x8x4x128xf32> into tensor<4x4096xf32>
+ %offset = affine.apply affine_map<(d0) -> (d0 * 512)>(%iv)
+ iree_tensor_ext.dispatch.tensor.store %2, %1, offsets = [0, %offset], sizes = [4, 4096], strides = [1, 1]
+ : tensor<4x4096xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x20480xf32>>
+ }
+ return
+}
+// CHECK-LABEL: func @fold_collapse_dynamic_offset_in_nested_region(
+// CHECK: %[[SUBSPAN:.+]] = hal.interface.binding.subspan
+// CHECK-SAME: !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x40x4x128xf32>>
+// CHECK: scf.forall (%[[IV:.+]]) in
+// CHECK-NOT: tensor.collapse_shape
+// CHECK: iree_tensor_ext.dispatch.tensor.store %{{.+}}, %[[SUBSPAN]]
+// CHECK-SAME: offsets = [0, %[[IV]], 0, 0], sizes = [4, 8, 4, 128], strides = [1, 1, 1, 1]
+// CHECK-SAME: !iree_tensor_ext.dispatch.tensor<writeonly:tensor<4x40x4x128xf32>>
