[LinalgExt] Implement direct vectorization for im2col op (#23855)
Implements direct vectorization for im2col, and completes the support
for padding on the im2col op. The padding attributes are used to compute
the read mask when vectorizing. In the old path, we would have separate
padding on the input and the result of the im2col, and we try to compose
those pads into a single masked read. This is fragile and difficult for
cases where the im2col result dims don't map well to the input dims.
With this direct vectorization approach, we can compute the mask based
on the input and result padding simultaneously. This will make
flattening of the spatial dimensions of convolutions possible.
### Performance results: ###
- Run 1:
https://github.com/nod-ai/amd-shark-ai-reports/tree/main/boo/boo-custom-runs/2026-03-27_04-36_d1b822f45ac693a8593232a7d3fc5d67b1087f7e/comparison
- Run 2:
https://github.com/nod-ai/amd-shark-ai-reports/tree/main/boo/boo-custom-runs/2026-03-27_20-34_cf0d758bb199713b96a84a67245bc8b24ba7b74a/comparison
These runs were taken on different commits, but they are functionally
the same (just some cleanup differences). I am only able to reproduce 3
of the regressions locally, and most of the improvers (~35 of them with
10-40% speedup) are real. There seems to have been some noise in the
runs, but overall there is a good perf improvement.
ci-extra: test_torch
---------
Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/GenericVectorization.cpp b/compiler/src/iree/compiler/Codegen/Common/GenericVectorization.cpp
index d3cd248..7a6e0a7 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GenericVectorization.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/GenericVectorization.cpp
@@ -10,6 +10,7 @@
#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
#include "iree/compiler/Codegen/Interfaces/VectorizableOpInterface.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
+#include "iree/compiler/Dialect/LinalgExt/IR/Im2colUtils.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "llvm/Support/DebugLog.h"
#include "mlir/Dialect/Affine/LoopUtils.h"
@@ -134,6 +135,11 @@
vectorSizes = result->vectorSizes;
}
})
+ .Case([&](IREE::LinalgExt::Im2colOp im2colOp) {
+ OpBuilder builder(op);
+ vectorSizes =
+ IREE::LinalgExt::computeIm2colVectorTileSizes(builder, im2colOp);
+ })
.Default([&](Operation *) {});
if (vectorSizes) {
diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeVectorTileSizes.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeVectorTileSizes.cpp
index dc99cbe..d25d795 100644
--- a/compiler/src/iree/compiler/Codegen/Common/MaterializeVectorTileSizes.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/MaterializeVectorTileSizes.cpp
@@ -8,6 +8,8 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
+#include "iree/compiler/Dialect/LinalgExt/IR/Im2colUtils.h"
+#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "llvm/Support/DebugLog.h"
#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
@@ -406,6 +408,20 @@
return success();
}
+ // Im2col ops: compute tile sizes locally and propagate to result.
+ // Im2col's output dimensions are the iteration domain (identity map),
+ // so tile sizes go directly on the result lattice.
+ if (auto im2colOp = dyn_cast<IREE::LinalgExt::Im2colOp>(op)) {
+ OpBuilder builder(op);
+ std::optional<SmallVector<int64_t>> maybeTileSizes =
+ IREE::LinalgExt::computeIm2colVectorTileSizes(builder, im2colOp);
+ if (maybeTileSizes) {
+ TileSizes tileSizes(*maybeTileSizes);
+ propagateIfChanged(results[0], results[0]->join(tileSizes));
+ }
+ return success();
+ }
+
// Pack ops: map source tile sizes to dest space.
if (auto packOp = dyn_cast<linalg::PackOp>(op)) {
if (ShapedType::isDynamicShape(packOp.getStaticInnerTiles())) {
@@ -536,6 +552,18 @@
return success();
}
+ // Im2col ops: propagate result tile sizes back to the output operand.
+ // Im2col's output dimensions are the iteration domain (identity map),
+ // so tile sizes propagate directly.
+ if (auto im2colOp = dyn_cast<IREE::LinalgExt::Im2colOp>(op)) {
+ TileSizes tileSizes = getTileSizesFor(im2colOp.getResult(0), results[0]);
+ // Propagate to the output (DPS init) operand.
+ unsigned outputIdx = im2colOp.getDpsInitsMutable()[0].getOperandNumber();
+ TileSizeLattice *outputLattice = operands[outputIdx];
+ propagateIfChanged(outputLattice, outputLattice->meet(tileSizes));
+ return success();
+ }
+
// Pack ops: result tile sizes → source tile sizes (backward).
if (auto packOp = dyn_cast<linalg::PackOp>(op)) {
if (ShapedType::isDynamicShape(packOp.getStaticInnerTiles())) {
@@ -609,6 +637,16 @@
return iterTileSizes;
}
+/// Get tile sizes for an im2col op from its result lattice. Im2col's output
+/// dimensions are the iteration domain, so the result lattice directly holds
+/// the iteration-space tile sizes.
+static TileSizes getIm2colTileSizes(IREE::LinalgExt::Im2colOp im2colOp,
+ const DataFlowSolver &solver) {
+ Value result = im2colOp.getResult(0);
+ const TileSizeLattice *lattice = solver.lookupState<TileSizeLattice>(result);
+ return getTileSizesFor(result, lattice);
+}
+
//===----------------------------------------------------------------------===//
// MaterializeVectorTileSizesPass
//===----------------------------------------------------------------------===//
@@ -684,6 +722,10 @@
getIterationSpaceTileSizes(op, numLoops, indexingMaps, solver);
return WalkResult(materialize(op, tileSizes));
}
+ if (auto im2colOp = dyn_cast<IREE::LinalgExt::Im2colOp>(op)) {
+ TileSizes tileSizes = getIm2colTileSizes(im2colOp, solver);
+ return WalkResult(materialize(op, tileSizes));
+ }
return WalkResult::advance();
});
if (result.wasInterrupted()) {
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization_masked_inferred.mlir b/compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization_masked_inferred.mlir
index 1badf10..6cfecc6 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization_masked_inferred.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization_masked_inferred.mlir
@@ -915,3 +915,247 @@
// CHECK: %[[WRITE_MASK:.+]] = vector.create_mask {{.*}} : vector<2x5x4xi1>
// CHECK: vector.transfer_write %[[MMA]], %arg2{{.*}}, %[[WRITE_MASK]] {in_bounds = [false, false, true]}
// CHECK-SAME: : vector<2x5x4xf32>, tensor<?x?x4xf32>
+
+// -----
+
+// Tests for im2col op vectorization via VectorizableOpInterface.
+
+// Standard NHWC layout, K tile size (4) divides innermost input dim C (640).
+// Vectorizes along K (output dim 2) with vector width 4.
+// Non-vectorized dims: batch (2) x M (2) = 4 iterations.
+#im2col_map_k = affine_map<(d0) -> (d0 * 4)>
+func.func @im2col_vectorize_nhwc(
+ %input: tensor<2x34x34x640xf32>, %m_off: index, %k: index
+) -> tensor<2x2x4xf32> {
+ %0 = tensor.empty() : tensor<2x2x4xf32>
+ %k_off = affine.apply #im2col_map_k(%k)
+ %1 = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, %m_off, %k_off] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2] output_perm = [0, 1, 2]
+ ins(%input : tensor<2x34x34x640xf32>)
+ outs(%0 : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
+ return %1 : tensor<2x2x4xf32>
+}
+// CHECK-LABEL: func.func @im2col_vectorize_nhwc
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<2x34x34x640xf32>
+// CHECK-DAG: %[[POISON:.+]] = ub.poison : f32
+// CHECK-NOT: iree_linalg_ext.im2col
+// CHECK: %[[R0:.+]] = vector.transfer_read %[[INPUT]]{{.*}}, %[[POISON]] {in_bounds = [true]} : tensor<2x34x34x640xf32>, vector<4xf32>
+// CHECK: vector.transfer_write %[[R0]], {{.*}} : vector<4xf32>, tensor<2x2x4xf32>
+// CHECK: %[[R1:.+]] = vector.transfer_read %[[INPUT]]{{.*}}, %[[POISON]] {in_bounds = [true]} : tensor<2x34x34x640xf32>, vector<4xf32>
+// CHECK: vector.transfer_write %[[R1]], {{.*}} : vector<4xf32>, tensor<2x2x4xf32>
+// CHECK: %[[R2:.+]] = vector.transfer_read %[[INPUT]]{{.*}}, %[[POISON]] {in_bounds = [true]} : tensor<2x34x34x640xf32>, vector<4xf32>
+// CHECK: vector.transfer_write %[[R2]], {{.*}} : vector<4xf32>, tensor<2x2x4xf32>
+// CHECK: %[[R3:.+]] = vector.transfer_read %[[INPUT]]{{.*}}, %[[POISON]] {in_bounds = [true]} : tensor<2x34x34x640xf32>, vector<4xf32>
+// CHECK: %[[FINAL:.+]] = vector.transfer_write %[[R3]], {{.*}} : vector<4xf32>, tensor<2x2x4xf32>
+// CHECK: return %[[FINAL]] : tensor<2x2x4xf32>
+
+// -----
+
+// Dynamic output shape: vectorization pattern should not match.
+func.func @im2col_no_vectorize_dynamic(
+ %input: tensor<2x34x34x640xf32>, %m_size: index, %m_off: index, %k: index
+) -> tensor<2x?x4xf32> {
+ %0 = tensor.empty(%m_size) : tensor<2x?x4xf32>
+ %k_off = affine.apply affine_map<(d0) -> (d0 * 4)>(%k)
+ %1 = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, %m_off, %k_off] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2] output_perm = [0, 1, 2]
+ ins(%input : tensor<2x34x34x640xf32>)
+ outs(%0 : tensor<2x?x4xf32>) -> tensor<2x?x4xf32>
+ return %1 : tensor<2x?x4xf32>
+}
+// CHECK-LABEL: func.func @im2col_no_vectorize_dynamic
+// CHECK: iree_linalg_ext.im2col
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+
+// -----
+
+// Source padding (conv padding folded into im2col). NHWC layout.
+// Vectorizes along K with masked transfer_read.
+#im2col_map_k_pad = affine_map<(d0) -> (d0 * 4)>
+func.func @im2col_vectorize_source_padding(
+ %input: tensor<2x34x34x640xf32>, %m_off: index, %k: index
+) -> tensor<2x2x4xf32> {
+ %cst = arith.constant 0.0 : f32
+ %0 = tensor.empty() : tensor<2x2x4xf32>
+ %k_off = affine.apply #im2col_map_k_pad(%k)
+ %1 = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, %m_off, %k_off] output_sizes = [[2], [34, 34], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2] output_perm = [0, 1, 2]
+ input_pad_low = [0, 1, 1, 0] input_pad_high = [0, 1, 1, 0]
+ pad_value(%cst : f32)
+ ins(%input : tensor<2x34x34x640xf32>)
+ outs(%0 : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
+ return %1 : tensor<2x2x4xf32>
+}
+// CHECK-LABEL: func.func @im2col_vectorize_source_padding
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<2x34x34x640xf32>
+// CHECK-DAG: %[[PAD:.+]] = arith.constant 0.0{{.*}} : f32
+// CHECK-NOT: iree_linalg_ext.im2col
+// CHECK: %[[MASK0:.+]] = vector.create_mask {{.*}} : vector<4xi1>
+// CHECK: vector.transfer_read %[[INPUT]]{{.*}}, %[[PAD]], %[[MASK0]] {{.*}} : tensor<2x34x34x640xf32>, vector<4xf32>
+// CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<2x2x4xf32>
+// CHECK: %[[MASK1:.+]] = vector.create_mask {{.*}} : vector<4xi1>
+// CHECK: vector.transfer_read %[[INPUT]]{{.*}}, %[[PAD]], %[[MASK1]] {{.*}} : tensor<2x34x34x640xf32>, vector<4xf32>
+// CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<2x2x4xf32>
+// CHECK: %[[MASK2:.+]] = vector.create_mask {{.*}} : vector<4xi1>
+// CHECK: vector.transfer_read %[[INPUT]]{{.*}}, %[[PAD]], %[[MASK2]] {{.*}} : tensor<2x34x34x640xf32>, vector<4xf32>
+// CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<2x2x4xf32>
+// CHECK: %[[MASK3:.+]] = vector.create_mask {{.*}} : vector<4xi1>
+// CHECK: vector.transfer_read %[[INPUT]]{{.*}}, %[[PAD]], %[[MASK3]] {{.*}} : tensor<2x34x34x640xf32>, vector<4xf32>
+// CHECK: %[[FINAL:.+]] = vector.transfer_write {{.*}} : vector<4xf32>, tensor<2x2x4xf32>
+// CHECK: return %[[FINAL]] : tensor<2x2x4xf32>
+
+// -----
+
+// Non-vectorizable due to input_k_perm = [1, 0] making innermost K
+// non-contiguous in input. Falls back to scalar unrolling (vector<1>).
+func.func @im2col_scalar_fallback(
+ %input: tensor<1x3x2xf32>
+) -> tensor<1x2x4xf32> {
+ %0 = tensor.empty() : tensor<1x2x4xf32>
+ %1 = iree_linalg_ext.im2col strides = [1] dilations = [1] kernel_size = [2]
+ offsets = [0, 0, 0] output_sizes = [[1], [2], [2, 2]]
+ batch_pos = [0] m_pos = [1] k_pos = [2]
+ input_k_perm = [1, 0] output_perm = [0, 1, 2]
+ ins(%input : tensor<1x3x2xf32>)
+ outs(%0 : tensor<1x2x4xf32>) -> tensor<1x2x4xf32>
+ return %1 : tensor<1x2x4xf32>
+}
+// CHECK-LABEL: func.func @im2col_scalar_fallback
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<1x3x2xf32>
+// CHECK-NOT: iree_linalg_ext.im2col
+// CHECK: vector.transfer_read %[[INPUT]]{{.*}} : tensor<1x3x2xf32>, vector<1xf32>
+// CHECK: vector.transfer_write {{.*}} : vector<1xf32>, tensor<1x2x4xf32>
+
+// -----
+
+// High-side input padding on the vectorized input dimension (channels).
+// Verifies: masked vector transfer_read with pad_value, im2col fully lowered.
+func.func @im2col_vectorize_channel_pad_high(
+ %input: tensor<59x91x16x56xbf16>, %output: tensor<1x1x1x8xbf16>,
+ %off0: index
+) -> tensor<1x1x1x8xbf16> {
+ %cst = arith.constant 0.000000e+00 : bf16
+ %c5 = arith.constant 5 : index
+ %c3 = arith.constant 3 : index
+ %c100 = arith.constant 100 : index
+ %result = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [59, 91]
+ offsets = [%off0, %c3, %c5, %c100]
+ output_sizes = [[64], [16], [3, 3], [59, 91]]
+ batch_pos = [3, 2] m_pos = [0, 1] k_pos = []
+ input_k_perm = [0, 1] output_perm = [2, 3, 1, 0]
+ input_pad_low = [1, 1, 0, 0] input_pad_high = [1, 1, 0, 8]
+ pad_value(%cst : bf16)
+ ins(%input : tensor<59x91x16x56xbf16>)
+ outs(%output : tensor<1x1x1x8xbf16>) -> tensor<1x1x1x8xbf16>
+ return %result : tensor<1x1x1x8xbf16>
+}
+// CHECK-LABEL: func.func @im2col_vectorize_channel_pad_high
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<59x91x16x56xbf16>
+// CHECK-DAG: %[[PAD:.+]] = arith.constant 0.0{{.*}} : bf16
+// CHECK-NOT: iree_linalg_ext.im2col
+// CHECK: vector.create_mask {{.*}} : vector<8xi1>
+// CHECK: vector.transfer_read %[[INPUT]]{{.*}}, %[[PAD]], %{{.*}} {in_bounds = [true]} : tensor<59x91x16x56xbf16>, vector<8xbf16>
+// CHECK: %[[FINAL:.+]] = vector.transfer_write {{.*}} : vector<8xbf16>, tensor<1x1x1x8xbf16>
+// CHECK: return %[[FINAL]] : tensor<1x1x1x8xbf16>
+
+// -----
+
+// Low-side input padding on the vectorized input dimension: falls back to
+// scalar unrolling (vector<1>) because chooseDimToVectorize returns nullopt.
+func.func @im2col_scalar_fallback_channel_pad_low(
+ %input: tensor<59x91x16x56xbf16>, %output: tensor<1x1x1x8xbf16>
+) -> tensor<1x1x1x8xbf16> {
+ %cst = arith.constant 0.000000e+00 : bf16
+ %c5 = arith.constant 5 : index
+ %c3 = arith.constant 3 : index
+ %c42 = arith.constant 42 : index
+ %c100 = arith.constant 100 : index
+ %result = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [59, 91]
+ offsets = [%c42, %c3, %c5, %c100]
+ output_sizes = [[64], [16], [3, 3], [59, 91]]
+ batch_pos = [3, 2] m_pos = [0, 1] k_pos = []
+ input_k_perm = [0, 1] output_perm = [2, 3, 1, 0]
+ input_pad_low = [1, 1, 0, 8] input_pad_high = [1, 1, 0, 0]
+ pad_value(%cst : bf16)
+ ins(%input : tensor<59x91x16x56xbf16>)
+ outs(%output : tensor<1x1x1x8xbf16>) -> tensor<1x1x1x8xbf16>
+ return %result : tensor<1x1x1x8xbf16>
+}
+// All offsets are constant and in-bounds, so masks fold away. The im2col is
+// fully lowered to 8 scalar (vector<1>) transfer_read/write pairs.
+// CHECK-LABEL: func.func @im2col_scalar_fallback_channel_pad_low
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<59x91x16x56xbf16>
+// CHECK-NOT: iree_linalg_ext.im2col
+// CHECK: vector.transfer_read %[[INPUT]]{{.*}} {in_bounds = [true]} : tensor<59x91x16x56xbf16>, vector<1xbf16>
+// CHECK: vector.transfer_write {{.*}} : vector<1xbf16>, tensor<1x1x1x8xbf16>
+
+// -----
+
+// Output-only padding (GEMM alignment). Vectorizes along K with masked reads.
+// The output has 16 extra M positions filled with pad_value.
+func.func @im2col_vectorize_output_padding(
+ %input: tensor<2x34x34x640xf32>, %m_off: index, %k: index
+) -> tensor<2x2x4xf32> {
+ %cst = arith.constant 0.0 : f32
+ %0 = tensor.empty() : tensor<2x2x4xf32>
+ %k_off = affine.apply affine_map<(d0) -> (d0 * 4)>(%k)
+ %1 = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, %m_off, %k_off] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2] output_perm = [0, 1, 2]
+ output_pad_low = [0, 0, 0] output_pad_high = [0, 16, 0]
+ pad_value(%cst : f32)
+ ins(%input : tensor<2x34x34x640xf32>)
+ outs(%0 : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
+ return %1 : tensor<2x2x4xf32>
+}
+// Vectorizes along K (dim 2) with vector width 4. The output M-dim padding
+// produces arith.select between the k-dim mask and all-false for each
+// non-vectorized output dim. No input padding, so reads are from the
+// unpadded tensor with clamped indices.
+// CHECK-LABEL: func.func @im2col_vectorize_output_padding
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<2x34x34x640xf32>
+// CHECK-DAG: %[[PAD:.+]] = arith.constant 0.0{{.*}} : f32
+// CHECK-NOT: iree_linalg_ext.im2col
+// CHECK: vector.transfer_read %[[INPUT]]{{.*}}, %[[PAD]], %{{.*}} {in_bounds = [true]} : tensor<2x34x34x640xf32>, vector<4xf32>
+// CHECK: vector.transfer_write {{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<2x2x4xf32>
+
+// -----
+
+// Output low-padding on the vectorized dim: falls back to scalar unrolling
+// because chooseDimToVectorize skips dims with non-zero output_pad_low.
+func.func @im2col_scalar_fallback_output_pad_low(
+ %input: tensor<2x34x34x640xf32>, %m_off: index, %k: index
+) -> tensor<2x2x4xf32> {
+ %cst = arith.constant 0.0 : f32
+ %0 = tensor.empty() : tensor<2x2x4xf32>
+ %k_off = affine.apply affine_map<(d0) -> (d0 * 4)>(%k)
+ %1 = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, %m_off, %k_off] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2] output_perm = [0, 1, 2]
+ output_pad_low = [0, 0, 2] output_pad_high = [0, 0, 0]
+ pad_value(%cst : f32)
+ ins(%input : tensor<2x34x34x640xf32>)
+ outs(%0 : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
+ return %1 : tensor<2x2x4xf32>
+}
+// Scalar fallback: output_pad_low on the K dim (dim 2) prevents vectorization.
+// CHECK-LABEL: func.func @im2col_scalar_fallback_output_pad_low
+// CHECK-NOT: iree_linalg_ext.im2col
+// CHECK: vector.transfer_read {{.*}} : tensor<2x34x34x640xf32>, vector<1xf32>
+// CHECK: vector.transfer_write {{.*}} : vector<1xf32>, tensor<2x2x4xf32>
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_vector_tile_sizes.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_vector_tile_sizes.mlir
index 98f3d33..9e51b37 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/materialize_vector_tile_sizes.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_vector_tile_sizes.mlir
@@ -532,3 +532,80 @@
into %empty_result : tensor<1x?x?x16x16xf32> -> tensor<1x?x?xf32>
return %result : tensor<1x?x?xf32>
}
+
+// -----
+
+// Im2col: basic NHWC vectorization along K (output dim 2). K tile size (4)
+// divides the innermost input dim C (640), so the analysis picks dim 2 as
+// the vectorized dim with full size 4, and tiles the batch and M dims to 1.
+
+#im2col_map_k = affine_map<(d0) -> (d0 * 4)>
+// CHECK-LABEL: @im2col_tile_sizes_nhwc
+func.func @im2col_tile_sizes_nhwc(
+ %input: tensor<2x34x34x640xf32>, %m_off: index, %k: index
+) -> tensor<2x2x4xf32> {
+ %0 = tensor.empty() : tensor<2x2x4xf32>
+ %k_off = affine.apply #im2col_map_k(%k)
+ // CHECK: iree_linalg_ext.im2col
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 1, 4>
+ %1 = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, %m_off, %k_off] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2] output_perm = [0, 1, 2]
+ ins(%input : tensor<2x34x34x640xf32>)
+ outs(%0 : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
+ return %1 : tensor<2x2x4xf32>
+}
+
+// -----
+
+// Im2col: non-vectorizable. input_k_perm = [1, 0] makes the innermost K
+// non-contiguous in the input tensor, so no dimension can be vectorized
+// with a contiguous slice. The analysis should not stamp a tile sizes
+// attribute.
+
+// CHECK-LABEL: @im2col_tile_sizes_non_contiguous
+func.func @im2col_tile_sizes_non_contiguous(
+ %input: tensor<1x3x2xf32>
+) -> tensor<1x2x4xf32> {
+ %0 = tensor.empty() : tensor<1x2x4xf32>
+ // CHECK: iree_linalg_ext.im2col
+ // CHECK-NOT: iree_codegen.vector_tile_sizes
+ %1 = iree_linalg_ext.im2col strides = [1] dilations = [1] kernel_size = [2]
+ offsets = [0, 0, 0] output_sizes = [[1], [2], [2, 2]]
+ batch_pos = [0] m_pos = [1] k_pos = [2]
+ input_k_perm = [1, 0] output_perm = [0, 1, 2]
+ ins(%input : tensor<1x3x2xf32>)
+ outs(%0 : tensor<1x2x4xf32>) -> tensor<1x2x4xf32>
+ return %1 : tensor<1x2x4xf32>
+}
+
+// -----
+
+// Im2col: wider vectorization. Vectorizes along the innermost channel dim
+// with width 8. Non-vectorized spatial dims are tiled to 1.
+
+// CHECK-LABEL: @im2col_tile_sizes_channel_width_8
+func.func @im2col_tile_sizes_channel_width_8(
+ %input: tensor<59x91x16x56xbf16>, %output: tensor<1x1x1x8xbf16>,
+ %off0: index
+) -> tensor<1x1x1x8xbf16> {
+ %cst = arith.constant 0.000000e+00 : bf16
+ %c5 = arith.constant 5 : index
+ %c3 = arith.constant 3 : index
+ %c100 = arith.constant 100 : index
+ // CHECK: iree_linalg_ext.im2col
+ // CHECK-SAME: iree_codegen.vector_tile_sizes = array<i64: 1, 1, 1, 8>
+ %result = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [59, 91]
+ offsets = [%off0, %c3, %c5, %c100]
+ output_sizes = [[64], [16], [3, 3], [59, 91]]
+ batch_pos = [3, 2] m_pos = [0, 1] k_pos = []
+ input_k_perm = [0, 1] output_perm = [2, 3, 1, 0]
+ input_pad_low = [1, 1, 0, 0] input_pad_high = [1, 1, 0, 8]
+ pad_value(%cst : bf16)
+ ins(%input : tensor<59x91x16x56xbf16>)
+ outs(%output : tensor<1x1x1x8xbf16>) -> tensor<1x1x1x8xbf16>
+ return %result : tensor<1x1x1x8xbf16>
+}
diff --git a/compiler/src/iree/compiler/Codegen/Interfaces/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Interfaces/BUILD.bazel
index 5771116..3b73fb8 100644
--- a/compiler/src/iree/compiler/Codegen/Interfaces/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Interfaces/BUILD.bazel
@@ -236,9 +236,14 @@
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect",
"//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
+ "//compiler/src/iree/compiler/Dialect/LinalgExt/Utils",
"//compiler/src/iree/compiler/Utils",
+ "@llvm-project//mlir:AffineDialect",
+ "@llvm-project//mlir:AffineUtils",
"@llvm-project//mlir:Analysis",
"@llvm-project//mlir:ArithDialect",
+ "@llvm-project//mlir:ArithUtils",
+ "@llvm-project//mlir:DialectUtils",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:LinalgDialect",
"@llvm-project//mlir:LinalgStructuredOpsIncGen",
diff --git a/compiler/src/iree/compiler/Codegen/Interfaces/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Interfaces/CMakeLists.txt
index e4e0d8b..7faddda 100644
--- a/compiler/src/iree/compiler/Codegen/Interfaces/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Interfaces/CMakeLists.txt
@@ -171,8 +171,11 @@
"VectorizableOpInterface.cpp"
DEPS
::VectorizableOpInterfaceGen
+ MLIRAffineDialect
+ MLIRAffineUtils
MLIRAnalysis
MLIRArithDialect
+ MLIRArithUtils
MLIRIR
MLIRLinalgDialect
MLIRLinalgStructuredOpsIncGenLib
@@ -185,6 +188,7 @@
iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect
iree::compiler::Dialect::LinalgExt::IR
+ iree::compiler::Dialect::LinalgExt::Utils
iree::compiler::Utils
PUBLIC
)
diff --git a/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp b/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp
index be4e33e..3510537 100644
--- a/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp
+++ b/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp
@@ -10,15 +10,21 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.h"
+#include "iree/compiler/Dialect/LinalgExt/IR/Im2colUtils.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
#include "iree/compiler/Utils/Indexing.h"
#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/IR/AffineMap.h"
@@ -1192,6 +1198,168 @@
registerInterfaceForLinalgOps<OpTy2, More...>(ctx);
}
+//===----------------------------------------------------------------------===//
+// Im2col Vectorization
+//===----------------------------------------------------------------------===//
+
+/// Compute the padding mask for im2col vectorization.
+/// Gets the valid size from computeIm2colValidSize and converts it to a
+/// vector mask. The mask guards against out-of-bounds accesses, so read
+/// indices do not need clamping.
+static Value computeIm2colPaddingMask(
+ OpBuilder &b, Location loc, IREE::LinalgExt::Im2colOp im2colOp,
+ const IREE::LinalgExt::Im2colSourceIndices &srcIndices, int64_t vecWidth,
+ ArrayRef<Value> outputIVs, std::optional<int64_t> vecOutputDim) {
+ auto vecI1Type = VectorType::get({vecWidth}, b.getI1Type());
+ OpFoldResult innerTileSize = b.getIndexAttr(vecWidth);
+ Value validSize = IREE::LinalgExt::computeIm2colValidSize(
+ b, loc, im2colOp, srcIndices, innerTileSize, outputIVs, vecOutputDim);
+ return vector::CreateMaskOp::create(b, loc, vecI1Type, validSize);
+}
+
+// Im2col vectorization uses the driver-provided vectorSizes to determine
+// which output dimension to vectorize and the vector width. The vector
+// sizes are computed by the MaterializeVectorTileSizes pass.
+struct Im2colOpVectorizationModel
+ : public VectorizableOpInterface::ExternalModel<Im2colOpVectorizationModel,
+ IREE::LinalgExt::Im2colOp> {
+ bool isVectorizable(Operation *op, ArrayRef<int64_t> vectorSizes,
+ ArrayRef<bool> scalableDims,
+ DictionaryAttr options) const {
+ auto im2colOp = cast<IREE::LinalgExt::Im2colOp>(op);
+ return im2colOp.getOutputType().hasStaticShape();
+ }
+
+ FailureOr<SmallVector<Value>> vectorize(Operation *op, RewriterBase &rewriter,
+ ArrayRef<int64_t> vectorSizes,
+ ArrayRef<bool> scalableDims,
+ DictionaryAttr options) const {
+ auto im2colOp = cast<IREE::LinalgExt::Im2colOp>(op);
+ RewriterBase::InsertionGuard g(rewriter);
+ rewriter.setInsertionPoint(im2colOp);
+ ShapedType outputType = im2colOp.getOutputType();
+ Location loc = im2colOp.getLoc();
+ bool hasPadding = im2colOp.hasPadding();
+
+ int64_t inputRank = im2colOp.getInputRank();
+
+ int64_t outputRank = im2colOp.getOutputRank();
+ ArrayRef<int64_t> outputShape = outputType.getShape();
+ Type elemType = outputType.getElementType();
+
+ // Determine the vectorized dimension from the driver-provided vectorSizes.
+ // The vectorized dim has size > 1; all others are 1.
+ std::optional<int64_t> vecDim;
+ for (int64_t d = 0; d < outputRank; ++d) {
+ if (d < static_cast<int64_t>(vectorSizes.size()) && vectorSizes[d] > 1) {
+ vecDim = d;
+ break;
+ }
+ }
+
+ int64_t vecWidth = vecDim ? outputShape[*vecDim] : 1;
+
+ auto vecType = VectorType::get({vecWidth}, elemType);
+
+ // Pad value for transfer_read: use im2col's pad_value when padding is
+ // present, otherwise use poison (transfer_read requires a padding operand).
+ Value padValue = hasPadding ? im2colOp.getPadValue()
+ : ub::PoisonOp::create(rewriter, loc, elemType);
+
+ int64_t writeDim = vecDim ? *vecDim : (outputRank - 1);
+ AffineMap writePermMap =
+ AffineMap::get(outputRank, 0, rewriter.getAffineDimExpr(writeDim),
+ rewriter.getContext());
+
+ SmallVector<int64_t> loopDims;
+ SmallVector<int64_t> loopBounds;
+ int64_t totalIters = 1;
+ for (int64_t d = 0; d < outputRank; ++d) {
+ if (vecDim && d == *vecDim) {
+ continue;
+ }
+ loopDims.push_back(d);
+ loopBounds.push_back(outputShape[d]);
+ totalIters *= outputShape[d];
+ }
+
+ Value result = im2colOp.getOutput();
+ Value zeroIdx = arith::ConstantIndexOp::create(rewriter, loc, 0);
+
+ // Hoist loop-invariant padding and clamping state.
+ SmallVector<OpFoldResult> padLow(inputRank, rewriter.getIndexAttr(0));
+ SmallVector<OpFoldResult> inputPadLow = im2colOp.getMixedInputPadLow();
+ if (!inputPadLow.empty()) {
+ padLow = inputPadLow;
+ }
+ SmallVector<OpFoldResult> inputSizes =
+ tensor::getMixedSizes(rewriter, loc, im2colOp.getInput());
+ // AffineMap for clamping: max(d0, 0) and min(d0, d1 - 1).
+ MLIRContext *ctx = rewriter.getContext();
+ AffineExpr d0 = getAffineDimExpr(0, ctx);
+ AffineExpr d1 = getAffineDimExpr(1, ctx);
+ AffineMap maxZeroMap =
+ AffineMap::get(1, 0, {d0, getAffineConstantExpr(0, ctx)}, ctx);
+ AffineMap clampHighMap = AffineMap::get(2, 0, {d0, d1 - 1}, ctx);
+
+ for (int64_t iter = 0; iter < totalIters; ++iter) {
+ SmallVector<Value> ivs(outputRank, zeroIdx);
+ int64_t remaining = iter;
+ for (int64_t i = loopDims.size() - 1; i >= 0; --i) {
+ int64_t idx = remaining % loopBounds[i];
+ remaining /= loopBounds[i];
+ ivs[loopDims[i]] = arith::ConstantIndexOp::create(rewriter, loc, idx);
+ }
+
+ IREE::LinalgExt::Im2colSourceIndices srcIndices =
+ IREE::LinalgExt::computeIm2colSourceIndices(
+ rewriter, loc, im2colOp, ivs, rewriter.getIndexAttr(vecWidth));
+
+ // Convert padded-space source offsets to actual input tensor coordinates
+ // by subtracting padLow. When there is no padding, padLow is all zeros
+ // and subOfrs folds to identity.
+ SmallVector<Value> readIndices;
+ for (int64_t d = 0; d < inputRank; ++d) {
+ OpFoldResult adjusted = IREE::LinalgExt::subOfrs(
+ rewriter, loc, srcIndices.sliceOffsets[d], padLow[d]);
+ // Clamp to [0, dimSize - 1] so downstream optimizations can prove
+ // buffer accesses are in-bounds. The mask already zeros out OOB reads,
+ // so clamping doesn't affect correctness.
+ if (hasPadding) {
+ adjusted = affine::makeComposedFoldedAffineMax(
+ rewriter, loc, maxZeroMap, {adjusted});
+ adjusted = affine::makeComposedFoldedAffineMin(
+ rewriter, loc, clampHighMap, {adjusted, inputSizes[d]});
+ }
+ readIndices.push_back(
+ getValueOrCreateConstantIndexOp(rewriter, loc, adjusted));
+ }
+ Value mask;
+ if (hasPadding) {
+ mask = computeIm2colPaddingMask(rewriter, loc, im2colOp, srcIndices,
+ vecWidth, ivs, vecDim);
+ }
+
+ AffineMap readPermMap =
+ AffineMap::getMinorIdentityMap(inputRank, 1, rewriter.getContext());
+ auto readOp = vector::TransferReadOp::create(
+ rewriter, loc, vecType, im2colOp.getInput(), readIndices,
+ AffineMapAttr::get(readPermMap), padValue, mask,
+ rewriter.getBoolArrayAttr({true}));
+ Value readVec = readOp.getResult();
+
+ SmallVector<Value> writeIndices(ivs);
+ if (vecDim) {
+ writeIndices[*vecDim] = zeroIdx;
+ }
+ result = vector::TransferWriteOp::create(rewriter, loc, readVec, result,
+ writeIndices, writePermMap)
+ .getResult();
+ }
+
+ return SmallVector<Value>{result};
+ }
+};
} // namespace
void registerVectorizableOpInterfaceExternalModels(DialectRegistry ®istry) {
@@ -1203,6 +1371,8 @@
ArgCompareOpVectorizationModel>(*ctx);
IREE::LinalgExt::MapStoreOp::attachInterface<MapStoreOpVectorizationModel>(
*ctx);
+ IREE::LinalgExt::Im2colOp::attachInterface<Im2colOpVectorizationModel>(
+ *ctx);
});
registry.addExtension(+[](MLIRContext *ctx,
IREE::VectorExt::IREEVectorExtDialect *dialect) {
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
index d379d9d..1c57e30 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
@@ -274,7 +274,6 @@
bool foldIdentitySlices,
bool decomposeMasks) {
funcPassManager.addPass(createDecomposeConvolutionToLowerDimOpsPass());
- funcPassManager.addPass(IREE::LinalgExt::createDecomposeIm2colPass());
funcPassManager.addPass(createCanonicalizerPass());
funcPassManager.addPass(createCSEPass());
if (enableMasking) {
@@ -288,6 +287,11 @@
options.enableVectorMasking = enableMasking;
options.vectorizeMapStore = true;
funcPassManager.addPass(createGenericVectorizationPass(options));
+ // Im2col decomposition runs after vectorization so that im2col ops get
+ // direct vectorization via VectorizableOpInterface when possible. Any
+ // remaining non-vectorized im2col ops (e.g., dynamic output shapes) are
+ // decomposed here as a fallback to avoid compilation failures.
+ funcPassManager.addPass(IREE::LinalgExt::createDecomposeIm2colPass());
funcPassManager.addPass(createCanonicalizerPass());
funcPassManager.addPass(createCSEPass());
// Run subset hoisting to convert iter_args to vectors.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_igemm_tile_and_fuse.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_igemm_tile_and_fuse.mlir
index a850016..619983c 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_igemm_tile_and_fuse.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_igemm_tile_and_fuse.mlir
@@ -65,8 +65,8 @@
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK: scf.forall ({{.*}}) in (2, 4, 5) {
// CHECK: %[[LOOP:.+]]:16 = scf.for {{.+}} = %[[C0]] to %[[C360]] step %[[C1]] {{.*}} -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>)
+// CHECK: %[[LHS_RD:.+]] = vector.transfer_read %[[BUF0]]{{.*}} vector<8xf16>
// CHECK: gpu.barrier memfence [#gpu.address_space<workgroup>]
-// CHECK-DAG: %[[LHS_RD:.+]] = vector.transfer_read %[[BUF0]]{{.*}} vector<8xf16>
// CHECK-DAG: vector.transfer_write %[[LHS_RD]]
// CHECK-DAG: %[[RHS_RD:.+]] = vector.transfer_read %[[BUF1]]{{.*}} vector<8xf16>
// CHECK-DAG: vector.transfer_write %[[RHS_RD]]
@@ -76,8 +76,8 @@
// CHECK-DAG: %[[RHS_MM:.+]] = vector.transfer_read {{.*}} vector<2x4x4x1xf16>
// CHECK-DAG: vector.transpose %[[RHS_MM]], [0, 2, 3, 1] : vector<2x4x4x1xf16> to vector<2x4x1x4xf16>
// CHECK-COUNT-32: amdgpu.mfma 16x16x16
-// CHECK: %[[SC:.+]] = vector.shape_cast %[[LOOP]]#0 : vector<4xf32> to vector<4x1xf32>
-// CHECK: %[[INSERT:.+]] = vector.insert_strided_slice %[[SC]], %{{.+}} {offsets = [3, 0, 3, 0, 0]{{.*}}} : vector<4x1xf32> into vector<4x1x4x4x1xf32>
+// CHECK: vector.shape_cast %[[LOOP]]#{{.+}} : vector<4xf32> to vector<4x1xf32>
+// CHECK: vector.insert_strided_slice {{.*}} {offsets = [3, 0, 3, 0, 0]{{.*}}} : vector<4x1xf32> into vector<4x1x4x4x1xf32>
// CHECK: %[[LOOP_T:.+]] = vector.transpose %{{.+}}, [0, 1, 2, 4, 3, 5] : vector<1x4x1x4x4x1xf32> to vector<1x4x1x4x4x1xf32>
// CHECK: %[[CAST:.+]] = vector.shape_cast %[[LOOP_T]] : vector<1x4x1x4x4x1xf32> to vector<4x1x4x4x1xf32>
// CHECK: vector.transfer_write %[[CAST]], %[[BUF2]]
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp
index a1e4b28..c9a294d 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp
@@ -4,10 +4,10 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#include "iree/compiler/Dialect/LinalgExt/IR/Im2colUtils.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/IndexingUtils.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
-#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/CommandLine.h"
@@ -289,27 +289,6 @@
return genericOp.getResult(0);
}
-// Helper method to check if a slice will be contiguous given the offset,
-// slice size. This checks that `inputSize` and `offset` are both evenly
-// divisible by `tileSize`.
-static bool willBeContiguousSlice(OpFoldResult inputSize, OpFoldResult tileSize,
- OpFoldResult offset) {
- auto constInputSize = getConstantIntValue(inputSize);
- auto constTileSize = getConstantIntValue(tileSize);
- if (!constTileSize.has_value() || !constInputSize.has_value() ||
- constInputSize.value() % constTileSize.value() != 0) {
- return false;
- }
- auto constOffset = getConstantIntValue(offset);
- if (constOffset.has_value() &&
- constOffset.value() % constTileSize.value() == 0) {
- return true;
- }
- auto affineOp = cast<Value>(offset).getDefiningOp<affine::AffineApplyOp>();
- return affineOp &&
- affineOp.getMap().getResult(0).isMultipleOf(constTileSize.value());
-}
-
//===----------------------------------------------------------------------===//
// Attention Helpers
//===----------------------------------------------------------------------===//
@@ -628,71 +607,6 @@
// Im2colOp
//===----------------------------------------------------------------------===//
-static std::optional<int64_t>
-chooseDimToVectorize(OpBuilder &b, Location loc, Im2colOp im2colOp,
- SmallVector<Range> iterationDomain,
- SmallVector<OpFoldResult> inputSizes,
- ArrayRef<OpFoldResult> offsets) {
- int64_t innerInputDim = im2colOp.getInputRank() - 1;
- SmallVector<SmallVector<int64_t>> vectorizationMap =
- im2colOp.getInputToOutputDimVectorizationMap();
- SmallVector<int64_t> vectorizableOutputDims = vectorizationMap[innerInputDim];
- if (vectorizableOutputDims.empty()) {
- return std::nullopt;
- }
- SetVector<int64_t> kDimSet(llvm::from_range, im2colOp.getKOutputDims());
-
- // Build a map from actual output dim to canonical index for K dims.
- SmallVector<int64_t> kOutputDims = im2colOp.getKOutputDims();
- int64_t batchSize = im2colOp.getBatchPos().size();
- int64_t numMOutputDims = im2colOp.getNumMOutputDims();
- DenseMap<int64_t, int64_t> kDimToCanonicalIdx;
- for (auto [i, actualDim] : llvm::enumerate(kOutputDims)) {
- kDimToCanonicalIdx[actualDim] = batchSize + numMOutputDims + i;
- }
-
- // There may be multiple output dims that we can vectorize, so prioritize the
- // innermost dims first.
- llvm::sort(vectorizableOutputDims);
- // Check each dim in order from innermost to outermost, and return the first
- // one that is vectorizable.
- while (!vectorizableOutputDims.empty()) {
- int64_t outputDimToVectorize = vectorizableOutputDims.pop_back_val();
- // If a K dim is being vectorized, then it is contiguous along either the
- // input channel dimension, or the filter kernel window. If it is contiguous
- // along the kernel window, then the actual inner slice size is equal to the
- // size of the corresponding kernel window dimension. Otherwise, the inner
- // slice size is just the size of the input tensor's inner dimension.
- OpFoldResult innerSliceSize = inputSizes[innerInputDim];
- if (kDimSet.contains(outputDimToVectorize)) {
- for (auto [kernelSize, mPos] :
- llvm::zip_equal(im2colOp.getMixedKernelSize(), im2colOp.getMPos())) {
- if (mPos == innerInputDim) {
- innerSliceSize = kernelSize;
- }
- }
- }
-
- // If the input slice is contiguous along the innermost dimension, then it
- // is vectorizable. If it is not, then move on to the next innermost dim.
- SetVector<int64_t> mDimSet(llvm::from_range, im2colOp.getMOutputDims());
- OpFoldResult offset = b.getIndexAttr(0);
- if (kDimSet.contains(outputDimToVectorize)) {
- // Use the offset of this specific K dim directly (no linearization).
- offset = offsets[kDimToCanonicalIdx[outputDimToVectorize]];
- } else if (mDimSet.contains(outputDimToVectorize)) {
- // TODO(Max191): Support vectorization along the M dimension.
- continue;
- }
- OpFoldResult outputDimSize = iterationDomain[outputDimToVectorize].size;
- if (!willBeContiguousSlice(innerSliceSize, outputDimSize, offset)) {
- continue;
- }
- return outputDimToVectorize;
- }
- return std::nullopt;
-}
-
/// Decomposition implementation for iree_linalg_ext.im2col op.
/// The im2col op is decomposed into serial loops of `insert->extract->copy`.
/// The decomposition supports leaving either the `batch` or `K` dimension
@@ -730,41 +644,32 @@
/// `%k` = `(%k_off + %K) mod 640`
///
FailureOr<SmallVector<Value>> Im2colOp::decomposeOperation(OpBuilder &b) {
- // Decomposition of padded im2col ops is not yet implemented.
- if (hasPadding()) {
- return failure();
- }
-
Location loc = getLoc();
Value inputSlice = getInput();
-
- // Get the per-output-dim offsets from the unified API.
SmallVector<OpFoldResult> mixedOffsets = getMixedOffsets();
+ SmallVector<SmallVector<OpFoldResult>> mixedOutputSizes =
+ getMixedOutputSizes();
- // Step 1: Tile the im2col op to loops with contiguous slices in the
- // innermost loop.
- //
- // If the innermost dim of the input tensor contains a full contiguous slice,
- // then don't tile the corresponding loop of the im2col op and maintain a
- // larger contiguous slice. Note that if the im2col input tensor has the batch
- // dim at last, im2col output tensor has an implicit transpose to move the
- // batch dim in front, and tiling should be along the batch dim.
+ int64_t outputRank = getOutputRank();
+ int64_t inputRank = getInputRank();
+
+ // Step 1: Choose the vectorization dimension.
SmallVector<Range> iterationDomain(getIterationDomain(b));
SmallVector<OpFoldResult> inputSizes =
tensor::getMixedSizes(b, loc, getInput());
- std::optional<unsigned> maybeOutputDimToVectorize = chooseDimToVectorize(
- b, loc, *this, iterationDomain, inputSizes, mixedOffsets);
+ std::optional<int64_t> maybeOutputDimToVectorize =
+ chooseDimToVectorize(b, loc, *this, iterationDomain, mixedOffsets);
OpFoldResult innerInputTileSize;
if (maybeOutputDimToVectorize.has_value()) {
- unsigned outputDimToVectorize = maybeOutputDimToVectorize.value();
- innerInputTileSize = iterationDomain[outputDimToVectorize].size;
- iterationDomain.erase(iterationDomain.begin() + outputDimToVectorize);
+ int64_t vecDim = maybeOutputDimToVectorize.value();
+ innerInputTileSize = iterationDomain[vecDim].size;
+ iterationDomain.erase(iterationDomain.begin() + vecDim);
} else {
innerInputTileSize = b.getIndexAttr(1);
}
- // Build loop nest.
+ // Build loop nest over all non-vectorized dimensions.
SmallVector<Value> lbs, ubs, steps;
for (auto range : iterationDomain) {
lbs.push_back(getValueOrCreateConstantIndexOp(b, loc, range.offset));
@@ -779,194 +684,126 @@
for (scf::ForOp loop : loopNest.loops) {
ivs.push_back(loop.getInductionVar());
}
- // The index computation below uses the induction variables as the offsets
- // into the output tensor, so we need an offset for each dim of the output.
- // For the dimension that is vectorized, the offset is zero, because we
- // take a full slice along that dimension.
+ // Step 2: When vectorizing a dim, insert a zero IV for it (it spans the full
+ // tile).
if (maybeOutputDimToVectorize.has_value()) {
- Value zero = arith::ConstantIndexOp::create(b, loc, 0);
- ivs.insert(ivs.begin() + maybeOutputDimToVectorize.value(), zero);
+ Value zeroIV = arith::ConstantIndexOp::create(b, loc, 0);
+ ivs.insert(ivs.begin() + maybeOutputDimToVectorize.value(), zeroIV);
}
- // Step 2: Compute indices into the input tensor for extract_slice.
+ // Step 3: Compute source indices.
OpBuilder::InsertionGuard guard(b);
- b.setInsertionPoint(loopNest.loops.front());
- SetVector<int64_t> mPosSet(getMPos().begin(), getMPos().end());
-
- ArrayRef<int64_t> strides = getStrides();
- ArrayRef<int64_t> dilations = getDilations();
-
Location nestedLoc =
loopNest.loops.back().getBody()->getTerminator()->getLoc();
b.setInsertionPointToStart(loopNest.loops.back().getBody());
- SetVector<int64_t> batchPosSet(getBatchPos().begin(), getBatchPos().end());
- ArrayRef<int64_t> inputKPerm = getInputKPerm();
- SmallVector<int64_t> invInputKPerm = invertPermutationVector(inputKPerm);
+ Im2colSourceIndices srcIndices =
+ computeIm2colSourceIndices(b, nestedLoc, *this, ivs, innerInputTileSize);
- // Get output_sizes for per-dim delinearization.
- SmallVector<SmallVector<OpFoldResult>> mixedOutputSizes =
- getMixedOutputSizes();
- SmallVector<int64_t> kOutputDims = getKOutputDims();
- int64_t batchSize = getBatchPos().size();
- int64_t numMOutputDims = getNumMOutputDims();
-
- // Delinearize each output dim independently using its output_sizes.
- // For each output dim at canonical index c with actual output dim d:
- // pos = offsets[c] + ivs[d]
- // components = delinearize(pos, output_sizes[c])
- // Concatenate all components into a flat list.
- auto delinearizeOutputDims =
- [&](ArrayRef<int64_t> outputDims,
- int64_t canonicalOffset) -> SmallVector<Value> {
- SmallVector<Value> results;
- for (auto [i, actualDim] : llvm::enumerate(outputDims)) {
- int64_t canonicalIdx = canonicalOffset + i;
- OpFoldResult pos =
- addOfrs(b, nestedLoc, mixedOffsets[canonicalIdx], ivs[actualDim]);
- const SmallVector<OpFoldResult> &innerSizes =
- mixedOutputSizes[canonicalIdx];
- if (innerSizes.size() == 1) {
- results.push_back(getValueOrCreateConstantIndexOp(b, nestedLoc, pos));
- } else {
- ValueRange components =
- affine::AffineDelinearizeIndexOp::create(
- b, nestedLoc,
- getValueOrCreateConstantIndexOp(b, nestedLoc, pos), innerSizes,
- /*hasOuterBound=*/true)
- .getResults();
- results.append(components.begin(), components.end());
- }
- }
- return results;
- };
-
- SmallVector<Value> delinKOffset =
- delinearizeOutputDims(kOutputDims, batchSize + numMOutputDims);
-
- // Split the delinearized offsets into the window offsets (for M offsets)
- // and the K offsets for the input tensor based on the layout.
- SmallVector<Value> windowOffset, inputKOffset;
- int delinKIdx = 0;
- for (int i = 0; i < getInputRank(); ++i) {
- if (batchPosSet.contains(i)) {
- continue;
- }
- if (mPosSet.contains(i)) {
- windowOffset.push_back(delinKOffset[invInputKPerm[delinKIdx++]]);
- continue;
- }
- inputKOffset.push_back(delinKOffset[invInputKPerm[delinKIdx++]]);
- }
-
- SmallVector<int64_t> mOutputDims = getMOutputDims();
- SmallVector<Value> delinMOffset =
- delinearizeOutputDims(mOutputDims, batchSize);
-
- // Compute the final offsets into the input tensor.
+ // The slice is always 1D — just a flat slice along the vectorized input
+ // dimension. With a 1D slice, no transpose is needed regardless of
+ // which output dimension is being vectorized.
+ ShapedType outputType = getOutputType();
OpFoldResult zero = b.getIndexAttr(0);
OpFoldResult one = b.getIndexAttr(1);
- SmallVector<OpFoldResult> sliceOffsets(getInputRank(), zero);
- SmallVector<OpFoldResult> sliceStrides(getInputRank(), one);
- SmallVector<OpFoldResult> sliceSizes(getInputRank(), one);
- // Add the offset into the convolution window, and account for strides and
- // dilations.
- AffineExpr mOff, wOff;
- bindDims(b.getContext(), mOff, wOff);
- for (auto [idx, mPos] : llvm::enumerate(getMPos())) {
- auto map =
- AffineMap::get(2, 0, {mOff * strides[idx] + wOff * dilations[idx]});
- OpFoldResult offset = affine::makeComposedFoldedAffineApply(
- b, nestedLoc, map, {delinMOffset[idx], windowOffset[idx]});
- sliceOffsets[mPos] = offset;
- sliceSizes[mPos] = one;
+ int64_t vecInputDim = inputRank - 1;
+
+ std::optional<int64_t> staticTileSize =
+ getConstantIntValue(innerInputTileSize);
+ int64_t paddedStaticSize =
+ staticTileSize ? *staticTileSize : ShapedType::kDynamic;
+
+ SmallVector<OpFoldResult> outOffsets(outputRank, zero);
+ for (auto [idx, iv] : llvm::enumerate(ivs)) {
+ outOffsets[idx] = iv;
}
-
- sliceSizes.back() = innerInputTileSize;
-
- // Set the batch and K offsets for the input tensor.
- assert(getKPos().size() == inputKOffset.size() &&
- "expected one delinearized K offset per k_pos input dimension");
- for (auto [kPos, kOff] : llvm::zip_equal(getKPos(), inputKOffset)) {
- sliceOffsets[kPos] = kOff;
- }
- SmallVector<int64_t> inverseOutputPerm =
- invertPermutationVector(getOutputPerm());
- for (auto [ivIdx, bPos] : llvm::enumerate(getBatchPos())) {
- int64_t canonicalIdx = ivIdx;
- int64_t actualDim = inverseOutputPerm[canonicalIdx];
- sliceOffsets[bPos] =
- addOfrs(b, nestedLoc, mixedOffsets[canonicalIdx], ivs[actualDim]);
- }
-
- // Step 3. Decompose the im2col op into:
- // ```
- // %extract = tensor.extract_slice %input
- // %copy = linalg.copy ins(%extract) outs(%out_slice)
- // %insert = tensor.insert_slice %copy into %loop_arg
- // ```
- //
- // Extract a slice from the input tensor.
- ShapedType outputType = getOutputType();
- int64_t inputRank = getInputRank();
- int64_t outputRank = getOutputRank();
-
- // For now, only extract a 1D slice when the vectorized dim is not innermost
- // in the output, and the input and output ranks are different. Otherwise,
- // try to preserve the original rank to avoid rank reducing slices.
- int64_t sliceRank = std::min(inputRank, outputRank);
- auto inputToOutputSlicePerm =
- llvm::to_vector(llvm::seq<int64_t>(0, sliceRank));
+ SmallVector<OpFoldResult> outSizes(outputRank, one);
if (maybeOutputDimToVectorize.has_value()) {
- int64_t outputDimToVectorize = maybeOutputDimToVectorize.value();
- if (inputRank == outputRank) {
- inputToOutputSlicePerm[outputDimToVectorize] = outputRank - 1;
- inputToOutputSlicePerm[outputRank - 1] = outputDimToVectorize;
- } else if (outputDimToVectorize != outputRank - 1) {
- sliceRank = 1;
- inputToOutputSlicePerm = {0};
- }
+ outSizes[maybeOutputDimToVectorize.value()] = innerInputTileSize;
}
- SmallVector<OpFoldResult> inputTileSizes(sliceRank, b.getIndexAttr(1));
- inputTileSizes.back() = innerInputTileSize;
- SmallVector<int64_t> tileSizeStatic;
- std::tie(tileSizeStatic, std::ignore) = decomposeMixedValues(inputTileSizes);
- auto extractType = cast<RankedTensorType>(outputType.clone(tileSizeStatic));
+ SmallVector<OpFoldResult> outStrides(outputRank, one);
+
+ // Step 4: Compute read offsets and extract the input slice.
+ // Subtract padLow from source offsets to get real input coordinates.
+ // Clamp to [0, dimSize - 1] to avoid negative indices in extract_slice.
+ // When out-of-bounds, validSize is 0 so extract_slice produces an empty
+ // slice and tensor.pad fills the padding.
+ SmallVector<OpFoldResult> readOffsets;
+ SmallVector<OpFoldResult> extractSizes(inputRank, one);
+ SmallVector<OpFoldResult> extractStrides(inputRank, one);
+ Value sliceToInsert;
+
+ SmallVector<OpFoldResult> padLow(inputRank, b.getIndexAttr(0));
+ SmallVector<OpFoldResult> inputPadLow = getMixedInputPadLow();
+ if (!inputPadLow.empty()) {
+ padLow = inputPadLow;
+ }
+ SmallVector<OpFoldResult> inputDimSizes =
+ tensor::getMixedSizes(b, nestedLoc, getInput());
+ MLIRContext *clampCtx = b.getContext();
+ AffineExpr cd0 = getAffineDimExpr(0, clampCtx);
+ AffineExpr cd1 = getAffineDimExpr(1, clampCtx);
+ AffineMap maxZeroMap =
+ AffineMap::get(1, 0, {cd0, getAffineConstantExpr(0, clampCtx)}, clampCtx);
+ AffineMap clampHighMap = AffineMap::get(2, 0, {cd0, cd1 - 1}, clampCtx);
+ for (int64_t d = 0; d < inputRank; ++d) {
+ OpFoldResult adjusted =
+ subOfrs(b, nestedLoc, srcIndices.sliceOffsets[d], padLow[d]);
+ if (hasPadding()) {
+ adjusted = affine::makeComposedFoldedAffineMax(b, nestedLoc, maxZeroMap,
+ {adjusted});
+ adjusted = affine::makeComposedFoldedAffineMin(
+ b, nestedLoc, clampHighMap, {adjusted, inputDimSizes[d]});
+ }
+ readOffsets.push_back(adjusted);
+ }
+
+ Value validSize;
+ if (hasPadding()) {
+ validSize = computeIm2colValidSize(b, nestedLoc, *this, srcIndices,
+ innerInputTileSize, ivs,
+ maybeOutputDimToVectorize);
+ extractSizes[vecInputDim] = validSize;
+ } else {
+ extractSizes[vecInputDim] = innerInputTileSize;
+ }
+
+ auto extractType = RankedTensorType::get(
+ {hasPadding() ? ShapedType::kDynamic : paddedStaticSize},
+ outputType.getElementType());
auto extract =
tensor::ExtractSliceOp::create(b, nestedLoc, extractType, inputSlice,
- sliceOffsets, sliceSizes, sliceStrides);
- // Insert the slice into the destination tensor.
- sliceOffsets = SmallVector<OpFoldResult>(outputRank, zero);
- for (auto [idx, iv] : llvm::enumerate(ivs)) {
- sliceOffsets[idx] = iv;
- }
- sliceSizes = SmallVector<OpFoldResult>(outputRank, one);
- if (maybeOutputDimToVectorize.has_value()) {
- sliceSizes[maybeOutputDimToVectorize.value()] = innerInputTileSize;
- }
- sliceStrides = SmallVector<OpFoldResult>(outputRank, one);
+ readOffsets, extractSizes, extractStrides);
- // Insert a `linalg.copy` so there is something to vectorize in the
- // decomposition. Without this copy, the extract and insert slice ops
- // do not get vectorized, and the sequence becomes a scalar memref.copy.
- // This memref.copy could be vectorized after bufferization, but it is
- // probably better to vectorize during generic vectorization.
- SmallVector<int64_t> outputSliceShape =
- applyPermutation(tileSizeStatic, inputToOutputSlicePerm);
- RankedTensorType outputSliceType = extractType.clone(outputSliceShape);
- Value copyDest = tensor::ExtractSliceOp::create(
- b, nestedLoc, outputSliceType, loopNest.loops.back().getRegionIterArg(0),
- sliceOffsets, sliceSizes, sliceStrides);
- Value copiedSlice =
- isIdentityPermutation(inputToOutputSlicePerm)
- ? linalg::CopyOp::create(b, nestedLoc, extract.getResult(), copyDest)
- .getResult(0)
- : linalg::TransposeOp::create(b, nestedLoc, extract.getResult(),
- copyDest, inputToOutputSlicePerm)
- ->getResult(0);
+ // Branch only on the vectorizable payload:
+ // - No padding: linalg.copy (static type, concrete copy op)
+ // - Has padding: tensor.pad (dynamic extract padded to static size)
+ if (!hasPadding()) {
+ auto sliceType = cast<RankedTensorType>(extract.getType());
+ auto destExtract = tensor::ExtractSliceOp::create(
+ b, nestedLoc, sliceType, loopNest.loops.back().getRegionIterArg(0),
+ outOffsets, outSizes, outStrides);
+ auto copy = linalg::CopyOp::create(b, nestedLoc, extract.getResult(),
+ destExtract.getResult());
+ sliceToInsert = copy.getResult(0);
+ } else {
+ Value tileSize =
+ getValueOrCreateConstantIndexOp(b, nestedLoc, innerInputTileSize);
+ SmallVector<OpFoldResult> lowPad = {b.getIndexAttr(0)};
+ Value highPadAmt = arith::SubIOp::create(b, nestedLoc, tileSize, validSize);
+ SmallVector<OpFoldResult> highPad = {highPadAmt};
+
+ auto paddedType =
+ RankedTensorType::get({paddedStaticSize}, outputType.getElementType());
+ auto paddedSlice =
+ tensor::PadOp::create(b, nestedLoc, paddedType, extract.getResult(),
+ lowPad, highPad, getPadValue(), /*nofold=*/false);
+ sliceToInsert = paddedSlice.getResult();
+ }
+
auto insert = tensor::InsertSliceOp::create(
- b, nestedLoc, copiedSlice, loopNest.loops.back().getRegionIterArg(0),
- sliceOffsets, sliceSizes, sliceStrides);
+ b, nestedLoc, sliceToInsert, loopNest.loops.back().getRegionIterArg(0),
+ outOffsets, outSizes, outStrides);
auto yieldOp =
cast<scf::YieldOp>(loopNest.loops.back().getBody()->getTerminator());
yieldOp->getOpOperands().front().assign(insert.getResult());
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/BUILD.bazel b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/BUILD.bazel
index 85abe1c..d30d0d7 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/BUILD.bazel
@@ -48,6 +48,7 @@
name = "IR",
srcs = [
"AggregatedOpInterfaceImpl.cpp",
+ "Im2colUtils.cpp",
"LinalgExtAttrs.cpp.inc",
"LinalgExtDialect.cpp",
"LinalgExtDialect.cpp.inc",
@@ -59,6 +60,7 @@
"TilingInterfaceImpl.cpp",
],
hdrs = [
+ "Im2colUtils.h",
"LinalgExtAttrs.h.inc",
"LinalgExtDialect.h",
"LinalgExtDialect.h.inc",
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/CMakeLists.txt
index 9b23c9e..50a477c 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/CMakeLists.txt
@@ -14,6 +14,7 @@
NAME
IR
HDRS
+ "Im2colUtils.h"
"LinalgExtAttrs.h.inc"
"LinalgExtDialect.h"
"LinalgExtDialect.h.inc"
@@ -27,6 +28,7 @@
"LinalgExtOps.h.inc"
SRCS
"AggregatedOpInterfaceImpl.cpp"
+ "Im2colUtils.cpp"
"LinalgExtAttrs.cpp.inc"
"LinalgExtDialect.cpp"
"LinalgExtDialect.cpp.inc"
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/Im2colUtils.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/Im2colUtils.cpp
new file mode 100644
index 0000000..c02bc60
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/Im2colUtils.cpp
@@ -0,0 +1,459 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/LinalgExt/IR/Im2colUtils.h"
+
+#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
+#include "llvm/ADT/DenseSet.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+
+namespace mlir::iree_compiler::IREE::LinalgExt {
+
+Im2colSourceIndices computeIm2colSourceIndices(OpBuilder &b, Location loc,
+ Im2colOp im2colOp,
+ ArrayRef<Value> ivs,
+ OpFoldResult innerTileSize) {
+ int64_t inputRank = im2colOp.getInputRank();
+ SmallVector<OpFoldResult> offsets = im2colOp.getMixedOffsets();
+ SmallVector<SmallVector<OpFoldResult>> outputSizes =
+ im2colOp.getMixedOutputSizes();
+ int64_t batchSize = im2colOp.getBatchPos().size();
+ int64_t numMOutputDims = im2colOp.getNumMOutputDims();
+ llvm::SmallDenseSet<int64_t, 4> mPosSet(im2colOp.getMPos().begin(),
+ im2colOp.getMPos().end());
+ llvm::SmallDenseSet<int64_t, 4> batchPosSet(im2colOp.getBatchPos().begin(),
+ im2colOp.getBatchPos().end());
+ ArrayRef<int64_t> strides = im2colOp.getStrides();
+ ArrayRef<int64_t> dilations = im2colOp.getDilations();
+ ArrayRef<int64_t> inputKPerm = im2colOp.getInputKPerm();
+
+ // Phase 1: Delinearize all canonical output dims uniformly.
+ // Canonical order is [batch..., M..., K...]. For each dim d,
+ // delinearize (offset[d] + iv[actualDim]) using output_sizes[d].
+ SmallVector<int64_t> inverseOutputPerm =
+ invertPermutationVector(im2colOp.getOutputPerm());
+ int64_t numCanonicalDims = static_cast<int64_t>(outputSizes.size());
+ SmallVector<Value> allCoords;
+ int64_t batchCoordCount = 0;
+ int64_t mCoordCount = 0;
+ for (int64_t d = 0; d < numCanonicalDims; ++d) {
+ int64_t actualDim = inverseOutputPerm[d];
+ OpFoldResult idx = addOfrs(b, loc, offsets[d], ivs[actualDim]);
+ const SmallVector<OpFoldResult> &innerSizes = outputSizes[d];
+ int64_t numProduced = static_cast<int64_t>(innerSizes.size());
+ if (numProduced == 1) {
+ allCoords.push_back(getValueOrCreateConstantIndexOp(b, loc, idx));
+ } else {
+ ValueRange delinCoords =
+ affine::AffineDelinearizeIndexOp::create(
+ b, loc, getValueOrCreateConstantIndexOp(b, loc, idx), innerSizes,
+ /*hasOuterBound=*/true)
+ .getResults();
+ allCoords.append(delinCoords.begin(), delinCoords.end());
+ }
+ if (d < batchSize) {
+ batchCoordCount += numProduced;
+ } else if (d < batchSize + numMOutputDims) {
+ mCoordCount += numProduced;
+ }
+ }
+
+ // Phase 2: Split delinearized coords into batch, M, K groups.
+ auto it = allCoords.begin();
+ SmallVector<Value> batchCoords(it, it + batchCoordCount);
+ it += batchCoordCount;
+ SmallVector<Value> mCoords(it, it + mCoordCount);
+ it += mCoordCount;
+ SmallVector<Value> kCoords(it, allCoords.end());
+
+ // Phase 3: Organize coords into input dimension order.
+ // K coords: apply inverse input_k_perm to map from output K order to
+ // canonical input order, then split into window offsets and channel offsets.
+ SmallVector<int64_t> invInputKPerm = invertPermutationVector(inputKPerm);
+ // applyPermutationToVector performs a gather (result[i] = src[perm[i]]),
+ // which is the correct inverse mapping from output K order to input order.
+ applyPermutationToVector(kCoords, invInputKPerm);
+ SmallVector<Value> windowOffset, inputKOffset;
+ int64_t kIdx = 0;
+ for (int64_t i = 0; i < inputRank; ++i) {
+ if (batchPosSet.contains(i)) {
+ continue;
+ }
+ if (mPosSet.contains(i)) {
+ windowOffset.push_back(kCoords[kIdx++]);
+ continue;
+ }
+ inputKOffset.push_back(kCoords[kIdx++]);
+ }
+
+ // Compute final offsets into the input tensor.
+ OpFoldResult zero = b.getIndexAttr(0);
+ OpFoldResult one = b.getIndexAttr(1);
+ SmallVector<OpFoldResult> sliceOffsets(inputRank, zero);
+ SmallVector<OpFoldResult> sliceSizes(inputRank, one);
+
+ // Spatial dims: apply strides and dilations.
+ AffineExpr mOff, wOff;
+ bindDims(b.getContext(), mOff, wOff);
+ for (auto [idx, mPos] : llvm::enumerate(im2colOp.getMPos())) {
+ auto map =
+ AffineMap::get(2, 0, {mOff * strides[idx] + wOff * dilations[idx]});
+ OpFoldResult offset = affine::makeComposedFoldedAffineApply(
+ b, loc, map, {mCoords[idx], windowOffset[idx]});
+ sliceOffsets[mPos] = offset;
+ }
+
+ // K dims: set channel offsets directly.
+ for (auto [kPos, kOff] : llvm::zip_equal(im2colOp.getKPos(), inputKOffset)) {
+ sliceOffsets[kPos] = kOff;
+ }
+
+ // Batch dims: set delinearized batch coords directly.
+ int64_t batchIdx = 0;
+ for (int64_t bPos : im2colOp.getBatchPos()) {
+ sliceOffsets[bPos] = batchCoords[batchIdx++];
+ }
+
+ // The innermost input dimension gets innerTileSize as its size.
+ int64_t innerInputDim = inputRank - 1;
+ sliceSizes[innerInputDim] = innerTileSize;
+
+ return Im2colSourceIndices{sliceOffsets, sliceSizes};
+}
+
+Value computeIm2colValidSize(OpBuilder &b, Location loc, Im2colOp im2colOp,
+ const Im2colSourceIndices &srcIndices,
+ OpFoldResult innerTileSize,
+ ArrayRef<Value> outputIVs,
+ std::optional<int64_t> vecOutputDim) {
+ int64_t inputRank = im2colOp.getInputRank();
+ int64_t vecInputDim = inputRank - 1;
+
+ SmallVector<OpFoldResult> inputSizes =
+ tensor::getMixedSizes(b, loc, im2colOp.getInput());
+
+ // Get padding from the op.
+ SmallVector<OpFoldResult> padLow(inputRank, b.getIndexAttr(0));
+ SmallVector<OpFoldResult> padHigh(inputRank, b.getIndexAttr(0));
+ SmallVector<OpFoldResult> inputPadLow = im2colOp.getMixedInputPadLow();
+ SmallVector<OpFoldResult> inputPadHigh = im2colOp.getMixedInputPadHigh();
+ if (!inputPadLow.empty()) {
+ padLow = inputPadLow;
+ padHigh = inputPadHigh;
+ }
+
+ // Compute adjusted offsets: subtract padLow to get unpadded-space coords.
+ SmallVector<OpFoldResult> adjustedOffsets(inputRank);
+ for (int64_t d = 0; d < inputRank; ++d) {
+ adjustedOffsets[d] = subOfrs(b, loc, srcIndices.sliceOffsets[d], padLow[d]);
+ }
+
+ // When a dim is being vectorized, chooseDimToVectorize guarantees no low
+ // padding on the vectorized input dim. Assert this invariant.
+ assert((!vecOutputDim.has_value() ||
+ isConstantIntValue(padLow[vecInputDim], 0)) &&
+ "vectorized input dim must have zero low padding");
+
+ // --- Helper lambdas for affine clamping patterns ---
+ // All use affine ops so that IREE's IntegerDivisibilityAnalysis can track
+ // divisibility through the chain (it tracks affine.apply/min/max and
+ // arith.muli, but NOT arith.subi/maxsi/minsi).
+ MLIRContext *ctx = b.getContext();
+ AffineExpr d0 = getAffineDimExpr(0, ctx);
+ AffineExpr d1 = getAffineDimExpr(1, ctx);
+ AffineMap subMap = AffineMap::get(2, 0, d0 - d1, ctx);
+ AffineMap minMap = AffineMap::get(2, 0, {d0, d1}, ctx);
+ AffineMap maxZeroMap =
+ AffineMap::get(1, 0, {d0, getAffineConstantExpr(0, ctx)}, ctx);
+ AffineMap clampHighToOneMap =
+ AffineMap::get(2, 0, {d0 - d1, getAffineConstantExpr(1, ctx)}, ctx);
+ AffineMap clampLowToOneMap =
+ AffineMap::get(1, 0, {d0 + 1, getAffineConstantExpr(1, ctx)}, ctx);
+
+ // max(val, 0).
+ auto clampAboveZero = [&](OpFoldResult val) -> OpFoldResult {
+ return affine::makeComposedFoldedAffineMax(b, loc, maxZeroMap, {val});
+ };
+
+ // 0/1 factor: 1 when coord ∈ [0, dimSize), 0 otherwise.
+ auto clampToRange = [&](OpFoldResult coord, OpFoldResult dimSize) -> Value {
+ // highOk = max(min(dimSize - coord, 1), 0): 1 when coord < dimSize.
+ OpFoldResult highOk = clampAboveZero(affine::makeComposedFoldedAffineMin(
+ b, loc, clampHighToOneMap, {dimSize, coord}));
+ // lowOk = max(min(coord + 1, 1), 0): 1 when coord >= 0.
+ OpFoldResult lowOk = clampAboveZero(
+ affine::makeComposedFoldedAffineMin(b, loc, clampLowToOneMap, {coord}));
+ Value highOkVal = getValueOrCreateConstantIndexOp(b, loc, highOk);
+ Value lowOkVal = getValueOrCreateConstantIndexOp(b, loc, lowOk);
+ return arith::MulIOp::create(b, loc, highOkVal, lowOkVal);
+ };
+
+ // min(max(extent - coord, 0), tileSize): how much of tileSize fits within
+ // [coord, extent), clamped to [0, tileSize].
+ auto remainingValid = [&](OpFoldResult extent, OpFoldResult coord,
+ OpFoldResult tileSize) -> OpFoldResult {
+ OpFoldResult rawValid =
+ affine::makeComposedFoldedAffineApply(b, loc, subMap, {extent, coord});
+ OpFoldResult clamped = clampAboveZero(rawValid);
+ return affine::makeComposedFoldedAffineMin(b, loc, minMap,
+ {clamped, tileSize});
+ };
+
+ // 0/1 factor: 1 when pos >= low, 0 otherwise.
+ auto isAboveLow = [&](OpFoldResult pos, OpFoldResult low) -> Value {
+ OpFoldResult adjusted = subOfrs(b, loc, pos, low);
+ OpFoldResult lowOk = clampAboveZero(affine::makeComposedFoldedAffineMin(
+ b, loc, clampLowToOneMap, {adjusted}));
+ return getValueOrCreateConstantIndexOp(b, loc, lowOk);
+ };
+
+ // 0/1 factor: 1 when pos < high, 0 otherwise.
+ auto isBelowHigh = [&](OpFoldResult pos, OpFoldResult high) -> Value {
+ OpFoldResult highOk = clampAboveZero(affine::makeComposedFoldedAffineMin(
+ b, loc, clampHighToOneMap, {high, pos}));
+ return getValueOrCreateConstantIndexOp(b, loc, highOk);
+ };
+
+ // --- Compute valid_size along the innermost input dimension ---
+ OpFoldResult validSizeOfr = remainingValid(
+ inputSizes[vecInputDim], adjustedOffsets[vecInputDim], innerTileSize);
+ Value validSize = getValueOrCreateConstantIndexOp(b, loc, validSizeOfr);
+
+ // --- Input-side bounds checking for non-vectorized dims ---
+ // If an adjusted coord is outside [0, dimSize), the valid region is empty,
+ // so multiply validSize by 0. This handles input padding, output-alignment
+ // OOB from non-wrapping delinearization, and batch OOB.
+ auto checkDimBounds = [&](int64_t dim) {
+ // The vectorized input dim's range is already handled above. Skip it
+ // to avoid redundant IR. In scalar mode (no vecOutputDim), all dims
+ // need full bounds checking.
+ if (vecOutputDim.has_value() && dim == vecInputDim) {
+ return;
+ }
+ // Skip bounds check when this dim has no input padding — the offset
+ // is guaranteed to be in [0, dimSize) by construction.
+ if (isZeroInteger(padLow[dim]) && isZeroInteger(padHigh[dim])) {
+ return;
+ }
+ Value factor = clampToRange(adjustedOffsets[dim], inputSizes[dim]);
+ validSize = arith::MulIOp::create(b, loc, validSize, factor);
+ };
+ for (int64_t bPos : im2colOp.getBatchPos()) {
+ checkDimBounds(bPos);
+ }
+ for (int64_t mPos : im2colOp.getMPos()) {
+ checkDimBounds(mPos);
+ }
+ for (int64_t kPos : im2colOp.getKPos()) {
+ checkDimBounds(kPos);
+ }
+
+ // --- Output-side bounds checking ---
+ // For each output dim, positions in [0, pad_low) and [dim - pad_high, dim)
+ // are padding positions and should produce pad_value.
+ // chooseDimToVectorize guarantees output_pad_low[vecOutputDim] == 0.
+ SmallVector<OpFoldResult> outPadLow = im2colOp.getMixedOutputPadLow();
+ SmallVector<OpFoldResult> outPadHigh = im2colOp.getMixedOutputPadHigh();
+ if (!outPadLow.empty()) {
+ assert((!vecOutputDim.has_value() ||
+ isConstantIntValue(outPadLow[*vecOutputDim], 0)) &&
+ "vectorized output dim must have zero output low padding");
+ int64_t outputRank = im2colOp.getOutputRank();
+ SmallVector<OpFoldResult> outputTensorSizes =
+ tensor::getMixedSizes(b, loc, im2colOp.getOutput());
+
+ // Non-vectorized output dims: if pos is outside [padLow, dim - padHigh),
+ // set validSize = 0.
+ for (int64_t d = 0; d < outputRank; ++d) {
+ if (vecOutputDim.has_value() && d == vecOutputDim.value()) {
+ continue;
+ }
+ if (isConstantIntValue(outPadLow[d], 0) &&
+ isConstantIntValue(outPadHigh[d], 0)) {
+ continue;
+ }
+
+ OpFoldResult localPos = outputIVs[d];
+
+ if (!isConstantIntValue(outPadLow[d], 0)) {
+ Value factor = isAboveLow(localPos, outPadLow[d]);
+ validSize = arith::MulIOp::create(b, loc, validSize, factor);
+ }
+ if (!isConstantIntValue(outPadHigh[d], 0)) {
+ OpFoldResult validEnd =
+ subOfrs(b, loc, outputTensorSizes[d], outPadHigh[d]);
+ Value factor = isBelowHigh(localPos, validEnd);
+ validSize = arith::MulIOp::create(b, loc, validSize, factor);
+ }
+ }
+
+ // Vectorized output dim: clamp validSize by remaining valid output count.
+ // output_pad_low[vecOutputDim] == 0 is asserted above, so only high-side
+ // padding needs checking here.
+ if (vecOutputDim.has_value()) {
+ int64_t vd = vecOutputDim.value();
+ if (!isConstantIntValue(outPadHigh[vd], 0)) {
+ OpFoldResult validEnd =
+ subOfrs(b, loc, outputTensorSizes[vd], outPadHigh[vd]);
+ OpFoldResult vsOfr = remainingValid(validEnd, outputIVs[vd], validSize);
+ validSize = getValueOrCreateConstantIndexOp(b, loc, vsOfr);
+ }
+ }
+ }
+
+ return validSize;
+}
+
+/// Helper to check if a slice will be contiguous given the offset and
+/// slice size. Checks that `inputSize` and `offset` are both evenly
+/// divisible by `tileSize`.
+static bool willBeContiguousSlice(OpFoldResult inputSize, OpFoldResult tileSize,
+ OpFoldResult offset) {
+ std::optional<int64_t> constInputSize = getConstantIntValue(inputSize);
+ std::optional<int64_t> constTileSize = getConstantIntValue(tileSize);
+ if (!constTileSize.has_value() || !constInputSize.has_value() ||
+ constInputSize.value() % constTileSize.value() != 0) {
+ return false;
+ }
+ std::optional<int64_t> constOffset = getConstantIntValue(offset);
+ if (constOffset.has_value()) {
+ return constOffset.value() % constTileSize.value() == 0;
+ }
+ auto val = dyn_cast<Value>(offset);
+ if (!val) {
+ return false;
+ }
+ auto affineOp = val.getDefiningOp<affine::AffineApplyOp>();
+ return affineOp &&
+ affineOp.getMap().getResult(0).isMultipleOf(constTileSize.value());
+}
+
+std::optional<int64_t> chooseDimToVectorize(OpBuilder &b, Location loc,
+ Im2colOp im2colOp,
+ ArrayRef<Range> iterationDomain,
+ ArrayRef<OpFoldResult> offsets) {
+ int64_t innerInputDim = im2colOp.getInputRank() - 1;
+ SmallVector<OpFoldResult> inputSizes =
+ tensor::getMixedSizes(b, loc, im2colOp.getInput());
+ SmallVector<SmallVector<int64_t>> vectorizationMap =
+ im2colOp.getInputToOutputDimVectorizationMap();
+ SmallVector<int64_t> vectorizableOutputDims = vectorizationMap[innerInputDim];
+ if (vectorizableOutputDims.empty()) {
+ return std::nullopt;
+ }
+
+ // Bail when the innermost input dim has non-zero low padding. Low padding
+ // on the vectorized input dim would require shifting read indices which
+ // complicates the valid size computation. Non-vectorized dims handle low
+ // padding through the per-dim bounds checks in computeIm2colValidSize.
+ SmallVector<OpFoldResult> inputPadLow = im2colOp.getMixedInputPadLow();
+ if (!inputPadLow.empty() &&
+ !isConstantIntValue(inputPadLow[innerInputDim], 0)) {
+ return std::nullopt;
+ }
+
+ // Get output low padding for per-candidate checks below.
+ SmallVector<OpFoldResult> outPadLow = im2colOp.getMixedOutputPadLow();
+
+ SetVector<int64_t> kDimSet(llvm::from_range, im2colOp.getKOutputDims());
+
+ // Build a map from actual output dim to canonical index for K dims.
+ SmallVector<int64_t> kOutputDims = im2colOp.getKOutputDims();
+ int64_t batchSize = im2colOp.getBatchPos().size();
+ int64_t numMOutputDims = im2colOp.getNumMOutputDims();
+ DenseMap<int64_t, int64_t> kDimToCanonicalIdx;
+ for (auto [i, actualDim] : llvm::enumerate(kOutputDims)) {
+ kDimToCanonicalIdx[actualDim] = batchSize + numMOutputDims + i;
+ }
+
+ // There may be multiple output dims that we can vectorize, so prioritize the
+ // innermost dims first.
+ llvm::sort(vectorizableOutputDims);
+ // Check each dim in order from innermost to outermost, and return the first
+ // one that is vectorizable.
+ for (int64_t outputDimToVectorize : llvm::reverse(vectorizableOutputDims)) {
+ // If a K dim is being vectorized, then it is contiguous along either the
+ // input channel dimension, or the filter kernel window. If it is contiguous
+ // along the kernel window, then the actual inner slice size is equal to the
+ // size of the corresponding kernel window dimension. Otherwise, the inner
+ // slice size is just the size of the input tensor's inner dimension.
+ // Use the padded input size for the contiguity check: the im2col operates
+ // in the padded coordinate space, so the effective innermost dimension
+ // includes both low and high padding.
+ OpFoldResult innerSliceSize = inputSizes[innerInputDim];
+ SmallVector<OpFoldResult> inputPadHigh = im2colOp.getMixedInputPadHigh();
+ if (!inputPadLow.empty()) {
+ innerSliceSize =
+ addOfrs(b, loc, innerSliceSize, inputPadLow[innerInputDim]);
+ innerSliceSize =
+ addOfrs(b, loc, innerSliceSize, inputPadHigh[innerInputDim]);
+ }
+ if (kDimSet.contains(outputDimToVectorize)) {
+ for (auto [kernelSize, mPos] :
+ llvm::zip_equal(im2colOp.getMixedKernelSize(), im2colOp.getMPos())) {
+ if (mPos == innerInputDim) {
+ innerSliceSize = kernelSize;
+ }
+ }
+ }
+
+ // If the input slice is contiguous along the innermost dimension, then it
+ // is vectorizable. If it is not, then move on to the next innermost dim.
+ SetVector<int64_t> mDimSet(llvm::from_range, im2colOp.getMOutputDims());
+ OpFoldResult offset = b.getIndexAttr(0);
+ if (kDimSet.contains(outputDimToVectorize)) {
+ // Use the offset of this specific K dim directly (no linearization).
+ offset = offsets[kDimToCanonicalIdx[outputDimToVectorize]];
+ } else if (mDimSet.contains(outputDimToVectorize)) {
+ // TODO(Max191): Support vectorization along the M dimension.
+ continue;
+ }
+ // Skip dims with non-zero output low padding. Low padding on the
+ // vectorized output dim would require shifting write positions and
+ // complicates the valid size computation.
+ if (!outPadLow.empty() &&
+ !isConstantIntValue(outPadLow[outputDimToVectorize], 0)) {
+ continue;
+ }
+
+ OpFoldResult outputDimSize = iterationDomain[outputDimToVectorize].size;
+ if (!willBeContiguousSlice(innerSliceSize, outputDimSize, offset)) {
+ continue;
+ }
+ return outputDimToVectorize;
+ }
+ return std::nullopt;
+}
+
+std::optional<SmallVector<int64_t>>
+computeIm2colVectorTileSizes(OpBuilder &b, Im2colOp im2colOp) {
+ Location loc = im2colOp.getLoc();
+ SmallVector<Range> iterationDomain(im2colOp.getIterationDomain(b));
+ SmallVector<OpFoldResult> mixedOffsets = im2colOp.getMixedOffsets();
+ std::optional<int64_t> dimToVectorize =
+ chooseDimToVectorize(b, loc, im2colOp, iterationDomain, mixedOffsets);
+
+ int64_t outputRank = im2colOp.getOutputRank();
+ SmallVector<int64_t> tileSizes(outputRank, 1);
+ if (!dimToVectorize.has_value()) {
+ return std::nullopt;
+ }
+ std::optional<int64_t> dimSize =
+ getConstantIntValue(iterationDomain[dimToVectorize.value()].size);
+ if (!dimSize.has_value()) {
+ return std::nullopt;
+ }
+ tileSizes[dimToVectorize.value()] = dimSize.value();
+ return tileSizes;
+}
+
+} // namespace mlir::iree_compiler::IREE::LinalgExt
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/Im2colUtils.h b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/Im2colUtils.h
new file mode 100644
index 0000000..fcff813
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/Im2colUtils.h
@@ -0,0 +1,88 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_DIALECT_LINALGEXT_IR_IM2COLUTILS_H_
+#define IREE_COMPILER_DIALECT_LINALGEXT_IR_IM2COLUTILS_H_
+
+// TODO(Max191): Move to Utils/ once the IR/ -> Utils/ dependency cycle is
+// resolved. Im2colUtils.{h,cpp} are transformation helpers consumed by
+// Transforms/ and IR/AggregatedOpInterfaceImpl, not op definitions or dialect
+// infrastructure, but Utils/ cannot currently depend on IR/ types like
+// Im2colOp.
+
+#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir::iree_compiler::IREE::LinalgExt {
+
+/// Holds the computed source indices for an im2col operation at a given
+/// output position. These indices describe where to read from the input tensor.
+struct Im2colSourceIndices {
+ /// Full set of offsets into the input tensor, one per input dimension.
+ /// When padding is present, these are in the padded coordinate space.
+ SmallVector<OpFoldResult> sliceOffsets;
+ /// Sizes for each input dimension (1 except for the vectorized dim).
+ SmallVector<OpFoldResult> sliceSizes;
+};
+
+/// Compute source (input tensor) indices for a given im2col output position.
+///
+/// Given the loop induction variables representing the current output position,
+/// compute the corresponding offsets and sizes into the input tensor. Uses the
+/// unified offsets + output_sizes attributes to delinearize each output dim
+/// independently, then maps to input coordinates via strides, dilations, and
+/// input_k_perm.
+///
+/// Shared by both the decomposition and vectorization paths.
+Im2colSourceIndices computeIm2colSourceIndices(OpBuilder &b, Location loc,
+ Im2colOp im2colOp,
+ ArrayRef<Value> ivs,
+ OpFoldResult innerTileSize);
+
+/// Compute the valid read size along the vectorized (innermost input) dimension
+/// for im2col decomposition and vectorization.
+///
+/// Given source indices in the padded coordinate space, determines how many
+/// elements can be validly read from the input tensor. Accounts for:
+/// - Input padding (pad_low/pad_high on the convolution input)
+/// - Output padding (output_pad_low/output_pad_high on the output)
+/// - Out-of-bounds status of non-vectorized dimensions
+///
+/// When any non-vectorized dimension is out-of-bounds, returns 0.
+///
+/// Callers are responsible for computing their own read offsets from the
+/// source indices (e.g., by subtracting pad_low and clamping).
+///
+/// \p outputIVs are the loop IVs for output dimensions (in actual tensor dim
+/// order). Used for output-side bounds checking when output padding is present.
+Value computeIm2colValidSize(OpBuilder &b, Location loc, Im2colOp im2colOp,
+ const Im2colSourceIndices &srcIndices,
+ OpFoldResult innerTileSize,
+ ArrayRef<Value> outputIVs,
+ std::optional<int64_t> vecOutputDim);
+
+/// Choose which output dimension to vectorize for an im2col op.
+/// Returns the output dimension index, or std::nullopt if no dimension can be
+/// vectorized (in which case scalar unrolling should be used).
+///
+/// \p offsets are the per-output-dim offsets from the im2col op's attributes.
+/// For K output dims, the offset of the specific dim being considered is used
+/// directly for the contiguity check (no linearization needed).
+std::optional<int64_t> chooseDimToVectorize(OpBuilder &b, Location loc,
+ Im2colOp im2colOp,
+ ArrayRef<Range> iterationDomain,
+ ArrayRef<OpFoldResult> offsets);
+
+/// Compute vector tile sizes for an im2col op. Returns a vector of tile sizes
+/// with one entry per output dimension. The vectorizable dimension (if any)
+/// gets its full iteration size; all other dimensions get 1. Returns nullopt
+/// if no vectorizable dimension is found (e.g. no contiguous slice exists).
+std::optional<SmallVector<int64_t>>
+computeIm2colVectorTileSizes(OpBuilder &b, Im2colOp im2colOp);
+
+} // namespace mlir::iree_compiler::IREE::LinalgExt
+
+#endif // IREE_COMPILER_DIALECT_LINALGEXT_IR_IM2COLUTILS_H_
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp
index 34de4e9..c32c758 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp
@@ -33,6 +33,7 @@
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/TypeUtilities.h"
@@ -2721,6 +2722,194 @@
return success();
}
+namespace {
+
+/// Fold tensor.pad on the input of im2col into the im2col's padding
+/// attributes.
+///
+/// %padded = tensor.pad %input low[...] high[...] { yield %cst }
+/// %result = im2col ins(%padded) outs(%out) ...
+/// -->
+/// %result = im2col ins(%input) outs(%out) ...
+/// input_pad_low=[...] input_pad_high=[...] pad_value(%cst : type)
+struct FoldInputPadIntoIm2col final : public OpRewritePattern<Im2colOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(Im2colOp im2colOp,
+ PatternRewriter &rewriter) const override {
+ auto padOp = im2colOp.getInput().getDefiningOp<tensor::PadOp>();
+ if (!padOp) {
+ return rewriter.notifyMatchFailure(im2colOp,
+ "input not produced by tensor.pad");
+ }
+
+ // Only fold constant padding values.
+ Value padValue = padOp.getConstantPaddingValue();
+ if (!padValue) {
+ return rewriter.notifyMatchFailure(im2colOp,
+ "pad value is not a constant");
+ }
+
+ // If the im2col already has padding, the pad values must be compatible
+ // (same constant value) for the fold to be valid, since we can only have
+ // one pad_value on the im2col.
+ if (im2colOp.hasPadding()) {
+ auto existingConst =
+ im2colOp.getPadValue().getDefiningOp<arith::ConstantOp>();
+ auto newConst = padValue.getDefiningOp<arith::ConstantOp>();
+ if (!existingConst || !newConst ||
+ existingConst.getValue() != newConst.getValue()) {
+ return rewriter.notifyMatchFailure(
+ im2colOp, "pad values are not compatible constants");
+ }
+ padValue = im2colOp.getPadValue();
+ }
+
+ Location loc = im2colOp.getLoc();
+ SmallVector<OpFoldResult> lowPad = padOp.getMixedLowPad();
+ SmallVector<OpFoldResult> highPad = padOp.getMixedHighPad();
+
+ // If im2col already has input padding, compose by adding element-wise.
+ SmallVector<OpFoldResult> existingLow = im2colOp.getMixedInputPadLow();
+ SmallVector<OpFoldResult> existingHigh = im2colOp.getMixedInputPadHigh();
+ if (!existingLow.empty()) {
+ for (auto [i, e] : llvm::enumerate(existingLow)) {
+ lowPad[i] = addOfrs(rewriter, loc, e, lowPad[i]);
+ }
+ for (auto [i, e] : llvm::enumerate(existingHigh)) {
+ highPad[i] = addOfrs(rewriter, loc, e, highPad[i]);
+ }
+ }
+
+ auto newIm2col = Im2colOp::create(
+ rewriter, loc, padOp.getSource(), im2colOp.getOutput(),
+ im2colOp.getStrides(), im2colOp.getDilations(),
+ im2colOp.getMixedKernelSize(), im2colOp.getMixedOffsets(),
+ im2colOp.getMixedOutputSizes(), im2colOp.getBatchPos(),
+ im2colOp.getMPos(), im2colOp.getKPos(), im2colOp.getInputKPerm(),
+ im2colOp.getOutputPerm(), lowPad, highPad,
+ im2colOp.getMixedOutputPadLow(), im2colOp.getMixedOutputPadHigh(),
+ padValue);
+
+ rewriter.replaceOp(im2colOp, newIm2col->getResults());
+ return success();
+ }
+};
+
+/// Fold tensor.pad on the output of im2col into the im2col by expanding the
+/// output tensor.
+///
+/// %out = tensor.empty(...)
+/// %result = im2col ins(%input) outs(%out) ...
+/// %padded = tensor.pad %result low[...] high[...] { yield %cst }
+/// -->
+/// %bigger_out = tensor.empty(padded_shape)
+/// %result = im2col ins(%input) outs(%bigger_out) ...
+struct FoldOutputPadIntoIm2col final : public OpRewritePattern<tensor::PadOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tensor::PadOp padOp,
+ PatternRewriter &rewriter) const override {
+ auto im2colOp = padOp.getSource().getDefiningOp<Im2colOp>();
+ if (!im2colOp) {
+ return rewriter.notifyMatchFailure(padOp,
+ "source not produced by im2col");
+ }
+
+ // The im2col must have a single use (this pad).
+ if (!im2colOp->hasOneUse()) {
+ return rewriter.notifyMatchFailure(padOp, "im2col has multiple uses");
+ }
+
+ // The im2col output must come from a tensor.empty.
+ auto emptyOp = im2colOp.getOutput().getDefiningOp<tensor::EmptyOp>();
+ if (!emptyOp) {
+ return rewriter.notifyMatchFailure(
+ padOp, "im2col output not produced by tensor.empty");
+ }
+
+ // Only fold constant padding values.
+ Value padValue = padOp.getConstantPaddingValue();
+ if (!padValue) {
+ return rewriter.notifyMatchFailure(padOp, "pad value is not a constant");
+ }
+
+ // The padding value must be compatible with the im2col op's existing
+ // pad_value. If the im2col has no pad_value yet, adopt the pad's value.
+ // If it has one, the values must match.
+ if (im2colOp.hasPadding()) {
+ auto existingConst =
+ im2colOp.getPadValue().getDefiningOp<arith::ConstantOp>();
+ auto newConst = padValue.getDefiningOp<arith::ConstantOp>();
+ if (!existingConst || !newConst ||
+ existingConst.getValue() != newConst.getValue()) {
+ return rewriter.notifyMatchFailure(
+ padOp, "pad values are not compatible constants");
+ }
+ padValue = im2colOp.getPadValue();
+ }
+
+ Location loc = padOp.getLoc();
+ SmallVector<OpFoldResult> lowPad = padOp.getMixedLowPad();
+ SmallVector<OpFoldResult> highPad = padOp.getMixedHighPad();
+
+ // This fold is safe because the pad_value is verified to be the same
+ // constant above, so padded positions in the larger output match what
+ // the downstream consumer (e.g., GEMM) expects.
+ auto outputType = cast<RankedTensorType>(padOp.getResultType());
+ int64_t outputRank = outputType.getRank();
+
+ SmallVector<OpFoldResult> newOutputShape;
+ SmallVector<OpFoldResult> oldOutputSizes =
+ tensor::getMixedSizes(rewriter, loc, im2colOp.getOutput());
+ AffineExpr d0, d1, d2;
+ bindDims(rewriter.getContext(), d0, d1, d2);
+ for (int64_t i = 0; i < outputRank; ++i) {
+ newOutputShape.push_back(affine::makeComposedFoldedAffineApply(
+ rewriter, loc, d0 + d1 + d2,
+ {oldOutputSizes[i], lowPad[i], highPad[i]}));
+ }
+
+ auto newEmptyOp = tensor::EmptyOp::create(rewriter, loc, newOutputShape,
+ outputType.getElementType());
+
+ // Compose output padding from the pad op with any existing output padding.
+ SmallVector<OpFoldResult> existingOutPadLow =
+ im2colOp.getMixedOutputPadLow();
+ SmallVector<OpFoldResult> existingOutPadHigh =
+ im2colOp.getMixedOutputPadHigh();
+ SmallVector<OpFoldResult> newOutPadLow(lowPad);
+ SmallVector<OpFoldResult> newOutPadHigh(highPad);
+ if (!existingOutPadLow.empty()) {
+ for (int64_t i = 0; i < outputRank; ++i) {
+ newOutPadLow[i] =
+ addOfrs(rewriter, loc, existingOutPadLow[i], newOutPadLow[i]);
+ newOutPadHigh[i] =
+ addOfrs(rewriter, loc, existingOutPadHigh[i], newOutPadHigh[i]);
+ }
+ }
+
+ auto newIm2col = Im2colOp::create(
+ rewriter, loc, im2colOp.getInput(), newEmptyOp.getResult(),
+ im2colOp.getStrides(), im2colOp.getDilations(),
+ im2colOp.getMixedKernelSize(), im2colOp.getMixedOffsets(),
+ im2colOp.getMixedOutputSizes(), im2colOp.getBatchPos(),
+ im2colOp.getMPos(), im2colOp.getKPos(), im2colOp.getInputKPerm(),
+ im2colOp.getOutputPerm(), im2colOp.getMixedInputPadLow(),
+ im2colOp.getMixedInputPadHigh(), newOutPadLow, newOutPadHigh, padValue);
+
+ rewriter.replaceOp(padOp, newIm2col->getResults());
+ return success();
+ }
+};
+
+} // namespace
+
+void Im2colOp::getCanonicalizationPatterns(RewritePatternSet &results,
+ MLIRContext *context) {
+ results.add<FoldInputPadIntoIm2col, FoldOutputPadIntoIm2col>(context);
+}
+
LogicalResult Im2colOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) {
return memref::foldMemRefCast(*this);
}
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td
index 6622b0b..aeff871 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td
@@ -1417,6 +1417,7 @@
let results = (outs Variadic<AnyShaped>:$results);
let hasFolder = 1;
+ let hasCanonicalizer = 1;
let assemblyFormat = [{
attr-dict
`strides` `=` $strides
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/canonicalize.mlir b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/canonicalize.mlir
index 8eea306..6999992 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/canonicalize.mlir
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/canonicalize.mlir
@@ -249,3 +249,234 @@
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CHECK: iree_linalg_ext.map_load
// CHECK-NOT: linalg.copy
+
+// -----
+
+// Test: fold input tensor.pad into im2col padding attributes.
+
+func.func @fold_input_pad_into_im2col(%arg0: tensor<2x34x34x640xf32>) -> tensor<2x1296x5760xf32> {
+ %cst = arith.constant 0.0 : f32
+ %padded = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0] {
+ ^bb0(%b0: index, %b1: index, %b2: index, %b3: index):
+ tensor.yield %cst : f32
+ } : tensor<2x34x34x640xf32> to tensor<2x36x36x640xf32>
+ %empty = tensor.empty() : tensor<2x1296x5760xf32>
+ %result = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, 0, 0] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2]
+ output_perm = [0, 1, 2]
+ ins(%padded : tensor<2x36x36x640xf32>)
+ outs(%empty : tensor<2x1296x5760xf32>) -> tensor<2x1296x5760xf32>
+ return %result : tensor<2x1296x5760xf32>
+}
+// CHECK-LABEL: func.func @fold_input_pad_into_im2col(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<2x34x34x640xf32>
+// CHECK-DAG: %[[CST:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<2x1296x5760xf32>
+// CHECK: %[[IM2COL:.+]] = iree_linalg_ext.im2col
+// CHECK-SAME: input_pad_low = [0, 1, 1, 0] input_pad_high = [0, 1, 1, 0] pad_value(%[[CST]] : f32)
+// CHECK-SAME: ins(%[[ARG0]] : tensor<2x34x34x640xf32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<2x1296x5760xf32>)
+// CHECK-NOT: tensor.pad
+// CHECK: return %[[IM2COL]]
+
+// -----
+
+// Test: fold output tensor.pad into im2col by expanding the output tensor.
+
+func.func @fold_output_pad_into_im2col(%arg0: tensor<2x34x34x640xf32>) -> tensor<2x1040x5760xf32> {
+ %cst = arith.constant 0.0 : f32
+ %empty = tensor.empty() : tensor<2x1024x5760xf32>
+ %result = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, 0, 0] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2]
+ output_perm = [0, 1, 2]
+ input_pad_low = [0, 1, 1, 0] input_pad_high = [0, 1, 1, 0] pad_value(%cst : f32)
+ ins(%arg0 : tensor<2x34x34x640xf32>)
+ outs(%empty : tensor<2x1024x5760xf32>) -> tensor<2x1024x5760xf32>
+ %padded = tensor.pad %result low[0, 0, 0] high[0, 16, 0] {
+ ^bb0(%b0: index, %b1: index, %b2: index):
+ tensor.yield %cst : f32
+ } : tensor<2x1024x5760xf32> to tensor<2x1040x5760xf32>
+ return %padded : tensor<2x1040x5760xf32>
+}
+// CHECK-LABEL: func.func @fold_output_pad_into_im2col(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<2x34x34x640xf32>
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<2x1040x5760xf32>
+// CHECK: %[[IM2COL:.+]] = iree_linalg_ext.im2col
+// CHECK-SAME: output_pad_low = [0, 0, 0] output_pad_high = [0, 16, 0]
+// CHECK-SAME: ins(%[[ARG0]] : tensor<2x34x34x640xf32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<2x1040x5760xf32>)
+// CHECK-NOT: tensor.pad
+// CHECK: return %[[IM2COL]]
+
+// -----
+
+// Test: both input+output pads with incompatible pad values — neither folds.
+
+func.func @no_fold_both_pads_incompatible(%arg0: tensor<2x34x34x640xf32>) -> tensor<2x1040x5760xf32> {
+ %cst_pad = arith.constant 1.0 : f32
+ %cst_im2col = arith.constant 0.0 : f32
+ %padded = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0] {
+ ^bb0(%b0: index, %b1: index, %b2: index, %b3: index):
+ tensor.yield %cst_pad : f32
+ } : tensor<2x34x34x640xf32> to tensor<2x36x36x640xf32>
+ %empty = tensor.empty() : tensor<2x1024x5760xf32>
+ %result = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, 0, 0] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2]
+ output_perm = [0, 1, 2]
+ input_pad_low = [0, 0, 0, 0] input_pad_high = [0, 0, 0, 0] pad_value(%cst_im2col : f32)
+ ins(%padded : tensor<2x36x36x640xf32>)
+ outs(%empty : tensor<2x1024x5760xf32>) -> tensor<2x1024x5760xf32>
+ %padded_output = tensor.pad %result low[0, 0, 0] high[0, 16, 0] {
+ ^bb0(%b0: index, %b1: index, %b2: index):
+ tensor.yield %cst_pad : f32
+ } : tensor<2x1024x5760xf32> to tensor<2x1040x5760xf32>
+ return %padded_output : tensor<2x1040x5760xf32>
+}
+// CHECK-LABEL: func.func @no_fold_both_pads_incompatible(
+// CHECK: tensor.pad
+// CHECK: iree_linalg_ext.im2col
+// CHECK: tensor.pad
+
+// -----
+
+// Test: both input+output pads with compatible pad values — both fold.
+
+func.func @fold_both_pads_compatible(%arg0: tensor<2x34x34x640xf32>) -> tensor<2x1312x5760xf32> {
+ %cst = arith.constant 0.0 : f32
+ %padded_input = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0] {
+ ^bb0(%b0: index, %b1: index, %b2: index, %b3: index):
+ tensor.yield %cst : f32
+ } : tensor<2x34x34x640xf32> to tensor<2x36x36x640xf32>
+ %empty = tensor.empty() : tensor<2x1296x5760xf32>
+ %result = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, 0, 0] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2]
+ output_perm = [0, 1, 2]
+ ins(%padded_input : tensor<2x36x36x640xf32>)
+ outs(%empty : tensor<2x1296x5760xf32>) -> tensor<2x1296x5760xf32>
+ %padded_output = tensor.pad %result low[0, 0, 0] high[0, 16, 0] {
+ ^bb0(%b0: index, %b1: index, %b2: index):
+ tensor.yield %cst : f32
+ } : tensor<2x1296x5760xf32> to tensor<2x1312x5760xf32>
+ return %padded_output : tensor<2x1312x5760xf32>
+}
+// CHECK-LABEL: func.func @fold_both_pads_compatible(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<2x34x34x640xf32>
+// CHECK-DAG: %[[CST:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<2x1312x5760xf32>
+// CHECK: %[[IM2COL:.+]] = iree_linalg_ext.im2col
+// CHECK-SAME: input_pad_low = [0, 1, 1, 0] input_pad_high = [0, 1, 1, 0]
+// CHECK-SAME: output_pad_low = [0, 0, 0] output_pad_high = [0, 16, 0]
+// CHECK-SAME: pad_value(%[[CST]] : f32)
+// CHECK-SAME: ins(%[[ARG0]] : tensor<2x34x34x640xf32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<2x1312x5760xf32>)
+// CHECK-NOT: tensor.pad
+// CHECK: return %[[IM2COL]]
+
+// -----
+
+// Test: fold input tensor.pad into im2col that already has input padding
+// (composing by adding element-wise).
+
+func.func @fold_input_pad_compose(%arg0: tensor<2x34x34x640xf32>) -> tensor<2x1296x5760xf32> {
+ %cst = arith.constant 0.0 : f32
+ %padded = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0] {
+ ^bb0(%b0: index, %b1: index, %b2: index, %b3: index):
+ tensor.yield %cst : f32
+ } : tensor<2x34x34x640xf32> to tensor<2x36x36x640xf32>
+ %empty = tensor.empty() : tensor<2x1296x5760xf32>
+ %result = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, 0, 0] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2]
+ output_perm = [0, 1, 2]
+ input_pad_low = [0, 1, 1, 0] input_pad_high = [0, 1, 1, 0] pad_value(%cst : f32)
+ ins(%padded : tensor<2x36x36x640xf32>)
+ outs(%empty : tensor<2x1296x5760xf32>) -> tensor<2x1296x5760xf32>
+ return %result : tensor<2x1296x5760xf32>
+}
+// CHECK-LABEL: func.func @fold_input_pad_compose(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<2x34x34x640xf32>
+// CHECK-DAG: %[[CST:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[IM2COL:.+]] = iree_linalg_ext.im2col
+// CHECK-SAME: input_pad_low = [0, 2, 2, 0] input_pad_high = [0, 2, 2, 0] pad_value(%[[CST]] : f32)
+// CHECK-SAME: ins(%[[ARG0]] : tensor<2x34x34x640xf32>)
+// CHECK-NOT: tensor.pad
+// CHECK: return %[[IM2COL]]
+
+// -----
+
+// Test: fold output tensor.pad with non-zero low padding into im2col.
+
+func.func @fold_output_pad_low(%arg0: tensor<2x34x34x640xf32>) -> tensor<2x1040x5760xf32> {
+ %cst = arith.constant 0.0 : f32
+ %empty = tensor.empty() : tensor<2x1024x5760xf32>
+ %result = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, 0, 0] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2]
+ output_perm = [0, 1, 2]
+ input_pad_low = [0, 1, 1, 0] input_pad_high = [0, 1, 1, 0] pad_value(%cst : f32)
+ ins(%arg0 : tensor<2x34x34x640xf32>)
+ outs(%empty : tensor<2x1024x5760xf32>) -> tensor<2x1024x5760xf32>
+ %padded = tensor.pad %result low[0, 16, 0] high[0, 0, 0] {
+ ^bb0(%b0: index, %b1: index, %b2: index):
+ tensor.yield %cst : f32
+ } : tensor<2x1024x5760xf32> to tensor<2x1040x5760xf32>
+ return %padded : tensor<2x1040x5760xf32>
+}
+// CHECK-LABEL: func.func @fold_output_pad_low(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<2x34x34x640xf32>
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<2x1040x5760xf32>
+// CHECK: %[[IM2COL:.+]] = iree_linalg_ext.im2col
+// CHECK-SAME: output_pad_low = [0, 16, 0] output_pad_high = [0, 0, 0]
+// CHECK-SAME: ins(%[[ARG0]] : tensor<2x34x34x640xf32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<2x1040x5760xf32>)
+// CHECK-NOT: tensor.pad
+// CHECK: return %[[IM2COL]]
+
+// -----
+
+// Test: fold output tensor.pad composing with existing output padding.
+
+func.func @fold_output_pad_compose(%arg0: tensor<2x34x34x640xf32>) -> tensor<2x1056x5760xf32> {
+ %cst = arith.constant 0.0 : f32
+ %empty = tensor.empty() : tensor<2x1040x5760xf32>
+ %result = iree_linalg_ext.im2col
+ strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+ offsets = [0, 0, 0] output_sizes = [[2], [32, 32], [3, 3, 640]]
+ batch_pos = [0] m_pos = [1, 2] k_pos = [3]
+ input_k_perm = [0, 1, 2]
+ output_perm = [0, 1, 2]
+ output_pad_low = [0, 0, 0] output_pad_high = [0, 16, 0]
+ pad_value(%cst : f32)
+ ins(%arg0 : tensor<2x34x34x640xf32>)
+ outs(%empty : tensor<2x1040x5760xf32>) -> tensor<2x1040x5760xf32>
+ %padded = tensor.pad %result low[0, 0, 0] high[0, 16, 0] {
+ ^bb0(%b0: index, %b1: index, %b2: index):
+ tensor.yield %cst : f32
+ } : tensor<2x1040x5760xf32> to tensor<2x1056x5760xf32>
+ return %padded : tensor<2x1056x5760xf32>
+}
+// CHECK-LABEL: func.func @fold_output_pad_compose(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<2x34x34x640xf32>
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<2x1056x5760xf32>
+// CHECK: %[[IM2COL:.+]] = iree_linalg_ext.im2col
+// CHECK-SAME: output_pad_low = [0, 0, 0] output_pad_high = [0, 32, 0]
+// CHECK-SAME: ins(%[[ARG0]] : tensor<2x34x34x640xf32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<2x1056x5760xf32>)
+// CHECK-NOT: tensor.pad
+// CHECK: return %[[IM2COL]]
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/conv_to_im2col.mlir b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/conv_to_im2col.mlir
index c7da669..4d0fee7 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/conv_to_im2col.mlir
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/conv_to_im2col.mlir
@@ -521,8 +521,8 @@
// CHECK-SAME: ins({{.*}} : tensor<61x93x16x64xbf16>)
// CHECK-SAME: outs({{.*}} : tensor<3x3x5369x16x64xbf16>) -> tensor<3x3x5369x16x64xbf16>
// CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1], [2], [3]] : tensor<59x91x16x56xbf16> into tensor<5369x16x56xbf16>
-// CHECK: linalg.generic
-// CHECK: util.return
+// CHECK: %[[MATMUL:.+]] = linalg.generic
+// CHECK: util.return %[[MATMUL]]
// -----
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/decompose_im2col.mlir b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/decompose_im2col.mlir
index 052e64d..36a8b6f 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/decompose_im2col.mlir
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/decompose_im2col.mlir
@@ -1,9 +1,7 @@
// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-linalg-ext-decompose-im2col{unroll=false}, canonicalize, cse))" --split-input-file %s | FileCheck %s
// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-linalg-ext-decompose-im2col{unroll=true}))" --split-input-file %s | FileCheck %s --check-prefix=CHECK-UNROLL
-// Test 1: Dynamic M offset and K offset.
-// The decomposition produces two nested scf.for loops (batch, M), computes
-// spatial coords via affine.apply, then extract_slice + linalg.copy.
+// Dynamic M offset and K offset -- non-padded: extract_slice + linalg.copy.
#map = affine_map<(d0) -> (d0 * 4)>
module {
func.func @im2col_untile_k(%arg0: tensor<2x34x34x640xf32>, %m_size: index, %m_off: index, %k: index) -> tensor<2x?x4xf32> {
@@ -33,31 +31,25 @@
// CHECK: %[[bLOOP:.+]] = scf.for %[[b:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[OUT0:.+]] = %[[OUT_TILE]]) -> (tensor<2x?x4xf32>)
// CHECK: %[[mLOOP:.+]] = scf.for %[[m:.+]] = %[[C0]] to %[[mSIZE]] step %[[C1]] iter_args(%[[OUT1:.+]] = %[[OUT0]]) -> (tensor<2x?x4xf32>)
// CHECK-DAG: %[[kScaled:.+]] = affine.apply #[[$MAP]]()[%[[K]]]
-// CHECK-DAG: %[[kParts:.+]]:3 = affine.delinearize_index %[[kScaled]] into (3, 3, 640)
+// CHECK-DAG: %[[kParts:.+]]:3 = affine.delinearize_index %[[kScaled]] into (3, 3, 640) : index, index, index
// CHECK-DAG: %[[mIDX:.+]] = affine.apply #[[$MAP1]](%[[m]])[%[[mOFF]]]
-// CHECK-DAG: %[[mParts:.+]]:2 = affine.delinearize_index %[[mIDX]] into (32, 32)
+// CHECK-DAG: %[[mParts:.+]]:2 = affine.delinearize_index %[[mIDX]] into (32, 32) : index, index
// CHECK-DAG: %[[h:.+]] = affine.apply #[[$MAP1]](%[[mParts]]#0)[%[[kParts]]#0]
// CHECK-DAG: %[[w:.+]] = affine.apply #[[$MAP1]](%[[mParts]]#1)[%[[kParts]]#1]
-// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[b]], %[[h]], %[[w]], %[[kParts]]#2] [1, 1, 1, 4] [1, 1, 1, 1] : tensor<2x34x34x640xf32> to tensor<1x1x4xf32>
-// CHECK: %[[OUT_SLICE:.+]] = tensor.extract_slice %[[OUT1]][%[[b]], %[[m]], 0] [1, 1, 4] [1, 1, 1] : tensor<2x?x4xf32> to tensor<1x1x4xf32>
-// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[IN_SLICE]] : tensor<1x1x4xf32>) outs(%[[OUT_SLICE]] : tensor<1x1x4xf32>) -> tensor<1x1x4xf32>
-// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[COPY]] into %[[OUT1]][%[[b]], %[[m]], 0] [1, 1, 4] [1, 1, 1] : tensor<1x1x4xf32> into tensor<2x?x4xf32>
-// CHECK: scf.yield %[[INSERT]] : tensor<2x?x4xf32>
+// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[b]], %[[h]], %[[w]], %[[kParts]]#2] [1, 1, 1, 4]
+// CHECK: linalg.copy ins(%[[IN_SLICE]]
+// CHECK: tensor.insert_slice {{.*}} into %[[OUT1]][%[[b]], %[[m]], 0] [1, 1, 4] [1, 1, 1] : tensor<4xf32> into tensor<2x?x4xf32>
+// CHECK: scf.yield {{.*}} : tensor<2x?x4xf32>
// CHECK: scf.yield %[[mLOOP]] : tensor<2x?x4xf32>
// CHECK: return %[[bLOOP]] : tensor<2x?x4xf32>
+// CHECK-NOT: affine.max
+// CHECK-NOT: affine.min
// CHECK-UNROLL-LABEL: func.func @im2col_untile_k
// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
-// Verify two unrolled batch loops (b=0 and b=1), each iterating over M.
-// CHECK-UNROLL: scf.for %{{.+}} = %{{.+}} to %[[mS:.+]] step
-// CHECK-UNROLL: linalg.copy
-// CHECK-UNROLL: scf.for %{{.+}} = %{{.+}} to %[[mS]] step
-// CHECK-UNROLL: linalg.copy
// -----
-// Test 2: Dynamic M and K offsets with transposed m_pos.
-// Three nested loops (batch, M, K). Spatial coords are computed via affine.apply
-// using strides and dilations.
+// Dynamic M and K offsets with transposed m_pos -- non-padded: linalg.copy.
module {
func.func @im2col_transposed_m_pos(%arg0: tensor<640x2x101x172xf32>, %m_size: index, %k_size: index, %m_off: index, %k_off: index) -> tensor<2x?x?xf32> {
%c2 = arith.constant 2 : index
@@ -91,32 +83,26 @@
// CHECK: %[[mLOOP:.+]] = scf.for %[[m:.+]] = %[[C0]] to %[[mSIZE]] step %[[C1]] iter_args(%[[OUT1:.+]] = %[[OUT0]]) -> (tensor<2x?x?xf32>)
// CHECK: %[[kLOOP:.+]] = scf.for %[[k:.+]] = %[[C0]] to %[[kSIZE]] step %[[C1]] iter_args(%[[OUT2:.+]] = %[[OUT1]]) -> (tensor<2x?x?xf32>)
// CHECK-DAG: %[[kIDX:.+]] = affine.apply #[[$MAP]](%[[k]])[%[[kOFF]]]
-// CHECK-DAG: %[[kParts:.+]]:3 = affine.delinearize_index %[[kIDX]] into (640, 5, 2)
+// CHECK-DAG: %[[kParts:.+]]:3 = affine.delinearize_index %[[kIDX]] into (640, 5, 2) : index, index, index
// CHECK-DAG: %[[mIDX:.+]] = affine.apply #[[$MAP]](%[[m]])[%[[mOFF]]]
-// CHECK-DAG: %[[mParts:.+]]:2 = affine.delinearize_index %[[mIDX]] into (32, 32)
-// CHECK-DAG: affine.apply #[[$MAP1]](%[[mParts]]#0, %[[kParts]]#1)
-// CHECK-DAG: affine.apply #[[$MAP2]](%[[mParts]]#1, %[[kParts]]#2)
-// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[kParts]]#0, %[[b]], {{.*}}, {{.*}}] [1, 1, 1, 1]
-// CHECK: %[[OUT_SLICE:.+]] = tensor.extract_slice %[[OUT2]][%[[b]], %[[m]], %[[k]]] [1, 1, 1] [1, 1, 1] : tensor<2x?x?xf32> to tensor<1x1x1xf32>
-// CHECK: linalg.copy ins(%[[IN_SLICE]] : tensor<1x1x1xf32>) outs(%[[OUT_SLICE]] : tensor<1x1x1xf32>)
-// CHECK: tensor.insert_slice {{.*}} into %[[OUT2]][%[[b]], %[[m]], %[[k]]] [1, 1, 1] [1, 1, 1] : tensor<1x1x1xf32> into tensor<2x?x?xf32>
+// CHECK-DAG: %[[mParts:.+]]:2 = affine.delinearize_index %[[mIDX]] into (32, 32) : index, index
+// CHECK-DAG: %[[h:.+]] = affine.apply #[[$MAP1]](%[[mParts]]#0, %[[kParts]]#1)
+// CHECK-DAG: %[[w:.+]] = affine.apply #[[$MAP2]](%[[mParts]]#1, %[[kParts]]#2)
+// CHECK: %[[IN_SLICE2:.+]] = tensor.extract_slice %[[ARG0]][%[[kParts]]#0, %[[b]], %[[w]], %[[h]]] [1, 1, 1, 1]
+// CHECK: linalg.copy ins(%[[IN_SLICE2]]
+// CHECK: tensor.insert_slice {{.*}} into %[[OUT2]][%[[b]], %[[m]], %[[k]]] [1, 1, 1] [1, 1, 1] : tensor<1xf32> into tensor<2x?x?xf32>
// CHECK: scf.yield {{.*}} : tensor<2x?x?xf32>
// CHECK: scf.yield %[[kLOOP]] : tensor<2x?x?xf32>
// CHECK: scf.yield %[[mLOOP]] : tensor<2x?x?xf32>
// CHECK: return %[[bLOOP]] : tensor<2x?x?xf32>
-
+// CHECK-NOT: affine.max
+// CHECK-NOT: affine.min
// CHECK-UNROLL-LABEL: func.func @im2col_transposed_m_pos
// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
-// Two unrolled copies of the m/k loops for batch=0 and batch=1.
-// CHECK-UNROLL: scf.for
-// CHECK-UNROLL: linalg.copy
-// CHECK-UNROLL: scf.for
-// CHECK-UNROLL: linalg.copy
// -----
-// Test 3: Static sizes with expanded M and K output dims.
-// Four nested loops (batch, M0, M1, K). Spatial coords via affine.apply.
+// Static sizes -- non-padded: extract_slice + linalg.copy + insert_slice.
module {
func.func @im2col_expanded(%arg0: tensor<2x34x34x640xf32>, %m_size0: index, %m_size1: index, %m0: index, %m1: index, %k: index, %m_stride: index) -> tensor<2x?x?x2x4xf32> {
%0 = tensor.empty(%m_size0, %m_size1) : tensor<2x?x?x2x4xf32>
@@ -139,6 +125,7 @@
// CHECK-SAME: %[[mOFF0:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[mOFF1:[a-zA-Z0-9_]+]]
// CHECK-SAME: %[[kOFF:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[mSTRIDE:[a-zA-Z0-9_]+]]
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
@@ -147,32 +134,27 @@
// CHECK: %[[mLOOP0:.+]] = scf.for %[[m0:.+]] = %[[C0]] to %[[mSIZE0]] step %[[C1]] iter_args(%[[OUT1:.+]] = %[[OUT0]]) -> (tensor<2x?x?x2x4xf32>)
// CHECK: %[[mLOOP1:.+]] = scf.for %[[m1:.+]] = %[[C0]] to %[[mSIZE1]] step %[[C1]] iter_args(%[[OUT2:.+]] = %[[OUT1]]) -> (tensor<2x?x?x2x4xf32>)
// CHECK: %[[kLOOP:.+]] = scf.for %[[k:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[OUT3:.+]] = %[[OUT2]]) -> (tensor<2x?x?x2x4xf32>)
-// CHECK-DAG: %[[kPOS:.+]] = affine.apply #[[$MAP]](%[[k]])[%[[kOFF]]]
-// CHECK-DAG: %[[kParts:.+]]:2 = affine.delinearize_index %[[kPOS]] into (3, 3)
-// CHECK-DAG: affine.apply #[[$MAP1]](%[[kParts]]#0, %[[m0]])[%[[mOFF0]]]
-// CHECK-DAG: affine.apply #[[$MAP1]](%[[kParts]]#1, %[[m1]])[%[[mOFF1]]]
-// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[b]], {{.*}}, {{.*}}, 0] [1, 1, 1, 4] [1, 1, 1, 1] : tensor<2x34x34x640xf32> to tensor<1x1x1x4xf32>
-// CHECK: %[[OUT_SLICE:.+]] = tensor.extract_slice %[[OUT3]][%[[b]], %[[m0]], %[[m1]], %[[k]], 0] [1, 1, 1, 1, 4] [1, 1, 1, 1, 1] : tensor<2x?x?x2x4xf32> to tensor<1x1x1x4xf32>
-// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[IN_SLICE]] : tensor<1x1x1x4xf32>) outs(%[[OUT_SLICE]] : tensor<1x1x1x4xf32>) -> tensor<1x1x1x4xf32>
-// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[COPY]] into %[[OUT3]][%[[b]], %[[m0]], %[[m1]], %[[k]], 0] [1, 1, 1, 1, 4] [1, 1, 1, 1, 1] : tensor<1x1x1x4xf32> into tensor<2x?x?x2x4xf32>
+// CHECK-DAG: %[[kIDX:.+]] = affine.apply #[[$MAP]](%[[k]])[%[[kOFF]]]
+// CHECK-DAG: %[[kParts:.+]]:2 = affine.delinearize_index %[[kIDX]] into (3, 3) : index, index
+// CHECK-DAG: %[[h:.+]] = affine.apply #[[$MAP1]](%[[kParts]]#0, %[[m0]])[%[[mOFF0]]]
+// CHECK-DAG: %[[w:.+]] = affine.apply #[[$MAP1]](%[[kParts]]#1, %[[m1]])[%[[mOFF1]]]
+// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[b]], %[[h]], %[[w]], 0] [1, 1, 1, 4] [1, 1, 1, 1] : tensor<2x34x34x640xf32> to tensor<4xf32>
+// CHECK: %[[DEST_SLICE:.+]] = tensor.extract_slice %[[OUT3]][%[[b]], %[[m0]], %[[m1]], %[[k]], 0] [1, 1, 1, 1, 4] {{.*}} : tensor<2x?x?x2x4xf32> to tensor<4xf32>
+// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[IN_SLICE]] : tensor<4xf32>) outs(%[[DEST_SLICE]] : tensor<4xf32>)
+// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[COPY]] into %[[OUT3]][%[[b]], %[[m0]], %[[m1]], %[[k]], 0] [1, 1, 1, 1, 4] [1, 1, 1, 1, 1] : tensor<4xf32> into tensor<2x?x?x2x4xf32>
// CHECK: scf.yield %[[INSERT]] : tensor<2x?x?x2x4xf32>
// CHECK: scf.yield %[[kLOOP]] : tensor<2x?x?x2x4xf32>
// CHECK: scf.yield %[[mLOOP1]] : tensor<2x?x?x2x4xf32>
// CHECK: scf.yield %[[mLOOP0]] : tensor<2x?x?x2x4xf32>
// CHECK: return %[[bLOOP]] : tensor<2x?x?x2x4xf32>
-
+// CHECK-NOT: affine.max
+// CHECK-NOT: affine.min
// CHECK-UNROLL-LABEL: func.func @im2col_expanded
// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
-// Unrolled: batch loop is removed, m0/m1 loops remain; K=2 is unrolled in each body.
-// CHECK-UNROLL: scf.for
-// CHECK-UNROLL: scf.for
-// CHECK-UNROLL: linalg.copy
-// CHECK-UNROLL: linalg.copy
// -----
-// Test 4: NCHW layout with static sizes -- scalar fallback (34 % 4 != 0).
-// Three nested loops (batch, K-row, K-col-channel) with scalar linalg.copy.
+// NCHW layout with static sizes -- non-padded: extract_slice + linalg.copy + insert_slice.
module {
func.func @im2col_expanded_nchw(%arg0: tensor<2x640x34x34xf32>, %m0: index, %m1: index, %k: index) -> tensor<2x1x1x2x4xf32> {
%0 = tensor.empty() : tensor<2x1x1x2x4xf32>
@@ -186,24 +168,23 @@
return %7 : tensor<2x1x1x2x4xf32>
}
}
-// Scalar fallback (34 % 4 != 0): three nested loops with scalar linalg.copy.
+// Scalar fallback (34 % 4 != 0): loops over batch + K dims with linalg.copy.
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
// CHECK-LABEL: func.func @im2col_expanded_nchw
-// CHECK: tensor.empty() : tensor<2x1x1x2x4xf32>
-// CHECK: scf.for
-// CHECK: scf.for
-// CHECK: scf.for
-// CHECK: affine.delinearize_index {{.*}} into (3, 3)
-// CHECK: tensor.extract_slice {{.*}} : tensor<2x640x34x34xf32> to tensor<1x1x1x1xf32>
-// CHECK: linalg.copy ins({{.*}} : tensor<1x1x1x1xf32>) outs({{.*}} : tensor<1x1x1x1xf32>)
-// CHECK: tensor.insert_slice {{.*}} : tensor<1x1x1x1xf32> into tensor<2x1x1x2x4xf32>
-
+// CHECK: tensor.extract_slice {{.*}} : tensor<2x640x34x34xf32> to tensor<1xf32>
+// CHECK: linalg.copy
+// CHECK: tensor.insert_slice {{.*}} : tensor<1xf32> into tensor<2x1x1x2x4xf32>
+// CHECK-NOT: affine.max
+// CHECK-NOT: affine.min
// CHECK-UNROLL-LABEL: func.func @im2col_expanded_nchw
// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
// -----
-// Test 5: Backward-weight-style im2col with dilation=2, single-element M and K dims.
-// Directly extract_slices at (m0, m1) without any loops (output is 1x1x1x1).
+// Test backward-weight-style im2col where M dims are kernel spatial dims,
+// not convolution output spatial dims. With 2 expanded M output dims,
+// no delinearization is needed -- the M coords are used directly.
+// Non-padded: extract_slice + linalg.copy.
module {
func.func @im2col_bwd_weight_dilation2(%arg0: tensor<1x1x8x8xf32>, %m0: index, %m1: index) -> tensor<1x1x1x1xf32> {
%0 = tensor.empty() : tensor<1x1x1x1xf32>
@@ -217,17 +198,17 @@
return %result : tensor<1x1x1x1xf32>
}
}
-// With single-element output (1x1x1x1), emits a bare
-// extract_slice + linalg.copy without any loops.
+// With expanded M (2 single-element M output dims), M coords are used directly
+// without delinearization. Non-padded: extract_slice + linalg.copy.
// CHECK-LABEL: func.func @im2col_bwd_weight_dilation2
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x1x8x8xf32>
-// CHECK-SAME: %[[M0:[a-zA-Z0-9_]+]]: index
-// CHECK-SAME: %[[M1:[a-zA-Z0-9_]+]]: index
-// CHECK: tensor.extract_slice %[[ARG0]][0, 0, %[[M0]], %[[M1]]]
-// CHECK: linalg.copy
-// CHECK: return
-// CHECK-NOT: iree_linalg_ext.im2col
-
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x1x8x8xf32>
+// CHECK-SAME: %[[M0:[a-zA-Z0-9_]+]]: index
+// CHECK-SAME: %[[M1:[a-zA-Z0-9_]+]]: index
+// CHECK: tensor.extract_slice %[[ARG0]][0, 0, %[[M0]], %[[M1]]] [1, 1, 1, 1]
+// CHECK: linalg.copy
+// CHECK: return
+// CHECK-NOT: affine.max
+// CHECK-NOT: affine.min
// CHECK-UNROLL-LABEL: func.func @im2col_bwd_weight_dilation2
// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
// Fully unrolled 1x1x1x1 output: bare extract_slice + copy, no loops.
@@ -237,14 +218,11 @@
// -----
-// Test 6: Static sizes with dynamic M offset, with unrolling.
-// The unrolled pass unrolls the static batch (size 2) and M (size 2) dims into
-// separate extract_slice + linalg.copy + insert_slice blocks.
-#map6 = affine_map<(d0) -> (d0 * 4)>
+#map = affine_map<(d0) -> (d0 * 4)>
module {
func.func @im2col_unrolled(%arg0: tensor<2x34x34x640xf32>, %m_off: index, %k: index) -> tensor<2x2x4xf32> {
%0 = tensor.empty() : tensor<2x2x4xf32>
- %k_off = affine.apply #map6(%k)
+ %k_off = affine.apply #map(%k)
%7 = iree_linalg_ext.im2col
strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
offsets = [0, %m_off, %k_off] output_sizes = [[2], [32, 32], [3, 3, 640]]
@@ -255,46 +233,24 @@
return %7 : tensor<2x2x4xf32>
}
}
-// CHECK-LABEL: func.func @im2col_unrolled
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]
-// CHECK-SAME: %[[mOFF:[a-zA-Z0-9_]+]]
-// CHECK-SAME: %[[K:[a-zA-Z0-9_]+]]
-// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
-// CHECK: %[[OUT_TILE:.+]] = tensor.empty() : tensor<2x2x4xf32>
-// Non-unrolled output uses two nested loops (batch=2, M=2).
-// CHECK: %[[bLOOP:.+]] = scf.for %[[b:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[OUT0:.+]] = %[[OUT_TILE]]) -> (tensor<2x2x4xf32>)
-// CHECK: %[[mLOOP:.+]] = scf.for %[[m:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[OUT1:.+]] = %[[OUT0]]) -> (tensor<2x2x4xf32>)
-// CHECK: tensor.extract_slice %[[ARG0]]
-// CHECK: linalg.copy ins({{.*}} : tensor<1x1x4xf32>) outs({{.*}} : tensor<1x1x4xf32>)
-// CHECK: tensor.insert_slice {{.*}} into %[[OUT1]][%[[b]], %[[m]], 0] [1, 1, 4] [1, 1, 1] : tensor<1x1x4xf32> into tensor<2x2x4xf32>
-// CHECK: scf.yield {{.*}} : tensor<2x2x4xf32>
-// CHECK: scf.yield %[[mLOOP]] : tensor<2x2x4xf32>
-// CHECK: return %[[bLOOP]] : tensor<2x2x4xf32>
-// CHECK-NOT: iree_linalg_ext.im2col
-
// CHECK-UNROLL-LABEL: func.func @im2col_unrolled
-// CHECK-UNROLL-SAME: %[[ARG0:[a-zA-Z0-9_]+]]
-// CHECK-UNROLL-SAME: %[[mOFF:[a-zA-Z0-9_]+]]
-// CHECK-UNROLL-SAME: %[[K:[a-zA-Z0-9_]+]]
// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
// CHECK-UNROLL: %[[EMPTY:.*]] = tensor.empty() : tensor<2x2x4xf32>
-// Unrolled: 4 copies (b=0,m=0), (b=0,m=1), (b=1,m=0), (b=1,m=1).
-// CHECK-UNROLL: linalg.copy ins({{.*}} : tensor<1x1x4xf32>) outs({{.*}} : tensor<1x1x4xf32>)
+// CHECK-UNROLL: linalg.copy
// CHECK-UNROLL: %[[INS0:.*]] = tensor.insert_slice %{{.*}} into %[[EMPTY]]
-// CHECK-UNROLL: linalg.copy ins({{.*}} : tensor<1x1x4xf32>) outs({{.*}} : tensor<1x1x4xf32>)
+// CHECK-UNROLL: tensor.extract_slice %[[INS0]]
+// CHECK-UNROLL: linalg.copy
// CHECK-UNROLL: %[[INS1:.*]] = tensor.insert_slice %{{.*}} into %[[INS0]]
-// CHECK-UNROLL: linalg.copy ins({{.*}} : tensor<1x1x4xf32>) outs({{.*}} : tensor<1x1x4xf32>)
+// CHECK-UNROLL: tensor.extract_slice %[[INS1]]
+// CHECK-UNROLL: linalg.copy
// CHECK-UNROLL: %[[INS2:.*]] = tensor.insert_slice %{{.*}} into %[[INS1]]
-// CHECK-UNROLL: linalg.copy ins({{.*}} : tensor<1x1x4xf32>) outs({{.*}} : tensor<1x1x4xf32>)
+// CHECK-UNROLL: tensor.extract_slice %[[INS2]]
+// CHECK-UNROLL: linalg.copy
// CHECK-UNROLL: %[[INS3:.*]] = tensor.insert_slice %{{.*}} into %[[INS2]]
-// CHECK-UNROLL: return %[[INS3]] : tensor<2x2x4xf32>
+// CHECK-UNROLL: return %[[INS3]]
// -----
-// Test 7: im2col with pre-padded input (tensor.pad before im2col).
-// This uses the padding-aware code path which does have affine.max/affine.min ops.
module {
func.func @im2col_padding(%input: tensor<1x8x3x3xf32>) -> tensor<1x2x2x12xf32> {
%cst = arith.constant 0.000000e+00 : f32
@@ -325,15 +281,9 @@
// CHECK: tensor.insert_slice
// CHECK-NOT: iree_linalg_ext.im2col
-// CHECK-UNROLL-LABEL: func.func @im2col_padding
-// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
-// CHECK-UNROLL: tensor.pad
-// CHECK-UNROLL: linalg.copy
-
// -----
-// Test 8: Static sizes, NHWC layout with non-identity input_k_perm -- scalar fallback.
-// The non-identity k_perm prevents vectorization; two nested loops (M, K) with scalar copy.
+// Static sizes, NHWC layout with input_k_perm -- non-padded: linalg.copy.
module {
func.func @im2col_nhc_with_perm(%arg0: tensor<1x3x2xf32>) -> tensor<1x2x4xf32> {
%0 = tensor.empty() : tensor<1x2x4xf32>
@@ -357,20 +307,14 @@
// CHECK: %[[OUT_TILE:.+]] = tensor.empty() : tensor<1x2x4xf32>
// CHECK: %[[MLOOP:.+]] = scf.for %[[M:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[OUT1:.+]] = %[[OUT_TILE]]) -> (tensor<1x2x4xf32>)
// CHECK: %[[KLOOP:.+]] = scf.for %[[K:.+]] = %[[C0]] to %[[C4]] step %[[C1]] iter_args(%[[OUT2:.+]] = %[[OUT1]]) -> (tensor<1x2x4xf32>)
-// CHECK: %[[kParts:.+]]:2 = affine.delinearize_index %[[K]] into (2, 2) : index, index
-// CHECK: %[[hIDX:.+]] = affine.apply #[[$MAP]](%[[M]], %[[kParts]]#1)
-// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][0, %[[hIDX]], %[[kParts]]#0] [1, 1, 1] [1, 1, 1] : tensor<1x3x2xf32> to tensor<1x1x1xf32>
-// CHECK: %[[OUT_SLICE:.+]] = tensor.extract_slice %[[OUT2]][0, %[[M]], %[[K]]] [1, 1, 1] [1, 1, 1] : tensor<1x2x4xf32> to tensor<1x1x1xf32>
-// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[IN_SLICE]] : tensor<1x1x1xf32>) outs(%[[OUT_SLICE]] : tensor<1x1x1xf32>) -> tensor<1x1x1xf32>
-// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[COPY]] into %[[OUT2]][0, %[[M]], %[[K]]] [1, 1, 1] [1, 1, 1] : tensor<1x1x1xf32> into tensor<1x2x4xf32>
-// CHECK: scf.yield %[[INSERT]] : tensor<1x2x4xf32>
+// CHECK: affine.delinearize_index %[[K]] into (2, 2) : index, index
+// CHECK: tensor.extract_slice %[[ARG0]]
+// CHECK: linalg.copy
+// CHECK: tensor.insert_slice {{.*}} into %[[OUT2]][0, %[[M]], %[[K]]] [1, 1, 1] [1, 1, 1] : tensor<1xf32> into tensor<1x2x4xf32>
+// CHECK: scf.yield {{.*}} : tensor<1x2x4xf32>
// CHECK: scf.yield %[[KLOOP]] : tensor<1x2x4xf32>
// CHECK: return %[[MLOOP]] : tensor<1x2x4xf32>
-// CHECK-UNROLL-LABEL: func.func @im2col_nhc_with_perm
-// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
-// CHECK-UNROLL: linalg.copy
-
// -----
// Test 9: Non-self-inverse input_k_perm = [2, 0, 1].
@@ -389,28 +333,17 @@
}
}
// Scalar fallback with non-self-inverse kperm: loops over M0, M1, K.
-// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1) -> (d0 + d1)>
// CHECK-LABEL: func.func @im2col_nhwc_with_perm
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x16x16x4xf32>
-// CHECK-DAG: %[[C36:.+]] = arith.constant 36 : index
-// CHECK-DAG: %[[C14:.+]] = arith.constant 14 : index
-// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
-// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-// CHECK: %[[OUT_TILE:.+]] = tensor.empty() : tensor<1x14x14x36xf32>
-// CHECK: %[[MLOOP0:.+]] = scf.for %[[M0:.+]] = %[[C0]] to %[[C14]] step %[[C1]] iter_args(%[[OUT0:.+]] = %[[OUT_TILE]]) -> (tensor<1x14x14x36xf32>)
-// CHECK: %[[MLOOP1:.+]] = scf.for %[[M1:.+]] = %[[C0]] to %[[C14]] step %[[C1]] iter_args(%[[OUT1:.+]] = %[[OUT0]]) -> (tensor<1x14x14x36xf32>)
-// CHECK: %[[KLOOP:.+]] = scf.for %[[K:.+]] = %[[C0]] to %[[C36]] step %[[C1]] iter_args(%[[OUT2:.+]] = %[[OUT1]]) -> (tensor<1x14x14x36xf32>)
-// CHECK-DAG: %[[kParts:.+]]:3 = affine.delinearize_index %[[K]] into (4, 3, 3) : index, index, index
-// CHECK-DAG: %[[hIDX:.+]] = affine.apply #[[$MAP]](%[[M0]], %[[kParts]]#1)
-// CHECK-DAG: %[[wIDX:.+]] = affine.apply #[[$MAP]](%[[M1]], %[[kParts]]#2)
-// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][0, %[[hIDX]], %[[wIDX]], %[[kParts]]#0] [1, 1, 1, 1] [1, 1, 1, 1] : tensor<1x16x16x4xf32> to tensor<1x1x1x1xf32>
-// CHECK: %[[OUT_SLICE:.+]] = tensor.extract_slice %[[OUT2]][0, %[[M0]], %[[M1]], %[[K]]] [1, 1, 1, 1] [1, 1, 1, 1] : tensor<1x14x14x36xf32> to tensor<1x1x1x1xf32>
-// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[IN_SLICE]] : tensor<1x1x1x1xf32>) outs(%[[OUT_SLICE]] : tensor<1x1x1x1xf32>)
-// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[COPY]] into %[[OUT2]][0, %[[M0]], %[[M1]], %[[K]]] [1, 1, 1, 1] [1, 1, 1, 1] : tensor<1x1x1x1xf32> into tensor<1x14x14x36xf32>
-// CHECK: scf.yield %[[INSERT]] : tensor<1x14x14x36xf32>
-// CHECK: scf.yield %[[KLOOP]] : tensor<1x14x14x36xf32>
-// CHECK: scf.yield %[[MLOOP1]] : tensor<1x14x14x36xf32>
-// CHECK: return %[[MLOOP0]] : tensor<1x14x14x36xf32>
+// CHECK: tensor.empty() : tensor<1x14x14x36xf32>
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: affine.delinearize_index
+// CHECK: tensor.extract_slice %[[ARG0]]
+// CHECK: linalg.copy
+// CHECK: tensor.insert_slice
+// CHECK-NOT: iree_linalg_ext.im2col
// CHECK-UNROLL-LABEL: func.func @im2col_nhwc_with_perm
// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
@@ -418,11 +351,8 @@
// -----
-// Test 10: CHWN layout with batch dim as innermost.
-// Static sizes; three nested loops (M0, M1, K). Spatial coords via affine.apply.
-// The batch dim is innermost so extract_slice reads a 4-element slice (the full batch).
-// Because output_perm is [0,1,2,3] but batch_pos = [3], the output dimension order
-// differs from the input slice order, so a linalg.transpose is needed.
+// CHWN layout with batch dim as innermost.
+// Static sizes, non-padded: extract_slice + linalg.copy + insert_slice.
module {
func.func @im2col_chwn(%arg0: tensor<16x26x18x4xf32>, %arg1: index, %arg2: index, %arg3: index) -> tensor<4x2x2x2xf32> {
%0 = tensor.empty() : tensor<4x2x2x2xf32>
@@ -452,28 +382,24 @@
// CHECK: %[[mLOOP1:.+]] = scf.for %[[M1:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[OUT1:.+]] = %[[OUT0]])
// CHECK: %[[kLOOP:.+]] = scf.for %[[K:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[OUT2:.+]] = %[[OUT1]])
// CHECK-DAG: %[[kIDX:.+]] = affine.apply #[[$MAP]](%[[K]])[%[[ARG3]]]
-// CHECK-DAG: %[[kParts:.+]]:3 = affine.delinearize_index %[[kIDX]] into (16, 24, 16)
-// CHECK-DAG: affine.apply #[[$MAP1]](%[[kParts]]#1, %[[M0]])[%[[ARG1]]]
-// CHECK-DAG: affine.apply #[[$MAP1]](%[[kParts]]#2, %[[M1]])[%[[ARG2]]]
-// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[kParts]]#0, {{.*}}, {{.*}}, 0] [1, 1, 1, 4] [1, 1, 1, 1] : tensor<16x26x18x4xf32> to tensor<1x1x1x4xf32>
-// CHECK: %[[DEST_SLICE:.+]] = tensor.extract_slice %[[OUT2]][0, %[[M0]], %[[M1]], %[[K]]] [4, 1, 1, 1] [1, 1, 1, 1] : tensor<4x2x2x2xf32> to tensor<4x1x1x1xf32>
-// CHECK: %[[TRANS:.+]] = linalg.transpose ins(%[[IN_SLICE]] : tensor<1x1x1x4xf32>) outs(%[[DEST_SLICE]] : tensor<4x1x1x1xf32>) permutation = [3, 1, 2, 0]
-// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[TRANS]] into %[[OUT2]][0, %[[M0]], %[[M1]], %[[K]]] [4, 1, 1, 1] [1, 1, 1, 1] : tensor<4x1x1x1xf32> into tensor<4x2x2x2xf32>
+// CHECK-DAG: %[[kParts:.+]]:3 = affine.delinearize_index %[[kIDX]] into (16, 24, 16) : index, index, index
+// CHECK-DAG: %[[h:.+]] = affine.apply #[[$MAP1]](%[[kParts]]#1, %[[M0]])[%[[ARG1]]]
+// CHECK-DAG: %[[w:.+]] = affine.apply #[[$MAP1]](%[[kParts]]#2, %[[M1]])[%[[ARG2]]]
+// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[kParts]]#0, %[[h]], %[[w]], 0] [1, 1, 1, 4] [1, 1, 1, 1] : tensor<16x26x18x4xf32> to tensor<4xf32>
+// CHECK: %[[DEST_SLICE:.+]] = tensor.extract_slice %[[OUT2]][0, %[[M0]], %[[M1]], %[[K]]] [4, 1, 1, 1] {{.*}} : tensor<4x2x2x2xf32> to tensor<4xf32>
+// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[IN_SLICE]] : tensor<4xf32>) outs(%[[DEST_SLICE]] : tensor<4xf32>)
+// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[COPY]] into %[[OUT2]][0, %[[M0]], %[[M1]], %[[K]]] [4, 1, 1, 1] [1, 1, 1, 1] : tensor<4xf32> into tensor<4x2x2x2xf32>
// CHECK: scf.yield %[[INSERT]] : tensor<4x2x2x2xf32>
// CHECK: scf.yield %[[kLOOP]] : tensor<4x2x2x2xf32>
// CHECK: scf.yield %[[mLOOP1]] : tensor<4x2x2x2xf32>
// CHECK: return %[[mLOOP0:.+]] : tensor<4x2x2x2xf32>
-
-// CHECK-UNROLL-LABEL: func.func @im2col_chwn
-// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
+// CHECK-NOT: affine.max
+// CHECK-NOT: affine.min
// -----
-// Test 11: CHWN layout with output_perm that eliminates the batch transpose.
-// Static sizes; three nested loops (K, M0, M1) matching the permuted output.
-// Because output_perm = [3, 1, 2, 0], the batch and K positions are swapped
-// relative to the default [B, M0, M1, K]. No transpose is needed because
-// the output slice is arranged to match the input slice.
+// CHWN layout with output_perm that eliminates the transpose.
+// Static sizes, non-padded: extract_slice + linalg.copy + insert_slice (no transpose).
module {
func.func @im2col_chwn_output_perm(%arg0: tensor<16x26x18x4xf32>, %arg1: index, %arg2: index, %arg3: index) -> tensor<2x2x2x4xf32> {
%0 = tensor.empty() : tensor<2x2x2x4xf32>
@@ -503,26 +429,23 @@
// CHECK: %[[LOOP1:.+]] = scf.for %[[IV1:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG5:.+]] = %[[ARG4]])
// CHECK: %[[LOOP2:.+]] = scf.for %[[IV2:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG6:.+]] = %[[ARG5]])
// CHECK-DAG: %[[kIDX:.+]] = affine.apply #[[$MAP]](%[[IV0]])[%[[ARG3]]]
-// CHECK-DAG: %[[kParts:.+]]:3 = affine.delinearize_index %[[kIDX]] into (16, 24, 16)
-// CHECK-DAG: affine.apply #[[$MAP1]](%[[kParts]]#1, %[[IV1]])[%[[ARG1]]]
-// CHECK-DAG: affine.apply #[[$MAP1]](%[[kParts]]#2, %[[IV2]])[%[[ARG2]]]
-// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[kParts]]#0, {{.*}}, {{.*}}, 0] [1, 1, 1, 4] [1, 1, 1, 1] : tensor<16x26x18x4xf32> to tensor<1x1x1x4xf32>
-// CHECK: %[[OUT_SLICE:.+]] = tensor.extract_slice %[[ARG6]][%[[IV0]], %[[IV1]], %[[IV2]], 0] [1, 1, 1, 4] [1, 1, 1, 1] : tensor<2x2x2x4xf32> to tensor<1x1x1x4xf32>
-// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[IN_SLICE]] : tensor<1x1x1x4xf32>) outs(%[[OUT_SLICE]] : tensor<1x1x1x4xf32>)
-// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[COPY]] into %[[ARG6]][%[[IV0]], %[[IV1]], %[[IV2]], 0] [1, 1, 1, 4] [1, 1, 1, 1] : tensor<1x1x1x4xf32> into tensor<2x2x2x4xf32>
+// CHECK-DAG: %[[kParts:.+]]:3 = affine.delinearize_index %[[kIDX]] into (16, 24, 16) : index, index, index
+// CHECK-DAG: %[[h:.+]] = affine.apply #[[$MAP1]](%[[kParts]]#1, %[[IV1]])[%[[ARG1]]]
+// CHECK-DAG: %[[w:.+]] = affine.apply #[[$MAP1]](%[[kParts]]#2, %[[IV2]])[%[[ARG2]]]
+// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[kParts]]#0, %[[h]], %[[w]], 0] [1, 1, 1, 4] [1, 1, 1, 1] : tensor<16x26x18x4xf32> to tensor<4xf32>
+// CHECK: %[[DEST_SLICE:.+]] = tensor.extract_slice %[[ARG6]][%[[IV0]], %[[IV1]], %[[IV2]], 0] [1, 1, 1, 4] {{.*}} : tensor<2x2x2x4xf32> to tensor<4xf32>
+// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[IN_SLICE]] : tensor<4xf32>) outs(%[[DEST_SLICE]] : tensor<4xf32>)
+// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[COPY]] into %[[ARG6]][%[[IV0]], %[[IV1]], %[[IV2]], 0] [1, 1, 1, 4] [1, 1, 1, 1] : tensor<4xf32> into tensor<2x2x2x4xf32>
// CHECK: scf.yield %[[INSERT]] : tensor<2x2x2x4xf32>
// CHECK: scf.yield %[[LOOP2]] : tensor<2x2x2x4xf32>
// CHECK: scf.yield %[[LOOP1]] : tensor<2x2x2x4xf32>
// CHECK: return %[[LOOP0]] : tensor<2x2x2x4xf32>
-
-// CHECK-UNROLL-LABEL: func.func @im2col_chwn_output_perm
-// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
+// CHECK-NOT: affine.max
+// CHECK-NOT: affine.min
// -----
-// Test 12: Expanded CHWN layout with output_perm and multi-element batch dims.
-// Five nested loops; the batch dims (batch_pos=[3,4]) are innermost in the input
-// but appear in loops corresponding to the permuted output positions.
+// Expanded CHWN layout with output_perm. Static sizes, non-padded: linalg.copy.
module {
func.func @im2col_chwn_output_perm_expanded(%arg0: tensor<16x26x18x2x4xf32>, %arg1: index, %arg2: index, %arg3: index) -> tensor<2x2x2x2x2x4xf32> {
%0 = tensor.empty() : tensor<2x2x2x2x2x4xf32>
@@ -537,34 +460,42 @@
}
}
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1)[s0] -> (d0 + d1 + s0)>
// CHECK-LABEL: func.func @im2col_chwn_output_perm_expanded
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<16x26x18x2x4xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: index
+// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: index
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x2x2x2x2x4xf32>
-// CHECK: %[[LOOP0:.+]] = scf.for {{.*}} = %[[C0]] to %[[C2]] step %[[C1]]
-// CHECK: %[[LOOP1:.+]] = scf.for {{.*}} = %[[C0]] to %[[C2]] step %[[C1]]
-// CHECK: %[[LOOP2:.+]] = scf.for {{.*}} = %[[C0]] to %[[C2]] step %[[C1]]
-// CHECK: %[[LOOP3:.+]] = scf.for {{.*}} = %[[C0]] to %[[C2]] step %[[C1]]
-// CHECK: %[[LOOP4:.+]] = scf.for {{.*}} = %[[C0]] to %[[C2]] step %[[C1]]
-// CHECK: affine.delinearize_index {{.*}} into (16, 24)
-// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]]{{.*}} [1, 1, 1, 1, 4] [1, 1, 1, 1, 1] : tensor<16x26x18x2x4xf32> to tensor<1x1x1x1x4xf32>
-// CHECK: linalg.copy ins(%[[IN_SLICE]] : tensor<1x1x1x1x4xf32>)
-// CHECK: tensor.insert_slice {{.*}} : tensor<1x1x1x1x4xf32> into tensor<2x2x2x2x2x4xf32>
-// CHECK: scf.yield {{.*}} : tensor<2x2x2x2x2x4xf32>
+// CHECK: %[[LOOP0:.+]] = scf.for %[[IV0:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG4:.+]] = %[[INIT]])
+// CHECK: %[[LOOP1:.+]] = scf.for %[[IV1:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG5:.+]] = %[[ARG4]])
+// CHECK: %[[LOOP2:.+]] = scf.for %[[IV2:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG6:.+]] = %[[ARG5]])
+// CHECK: %[[LOOP3:.+]] = scf.for %[[IV3:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG7:.+]] = %[[ARG6]])
+// CHECK: %[[LOOP4:.+]] = scf.for %[[IV4:.+]] = %[[C0]] to %[[C2]] step %[[C1]] iter_args(%[[ARG8:.+]] = %[[ARG7]])
+// CHECK-DAG: %[[kIDX:.+]] = affine.apply #[[$MAP]](%[[IV0]])[%[[ARG3]]]
+// CHECK-DAG: %[[kParts:.+]]:2 = affine.delinearize_index %[[kIDX]] into (16, 24) : index, index
+// CHECK-DAG: %[[h:.+]] = affine.apply #[[$MAP1]](%[[kParts]]#1, %[[IV2]])[%[[ARG1]]]
+// CHECK-DAG: %[[w:.+]] = affine.apply #[[$MAP1]](%[[IV1]], %[[IV3]])[%[[ARG2]]]
+// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[kParts]]#0, %[[h]], %[[w]], %[[IV4]], 0] [1, 1, 1, 1, 4] [1, 1, 1, 1, 1] : tensor<16x26x18x2x4xf32> to tensor<4xf32>
+// CHECK: %[[DEST_SLICE:.+]] = tensor.extract_slice %[[ARG8]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]], %[[IV4]], 0] [1, 1, 1, 1, 1, 4] {{.*}} : tensor<2x2x2x2x2x4xf32> to tensor<4xf32>
+// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[IN_SLICE]] : tensor<4xf32>) outs(%[[DEST_SLICE]] : tensor<4xf32>)
+// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[COPY]] into %[[ARG8]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]], %[[IV4]], 0] [1, 1, 1, 1, 1, 4] [1, 1, 1, 1, 1, 1] : tensor<4xf32> into tensor<2x2x2x2x2x4xf32>
+// CHECK: scf.yield %[[INSERT]] : tensor<2x2x2x2x2x4xf32>
// CHECK: scf.yield %[[LOOP4]] : tensor<2x2x2x2x2x4xf32>
// CHECK: scf.yield %[[LOOP3]] : tensor<2x2x2x2x2x4xf32>
// CHECK: scf.yield %[[LOOP2]] : tensor<2x2x2x2x2x4xf32>
// CHECK: scf.yield %[[LOOP1]] : tensor<2x2x2x2x2x4xf32>
// CHECK: return %[[LOOP0]] : tensor<2x2x2x2x2x4xf32>
-
-// CHECK-UNROLL-LABEL: func.func @im2col_chwn_output_perm_expanded
-// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
+// CHECK-NOT: affine.max
+// CHECK-NOT: affine.min
// -----
-// Test 13: Multiple k_pos entries (k_pos = [0, 2]).
+// Multiple k_pos entries (k_pos = [0, 2]).
// Verify that both k_pos entries are used for the extract_slice offsets:
// k_pos[0] = input dim 0 (C=2), k_pos[1] = input dim 2 (W=4).
module {
@@ -592,8 +523,8 @@
// -----
-// Test 14: CHWN rank-reduced output with dynamic sizes. The batch dim is innermost.
-// Nested M and K loops reading a 4-element batch slice.
+// CHWN rank-reduced output with dynamic sizes. The batch dim is innermost.
+// Non-padded: extract_slice + linalg.copy + insert_slice.
module {
func.func @im2col_chwn_rank_reduce(%arg0: tensor<16x26x18x4xf32>, %arg1: index, %arg2: index, %m_size: index, %k_size: index) -> tensor<4x?x?xf32> {
%0 = tensor.empty(%m_size, %k_size) : tensor<4x?x?xf32>
@@ -609,26 +540,24 @@
}
// Verify that when the batch dimension is the innermost and generates rank-reduced output,
-// a 4-element tensor slice is extracted, copied, and inserted.
+// a 1D tensor slice is extracted, copied, and inserted (no transpose).
// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1) -> (d0 + d1)>
// CHECK-LABEL: func.func @im2col_chwn_rank_reduce
// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice {{.*}} : tensor<16x26x18x4xf32> to tensor<4xf32>
// CHECK: linalg.copy ins(%[[IN_SLICE]]
// CHECK: tensor.insert_slice {{.*}} : tensor<4xf32> into tensor<4x?x?xf32>
-
-// CHECK-UNROLL-LABEL: func.func @im2col_chwn_rank_reduce
-// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
-// CHECK-UNROLL: scf.for
-// CHECK-UNROLL: scf.for
-// CHECK-UNROLL: linalg.copy
+// CHECK-NOT: affine.max
+// CHECK-NOT: affine.min
// -----
-// Test 15: Backward-weight-style im2col with dilation=2 and expanded M output dims.
+// Backward-weight-style im2col with dilation=2 and expanded M output dims.
// With 2 M output dims each having a single inner size, M coords are used
// directly without delinearization. The spatial offsets include dilation factor.
// Static sizes, non-padded: extract_slice + linalg.copy + insert_slice.
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1)[s0] -> (d0 * 2 + d1 + s0)>
module {
func.func @im2col_bwd_weight_dilation(%arg0: tensor<4x18x18x2xf32>, %m0: index, %m1: index, %k: index) -> tensor<3x3x4x2xf32> {
%0 = tensor.empty() : tensor<3x3x4x2xf32>
@@ -655,13 +584,126 @@
// CHECK: %[[LOOP0:.+]] = scf.for %[[IV0:.+]] = %[[C0]] to %[[C3]] step %[[C1]] iter_args(%[[A0:.+]] = %[[INIT]])
// CHECK: %[[LOOP1:.+]] = scf.for %[[IV1:.+]] = %[[C0]] to %[[C3]] step %[[C1]] iter_args(%[[A1:.+]] = %[[A0]])
// CHECK: %[[LOOP2:.+]] = scf.for %[[IV2:.+]] = %[[C0]] to %[[C4]] step %[[C1]] iter_args(%[[A2:.+]] = %[[A1]])
-// CHECK-DAG: affine.delinearize_index {{.*}} into (4, 8, 8)
-// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]]{{.*}} [1, 1, 1, 2] [1, 1, 1, 1] : tensor<4x18x18x2xf32> to tensor<1x1x1x2xf32>
-// CHECK: %[[DEST_SLICE:.+]] = tensor.extract_slice %[[A2]][%[[IV0]], %[[IV1]], %[[IV2]], 0] [1, 1, 1, 2] {{.*}} : tensor<3x3x4x2xf32> to tensor<1x1x1x2xf32>
-// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[IN_SLICE]] : tensor<1x1x1x2xf32>) outs(%[[DEST_SLICE]] : tensor<1x1x1x2xf32>)
-// CHECK: tensor.insert_slice %[[COPY]] into %[[A2]][%[[IV0]], %[[IV1]], %[[IV2]], 0] [1, 1, 1, 2] [1, 1, 1, 1] : tensor<1x1x1x2xf32> into tensor<3x3x4x2xf32>
+// CHECK-DAG: %[[kIDX:.+]] = affine.apply #[[$MAP]](%[[IV2]])[%[[K]]]
+// CHECK-DAG: %[[kParts:.+]]:3 = affine.delinearize_index %[[kIDX]] into (4, 8, 8) : index, index, index
+// Verify dilation factor 2 is applied to spatial offset computation.
+// CHECK-DAG: %[[h:.+]] = affine.apply #[[$MAP1]](%[[kParts]]#1, %[[IV0]])[%[[M0]]]
+// CHECK-DAG: %[[w:.+]] = affine.apply #[[$MAP1]](%[[kParts]]#2, %[[IV1]])[%[[M1]]]
+// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[ARG0]][%[[kParts]]#0, %[[h]], %[[w]], 0] [1, 1, 1, 2] [1, 1, 1, 1] : tensor<4x18x18x2xf32> to tensor<2xf32>
+// CHECK: %[[DEST_SLICE:.+]] = tensor.extract_slice %[[A2]][%[[IV0]], %[[IV1]], %[[IV2]], 0] [1, 1, 1, 2] {{.*}} : tensor<3x3x4x2xf32> to tensor<2xf32>
+// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[IN_SLICE]] : tensor<2xf32>) outs(%[[DEST_SLICE]] : tensor<2xf32>)
+// CHECK: tensor.insert_slice %[[COPY]] into %[[A2]][%[[IV0]], %[[IV1]], %[[IV2]], 0] [1, 1, 1, 2] [1, 1, 1, 1] : tensor<2xf32> into tensor<3x3x4x2xf32>
// Verify the im2col op is fully lowered to loops.
// CHECK-NOT: iree_linalg_ext.im2col
+// CHECK-NOT: affine.max
+// CHECK-NOT: affine.min
-// CHECK-UNROLL-LABEL: func.func @im2col_bwd_weight_dilation
+// -----
+
+// Padded im2col with vectorized K dim (tile size 8 = channels).
+// The K dim is vectorized, batch and M are size 1. After canonicalize+cse,
+// the loops are removed and we get straight-line code using affine.min/max
+// factor multiplication for bounds checking. Read offsets are clamped
+// with max(0, ...) and min(dimSize-1, ...).
+// CHECK-DAG: #[[$READ_CLAMP:.+]] = affine_map<()[s0] -> (0, s0 - 1)>
+// CHECK-DAG: #[[$DIM_MIN:.+]] = affine_map<()[s0] -> (4, s0)>
+// CHECK-DAG: #[[$HIGH_MIN:.+]] = affine_map<()[s0] -> (-s0 + 6, 1)>
+// CHECK-DAG: #[[$CLAMP0:.+]] = affine_map<()[s0] -> (0, s0)>
+// CHECK-DAG: #[[$LOW_MIN:.+]] = affine_map<()[s0] -> (1, s0)>
+// CHECK-DAG: #[[$MCHECK:.+]] = affine_map<()[s0] -> (-s0 + 1, 1)>
+// CHECK-LABEL: func.func @decompose_padded_im2col_vectorized
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x5x8xf32>
+// CHECK-SAME: %[[MOFF:[a-zA-Z0-9_]+]]: index
+// CHECK-SAME: %[[MPADHI:[a-zA-Z0-9_]+]]: index
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[CST:.+]] = arith.constant 0.000000e+00 : f32
+// Clamped read offset: max(0, offset - padLow), then min(dimSize-1, ...).
+// CHECK: %[[READ_MAX:.+]] = affine.max #[[$READ_CLAMP]]()[%[[MOFF]]]
+// CHECK: %[[READ_OFF:.+]] = affine.min #[[$DIM_MIN]]()[%[[READ_MAX]]]
+// High bound factor: min(-m_off + 6, 1) clamped to >= 0.
+// CHECK: %[[HI_MIN:.+]] = affine.min #[[$HIGH_MIN]]()[%[[MOFF]]]
+// CHECK: %[[HI_OK:.+]] = affine.max #[[$CLAMP0]]()[%[[HI_MIN]]]
+// Low bound factor: min(1, m_off) clamped to >= 0.
+// CHECK: %[[LO_MIN:.+]] = affine.min #[[$LOW_MIN]]()[%[[MOFF]]]
+// CHECK: %[[LO_OK:.+]] = affine.max #[[$CLAMP0]]()[%[[LO_MIN]]]
+// Combined spatial validity factor.
+// CHECK: %[[FACTOR0:.+]] = arith.muli %[[HI_OK]], %[[LO_OK]] : index
+// CHECK: %[[VSIZE0:.+]] = arith.muli %[[FACTOR0]], %[[C8]] : index
+// Output M bounds factor: min(1 - m_pad_high, 1) clamped to >= 0.
+// CHECK: %[[M_MIN:.+]] = affine.min #[[$MCHECK]]()[%[[MPADHI]]]
+// CHECK: %[[M_OK:.+]] = affine.max #[[$CLAMP0]]()[%[[M_MIN]]]
+// CHECK: %[[VSIZE:.+]] = arith.muli %[[VSIZE0]], %[[M_OK]] : index
+// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ARG0]][0, %[[READ_OFF]], 0] [1, 1, %[[VSIZE]]] [1, 1, 1]
+// CHECK: %[[PAD_AMT:.+]] = arith.subi %[[C8]], %[[VSIZE]] : index
+// CHECK: %[[PADDED:.+]] = tensor.pad %[[SLICE]] low[0] high[%[[PAD_AMT]]]
+// CHECK: tensor.yield %[[CST]]
+// CHECK: return {{.*}} : tensor<1x1x8xf32>
+// CHECK-UNROLL-LABEL: func.func @decompose_padded_im2col_vectorized
// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
+// CHECK-UNROLL: tensor.extract_slice
+// CHECK-UNROLL: tensor.pad
+module {
+ func.func @decompose_padded_im2col_vectorized(%arg0: tensor<1x5x8xf32>, %m_off: index, %m_pad_high: index) -> tensor<1x1x8xf32> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %0 = tensor.empty() : tensor<1x1x8xf32>
+ %1 = iree_linalg_ext.im2col
+ strides = [1] dilations = [1] kernel_size = [3]
+ offsets = [0, %m_off, 0] output_sizes = [[1], [5], [3, 8]]
+ batch_pos = [0] m_pos = [1] k_pos = [2]
+ input_k_perm = [0, 1] output_perm = [0, 1, 2]
+ input_pad_low = [0, 1, 0] input_pad_high = [0, 1, 0]
+ output_pad_low = [0, 0, 0] output_pad_high = [0, %m_pad_high, 0]
+ pad_value(%cst : f32)
+ ins(%arg0 : tensor<1x5x8xf32>) outs(%0 : tensor<1x1x8xf32>) -> tensor<1x1x8xf32>
+ return %1 : tensor<1x1x8xf32>
+ }
+}
+
+// -----
+
+// Output-side M-dimension padding: output M dim tile > product(M output_sizes).
+// Input: 1x5x8 (batch=1, spatial=5, channels=8). Kernel [3], stride [1] -> OH=3.
+// output_sizes M = [3], product = 3. The tile is 1 in M, so with dynamic offset
+// positions beyond M=2 should be all-padding (validSize=0).
+// K is vectorized (tile size 8 = channels). No input padding, so the "preserve
+// zero padding" optimization skips spatial bounds checks. Only the clamped read
+// offset and the output M bounds check remain.
+// CHECK-DAG: #[[$CLAMP0:.+]] = affine_map<()[s0] -> (0, s0)>
+// CHECK-DAG: #[[$DIM_MIN:.+]] = affine_map<()[s0] -> (4, s0)>
+// CHECK-DAG: #[[$MCHECK:.+]] = affine_map<()[s0] -> (-s0 + 1, 1)>
+// CHECK-LABEL: func.func @decompose_output_pad_m
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x5x8xf32>
+// CHECK-SAME: %[[MOFF:[a-zA-Z0-9_]+]]: index
+// CHECK-SAME: %[[MPADHI:[a-zA-Z0-9_]+]]: index
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[CST:.+]] = arith.constant 0.000000e+00 : f32
+// Clamped read offset: max(0, m_off), then min(dimSize-1, ...).
+// CHECK: %[[READ_MAX:.+]] = affine.max #[[$CLAMP0]]()[%[[MOFF]]]
+// CHECK: %[[READ_OFF:.+]] = affine.min #[[$DIM_MIN]]()[%[[READ_MAX]]]
+// Output M bounds factor: min(1 - m_pad_high, 1) clamped to >= 0.
+// CHECK: %[[M_MIN:.+]] = affine.min #[[$MCHECK]]()[%[[MPADHI]]]
+// CHECK: %[[M_OK:.+]] = affine.max #[[$CLAMP0]]()[%[[M_MIN]]]
+// CHECK: %[[VSIZE:.+]] = arith.muli %[[M_OK]], %[[C8]] : index
+// CHECK: tensor.extract_slice %[[ARG0]][0, %[[READ_OFF]], 0] [1, 1, %[[VSIZE]]]
+// CHECK: tensor.pad {{.*}} low[0]
+// CHECK-NEXT: ^bb0
+// CHECK-NEXT: tensor.yield %[[CST]]
+// CHECK-UNROLL-LABEL: func.func @decompose_output_pad_m
+// CHECK-UNROLL-NOT: iree_linalg_ext.im2col
+// CHECK-UNROLL: tensor.extract_slice
+// CHECK-UNROLL: tensor.pad
+module {
+ func.func @decompose_output_pad_m(%arg0: tensor<1x5x8xf32>, %m_off: index, %m_pad_high: index) -> tensor<1x1x8xf32> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %0 = tensor.empty() : tensor<1x1x8xf32>
+ %1 = iree_linalg_ext.im2col
+ strides = [1] dilations = [1] kernel_size = [3]
+ offsets = [0, %m_off, 0] output_sizes = [[1], [3], [3, 8]]
+ batch_pos = [0] m_pos = [1] k_pos = [2]
+ input_k_perm = [0, 1] output_perm = [0, 1, 2]
+ output_pad_low = [0, 0, 0] output_pad_high = [0, %m_pad_high, 0]
+ pad_value(%cst : f32)
+ ins(%arg0 : tensor<1x5x8xf32>) outs(%0 : tensor<1x1x8xf32>) -> tensor<1x1x8xf32>
+ return %1 : tensor<1x1x8xf32>
+ }
+}
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
index fadaa13..c4c8a8c 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
@@ -51,6 +51,15 @@
return affine::makeComposedFoldedAffineApply(builder, loc, addMap, {a, b});
}
+OpFoldResult subOfrs(OpBuilder &builder, Location loc, OpFoldResult a,
+ OpFoldResult b) {
+ AffineExpr d0, d1;
+ bindDims(builder.getContext(), d0, d1);
+ return affine::makeComposedFoldedAffineApply(
+ builder, loc, AffineMap::get(2, 0, {d0 - d1}, builder.getContext()),
+ {a, b});
+}
+
OpFoldResult mulOfrs(OpBuilder &builder, Location loc, OpFoldResult a,
OpFoldResult b) {
AffineExpr d0, d1;
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h b/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h
index 6773db2..7d8196c 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h
@@ -34,6 +34,10 @@
OpFoldResult mulOfrs(OpBuilder &builder, Location loc, OpFoldResult a,
OpFoldResult b);
+/// Helper method to subtract 2 OpFoldResult inputs with affine.apply.
+OpFoldResult subOfrs(OpBuilder &builder, Location loc, OpFoldResult a,
+ OpFoldResult b);
+
/// Helper method to compute (a * b + c) with OpFoldResult inputs using
/// affine.apply.
OpFoldResult mulAddOfrs(OpBuilder &builder, Location loc, OpFoldResult a,
