[Mlir-commits] [mlir] [mlir][linalg] Fix depthwise im2col for multiple input batches (PR #225623)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Sep 23 19:56:20 PDT 2026
https://github.com/AkunCube updated https://github.com/llvm/llvm-project/pull/225623
>From 38c8dfd41c5f74651cefca78950f4cb11b814b14 Mon Sep 17 00:00:00 2001
From: cubehan3 <cubehan3 at gmail.com>
Date: Wed, 23 Sep 2026 07:41:31 +0000
Subject: [PATCH] [mlir][linalg] Fix depthwise im2col for multiple input
batches
---
.../Transforms/ConvertConv2DToImg2Col.cpp | 20 ++++++--
.../Linalg/convert-conv2d-to-img2col.mlir | 50 ++++++++++---------
2 files changed, 41 insertions(+), 29 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
index e2dd0a6d79a73..ddbf21dc841ef 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
@@ -333,6 +333,15 @@ rewriteInIm2Col(RewriterBase &rewriter,
int fh = filterTShape[1];
int fw = filterTShape[2];
+ // Replicate the filter across input batches to match the flattened (n, c)
+ // batch dimension of batch_matvec.
+ Value broadcastFilterInit = tensor::EmptyOp::create(
+ rewriter, loc, {n, c, fh, fw}, filterType.getElementType());
+ Value broadcastedFilter =
+ linalg::BroadcastOp::create(rewriter, loc, filterT, broadcastFilterInit,
+ ArrayRef<int64_t>{0})
+ ->getResult(0);
+
SmallVector<int64_t> colTensorShape = {n, c, oh, ow, fh, fw};
Value transposedOutputTensor = transposeOperand(output, {0, 3, 1, 2});
@@ -369,23 +378,24 @@ rewriteInIm2Col(RewriterBase &rewriter,
SmallVector<ReassociationIndices> img2ColTensorReassocIndices = {
{0, 1}, {2, 3}, {4, 5}};
- SmallVector<ReassociationIndices> filterReassociationIndice = {{0}, {1, 2}};
+ SmallVector<ReassociationIndices> filterReassociationIndices = {{0, 1},
+ {2, 3}};
SmallVector<ReassociationIndices> outputReassociationIndice = {{0, 1},
{2, 3}};
auto reshapedImg2ColTensorType = RankedTensorType::get(
{n * c, oh * ow, fh * fw}, inputType.getElementType());
auto reshapedFilterTensorType =
- RankedTensorType::get({c, fh * fw}, filterType.getElementType());
+ RankedTensorType::get({n * c, fh * fw}, filterType.getElementType());
auto reshapedOutputTensorType =
RankedTensorType::get({n * c, oh * ow}, outputType.getElementType());
Value reshapedImg2ColTensor = tensor::CollapseShapeOp::create(
rewriter, loc, reshapedImg2ColTensorType, img2ColTensor.getResult(0),
img2ColTensorReassocIndices);
- Value reshapedFilterTensor =
- tensor::CollapseShapeOp::create(rewriter, loc, reshapedFilterTensorType,
- filterT, filterReassociationIndice);
+ Value reshapedFilterTensor = tensor::CollapseShapeOp::create(
+ rewriter, loc, reshapedFilterTensorType, broadcastedFilter,
+ filterReassociationIndices);
Value reshapedoutputTensor = tensor::CollapseShapeOp::create(
rewriter, loc, reshapedOutputTensorType, transposedOutputTensor,
outputReassociationIndice);
diff --git a/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir b/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir
index 19b108014c7a1..7d26014061747 100644
--- a/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir
+++ b/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir
@@ -124,17 +124,17 @@ module attributes {transform.with_named_sequence} {
// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>
// CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1, d2)>
// CHECK: @depthwise_conv_hwc_114x16x3
-// CHECK-SAME: %[[INPUT:.+]]: tensor<1x114x114x16xf32>
+// CHECK-SAME: %[[INPUT:.+]]: tensor<2x114x114x16xf32>
// CHECK-SAME: %[[FILTER:.+]]: tensor<3x3x16xf32>
-// CHECK-SAME: %[[OUTPUT:.+]]: tensor<1x112x112x16xf32>
-// CHECK: %[[INPUT_T_INIT:.+]] = tensor.empty() : tensor<1x16x114x114xf32>
+// CHECK-SAME: %[[OUTPUT:.+]]: tensor<2x112x112x16xf32>
+// CHECK: %[[INPUT_T_INIT:.+]] = tensor.empty() : tensor<2x16x114x114xf32>
// CHECK: %[[INPUT_T:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]
-// CHECK-SAME: ins(%[[INPUT]] : tensor<1x114x114x16xf32>) outs(%[[INPUT_T_INIT]] : tensor<1x16x114x114xf32>) {
+// CHECK-SAME: ins(%[[INPUT]] : tensor<2x114x114x16xf32>) outs(%[[INPUT_T_INIT]] : tensor<2x16x114x114xf32>) {
// CHECK-NEXT: ^bb0(%[[ARG3:.+]]: f32, %[[ARG4:.+]]: f32):
// CHECK-NEXT: linalg.yield %[[ARG3]] : f32
-// CHECK-NEXT: } -> tensor<1x16x114x114xf32>
+// CHECK-NEXT: } -> tensor<2x16x114x114xf32>
// CHECK: %[[FILTER_T_INIT:.+]] = tensor.empty() : tensor<16x3x3xf32>
// CHECK: %[[FILTER_T:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]]
@@ -143,46 +143,48 @@ module attributes {transform.with_named_sequence} {
// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32):
// CHECK: linalg.yield
// CHECK: } -> tensor<16x3x3xf32>
-// CHECK: %[[INIT_OUTPUT_TENSOR:.+]] = tensor.empty() : tensor<1x16x112x112xf32>
+// CHECK: %[[FILTER_BROADCAST_INIT:.+]] = tensor.empty() : tensor<2x16x3x3xf32>
+// CHECK: %[[FILTER_BROADCAST:.+]] = linalg.broadcast ins(%[[FILTER_T]] : tensor<16x3x3xf32>) outs(%[[FILTER_BROADCAST_INIT]] : tensor<2x16x3x3xf32>) dimensions = [0]
+// CHECK: %[[INIT_OUTPUT_TENSOR:.+]] = tensor.empty() : tensor<2x16x112x112xf32>
// CHECK: %[[OUTPUT_T:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]
-// CHECK-SAME: ins(%[[OUTPUT]] : tensor<1x112x112x16xf32>) outs(%[[INIT_OUTPUT_TENSOR]] : tensor<1x16x112x112xf32>) {
+// CHECK-SAME: ins(%[[OUTPUT]] : tensor<2x112x112x16xf32>) outs(%[[INIT_OUTPUT_TENSOR]] : tensor<2x16x112x112xf32>) {
// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32):
// CHECK-NEXT: linalg.yield
-// CHECK-NEXT: } -> tensor<1x16x112x112xf32>
-// CHECK: %[[INIT_COL_TENSOR:.+]] = tensor.empty() : tensor<1x16x112x112x3x3xf32>
+// CHECK-NEXT: } -> tensor<2x16x112x112xf32>
+// CHECK: %[[INIT_COL_TENSOR:.+]] = tensor.empty() : tensor<2x16x112x112x3x3xf32>
// CHECK: %[[COL_TENSOR:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP4]], #[[MAP5]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]
-// CHECK-SAME: ins(%[[INPUT_T]] : tensor<1x16x114x114xf32>) outs(%[[INIT_COL_TENSOR]] : tensor<1x16x112x112x3x3xf32>) {
+// CHECK-SAME: ins(%[[INPUT_T]] : tensor<2x16x114x114xf32>) outs(%[[INIT_COL_TENSOR]] : tensor<2x16x112x112x3x3xf32>) {
// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32):
// CHECK-NEXT: linalg.yield
-// CHECK-NEXT: } -> tensor<1x16x112x112x3x3xf32>
+// CHECK-NEXT: } -> tensor<2x16x112x112x3x3xf32>
// CHECK: %[[COL_TENSOR_R:.+]] = tensor.collapse_shape %[[COL_TENSOR]]
-// CHECK-SAME: tensor<1x16x112x112x3x3xf32> into tensor<16x12544x9xf32>
-// CHECK: %[[FILTER_T_R:.+]] = tensor.collapse_shape %[[FILTER_T]]
-// CHECK-SAME: tensor<16x3x3xf32> into tensor<16x9xf32>
+// CHECK-SAME: tensor<2x16x112x112x3x3xf32> into tensor<32x12544x9xf32>
+// CHECK: %[[FILTER_T_R:.+]] = tensor.collapse_shape %[[FILTER_BROADCAST]]
+// CHECK-SAME: tensor<2x16x3x3xf32> into tensor<32x9xf32>
// CHECK: %[[OUTPUT_T_R:.+]] = tensor.collapse_shape %[[OUTPUT_T]]
-// CHECK-SAME: tensor<1x16x112x112xf32> into tensor<16x12544xf32>
-// CHECK: %[[BMV_RESULT:.+]] = linalg.batch_matvec ins(%[[COL_TENSOR_R]], %[[FILTER_T_R]] : tensor<16x12544x9xf32>, tensor<16x9xf32>) outs(%[[OUTPUT_T_R]] : tensor<16x12544xf32>) -> tensor<16x12544xf32>
+// CHECK-SAME: tensor<2x16x112x112xf32> into tensor<32x12544xf32>
+// CHECK: %[[BMV_RESULT:.+]] = linalg.batch_matvec ins(%[[COL_TENSOR_R]], %[[FILTER_T_R]] : tensor<32x12544x9xf32>, tensor<32x9xf32>) outs(%[[OUTPUT_T_R]] : tensor<32x12544xf32>) -> tensor<32x12544xf32>
// CHECK: %[[RESULT_R:.+]] = tensor.expand_shape %[[BMV_RESULT]]
-// CHECK-SAME: tensor<16x12544xf32> into tensor<1x16x112x112xf32>
-// CHECK: %[[RESULT_INIT:.+]] = tensor.empty() : tensor<1x112x112x16xf32>
+// CHECK-SAME: tensor<32x12544xf32> into tensor<2x16x112x112xf32>
+// CHECK: %[[RESULT_INIT:.+]] = tensor.empty() : tensor<2x112x112x16xf32>
// CHECK: %[[RESULT:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP6]], #[[MAP1]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]
-// CHECK-SAME: ins(%[[RESULT_R]] : tensor<1x16x112x112xf32>) outs(%[[RESULT_INIT]] : tensor<1x112x112x16xf32>) {
+// CHECK-SAME: ins(%[[RESULT_R]] : tensor<2x16x112x112xf32>) outs(%[[RESULT_INIT]] : tensor<2x112x112x16xf32>) {
// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32):
// CHECK-NEXT: linalg.yield
-// CHECK-NEXT: } -> tensor<1x112x112x16xf32>
-// CHECK: return %[[RESULT]] : tensor<1x112x112x16xf32>
-func.func @depthwise_conv_hwc_114x16x3(%input: tensor<1x114x114x16xf32>, %filter: tensor<3x3x16xf32>, %output: tensor<1x112x112x16xf32>) -> tensor<1x112x112x16xf32> {
+// CHECK-NEXT: } -> tensor<2x112x112x16xf32>
+// CHECK: return %[[RESULT]] : tensor<2x112x112x16xf32>
+func.func @depthwise_conv_hwc_114x16x3(%input: tensor<2x114x114x16xf32>, %filter: tensor<3x3x16xf32>, %output: tensor<2x112x112x16xf32>) -> tensor<2x112x112x16xf32> {
%0 = linalg.depthwise_conv_2d_nhwc_hwc {
dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>
- } ins(%input, %filter : tensor<1x114x114x16xf32>, tensor<3x3x16xf32>) outs(%output : tensor<1x112x112x16xf32>) -> tensor<1x112x112x16xf32>
- return %0 : tensor<1x112x112x16xf32>
+ } ins(%input, %filter : tensor<2x114x114x16xf32>, tensor<3x3x16xf32>) outs(%output : tensor<2x112x112x16xf32>) -> tensor<2x112x112x16xf32>
+ return %0 : tensor<2x112x112x16xf32>
}
module attributes {transform.with_named_sequence} {
More information about the Mlir-commits
mailing list