[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 00:43:56 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: Akun (AkunCube)
<details>
<summary>Changes</summary>
The depthwise im2col rewrite flattens the input batch and channel dimensions into the leading dimension of `linalg.batch_matvec`, but the filter retains only the channel dimension. With more than one input batch, the generated operands have different leading extents and fail verification.
Broadcast the transposed filter across input batches before collapsing it to `[batch * channel, filter_height * filter_width]`. This gives all three `batch_matvec` operands the same leading extent while sharing each channel's filter across batches. Update the existing depthwise im2col test to use two batches and check the broadcast, collapsed shapes, and contraction.
Fixes: #<!-- -->225252
---
Full diff: https://github.com/llvm/llvm-project/pull/225623.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp (+15-5)
- (modified) mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir (+26-24)
``````````diff
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} {
``````````
</details>
https://github.com/llvm/llvm-project/pull/225623
More information about the Mlir-commits
mailing list