[Mlir-commits] [mlir] [mlir][linalg] Support non-unit dilations in im2col decomposition (PR #208424)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 9 04:05:04 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: Zmicier Prybysh (dimp-pl)
<details>
<summary>Changes</summary>
When learning linang dialect, i noticed this TODO and decided to tackle it.
The im2col patterns for conv_2d_nhwc_hwcf, conv_2d_nchw_fchw and conv_2d_nhwc_fhwc avoided non-unit dilations, but the restriction doesn't seem to be fundamental: dilation only affects the gather step, which can fold it into its indexing map. Generalize the convolved index expression from `oh * stride + fh` to `oh * stride + fh * dilation` and drop the match failures.
Verified by comparing the results of a dilated convolution lowered directly to loops against the im2col decomposition with mlir-runner.
Assisted-By: Claude Code (for generating tests).
---
Full diff: https://github.com/llvm/llvm-project/pull/208424.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp (+14-27)
- (modified) mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir (+112)
``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
index 4cd6fdf4f7341..e2dd0a6d79a73 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
@@ -54,16 +54,16 @@ static Value createMul(Location loc, Value x, Value y, Type accType,
}
// Generate the affine expression to compute the convolved index
-// for the input as `oIndex * stride + fIndex`,
+// for the input as `oIndex * stride + fIndex * dilation`,
// where oIndex: output iterator; fIndex: filter iterator.
static AffineExpr getConvolvedExpr(OpBuilder &b, int64_t stride,
- bool useSymbols = true) {
+ int64_t dilation, bool useSymbols = true) {
AffineExpr oExpr, fExpr;
if (useSymbols)
bindSymbols(b.getContext(), oExpr, fExpr);
else
bindDims(b.getContext(), oExpr, fExpr);
- return AffineExpr(stride * oExpr + fExpr);
+ return AffineExpr(stride * oExpr + dilation * fExpr);
}
// Stores the affine expressions to map the iteration space of the im2col matrix
@@ -91,12 +91,14 @@ struct Im2ColToInputDimsExprs {
///
/// @param exprs Affine expressions for output and filter indices.
/// @param strides [height, width] stride values for the convolution.
+/// @param dilations [height, width] dilation values for the convolution.
/// @param rewriter Pattern rewriter.
/// @return Affine expressions mapping im2col matrix indices to input
/// offsets.
static Im2ColToInputDimsExprs
getIm2ColInputExpressions(Im2ColToOperandsExprs exprs,
- ArrayRef<int64_t> strides, RewriterBase &rewriter) {
+ ArrayRef<int64_t> strides,
+ ArrayRef<int64_t> dilations, RewriterBase &rewriter) {
// maps the iteration space of the im2col matrix to (output_y, filter_y)
auto hIndicesMap = AffineMap::inferFromExprList(
{ArrayRef{exprs.ohIndex, exprs.fhIndex}}, rewriter.getContext())[0];
@@ -110,10 +112,10 @@ getIm2ColInputExpressions(Im2ColToOperandsExprs exprs,
// then we compose them with the maps that map the im2col matrix elements to
// the (out_element, filter_element) pairs.
auto bIndexExpr = rewriter.getAffineDimExpr(0U);
- auto hIndexExpr = getConvolvedExpr(rewriter, strides[0],
+ auto hIndexExpr = getConvolvedExpr(rewriter, strides[0], dilations[0],
/*useSymbols*/ false);
hIndexExpr = hIndexExpr.compose(hIndicesMap);
- auto wIndexExpr = getConvolvedExpr(rewriter, strides[1],
+ auto wIndexExpr = getConvolvedExpr(rewriter, strides[1], dilations[1],
/*useSymbols*/ false);
wIndexExpr = wIndexExpr.compose(wIndicesMap);
auto cIndexExpr = exprs.icIndex;
@@ -138,11 +140,6 @@ rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcHwcfOp convOp) {
return rewriter.notifyMatchFailure(convOp,
"expected a static shape for the input");
- // TODO: Support dilation.
- if (!hasAllOneValues(convOp.getDilations()))
- return rewriter.notifyMatchFailure(convOp,
- "expected all ones for dilations");
-
MLIRContext *context = rewriter.getContext();
Value input = convOp.getInputs()[0];
Value filter = convOp.getInputs()[1];
@@ -203,10 +200,10 @@ rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcHwcfOp convOp) {
i2cToOperExprs.ohIndex = mIndicesExprs[0];
i2cToOperExprs.owIndex = mIndicesExprs[1];
- // im2col[n, oh*ow, fh*fw*ic] = input[n, sh*oh + fh, sw*ow + fw, ic]
+ // im2col[n, oh*ow, fh*fw*ic] = input[n, sh*oh + dh*fh, sw*ow + dw*fw, ic]
Im2ColToInputDimsExprs inExprs = getIm2ColInputExpressions(
i2cToOperExprs, llvm::to_vector(convOp.getStrides().getValues<int64_t>()),
- rewriter);
+ llvm::to_vector(convOp.getDilations().getValues<int64_t>()), rewriter);
auto inMap =
AffineMap::inferFromExprList({ArrayRef{inExprs.bIndex, inExprs.hIndex,
inExprs.wIndex, inExprs.cIndex}},
@@ -431,11 +428,6 @@ rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp) {
return rewriter.notifyMatchFailure(convOp,
"expected a static shape for the input");
- // TODO: Support dilation.
- if (!hasAllOneValues(convOp.getDilations()))
- return rewriter.notifyMatchFailure(convOp,
- "expected all ones for dilations");
-
Value input = convOp.getInputs()[0];
Value filter = convOp.getInputs()[1];
Value output = convOp.getOutputs()[0];
@@ -497,12 +489,12 @@ rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp) {
i2cToOperExprs.owIndex = mIndicesExprs[1];
Im2ColToInputDimsExprs inExprs = getIm2ColInputExpressions(
i2cToOperExprs, llvm::to_vector(convOp.getStrides().getValues<int64_t>()),
- rewriter);
+ llvm::to_vector(convOp.getDilations().getValues<int64_t>()), rewriter);
auto inMap =
AffineMap::inferFromExprList({ArrayRef{inExprs.bIndex, inExprs.cIndex,
inExprs.hIndex, inExprs.wIndex}},
rewriter.getContext())[0];
- // im2col[n, ic*fh*fw, oh*ow] = input[n, ic, sh*oh + fh, sw*ow + fw]
+ // im2col[n, ic*fh*fw, oh*ow] = input[n, ic, sh*oh + dh*fh, sw*ow + dw*fw]
SmallVector<AffineMap> img2colIndexingMaps = {
inMap, AffineMap::getMultiDimIdentityMap(nloops, context)};
@@ -565,11 +557,6 @@ rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcFhwcOp convOp) {
return rewriter.notifyMatchFailure(convOp,
"expected a static shape for the input");
- // TODO: Support dilation.
- if (!hasAllOneValues(convOp.getDilations()))
- return rewriter.notifyMatchFailure(convOp,
- "expected all ones for dilations");
-
MLIRContext *context = rewriter.getContext();
Value input = convOp.getInputs()[0];
Value filter = convOp.getInputs()[1];
@@ -632,10 +619,10 @@ rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcFhwcOp convOp) {
i2cToOperExprs.ohIndex = mIndicesExprs[0];
i2cToOperExprs.owIndex = mIndicesExprs[1];
- // im2col[n, oh*ow, fh*fw*ic] = input[n, sh*oh + fh, sw*ow + fw, ic]
+ // im2col[n, oh*ow, fh*fw*ic] = input[n, sh*oh + dh*fh, sw*ow + dw*fw, ic]
Im2ColToInputDimsExprs inExprs = getIm2ColInputExpressions(
i2cToOperExprs, llvm::to_vector(convOp.getStrides().getValues<int64_t>()),
- rewriter);
+ llvm::to_vector(convOp.getDilations().getValues<int64_t>()), rewriter);
auto inMap =
AffineMap::inferFromExprList({ArrayRef{inExprs.bIndex, inExprs.hIndex,
inExprs.wIndex, inExprs.cIndex}},
diff --git a/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir b/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir
index 152a392afe247..387f9762d518d 100644
--- a/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir
+++ b/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir
@@ -556,3 +556,115 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// Dilated NHWC-HWCF: the im2col gather reads input at
+// oh*stride + fh*dilation, i.e. d1 floordiv 12 + (d2 floordiv 12) * 2.
+
+// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 floordiv 12 + (d2 floordiv 12) * 2, d1 mod 12 + ((d2 mod 12) floordiv 4) * 2, d2 mod 4)>
+// CHECK-DAG: #[[MAPI2C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+// CHECK: func.func @conv_nhwc_hwcf_dilated
+// CHECK-SAME: (%[[INPUT:.+]]: tensor<1x16x16x4xf32>, %[[FILTER:.+]]: tensor<3x3x4x16xf32>, %[[INIT:.+]]: tensor<1x12x12x16xf32>)
+// CHECK-DAG: %[[CS_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0, 1, 2], [3]] : tensor<3x3x4x16xf32> into tensor<36x16xf32>
+// CHECK-DAG: %[[CS_RESULT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1, 2], [3]] : tensor<1x12x12x16xf32> into tensor<1x144x16xf32>
+// CHECK: %[[IT:.+]] = tensor.empty() : tensor<1x144x36xf32>
+// CHECK: %[[IMG2COL:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAPI2C]]]
+// CHECK-SAME: ins(%[[INPUT]] : tensor<1x16x16x4xf32>)
+// CHECK-SAME: outs(%[[IT]] : tensor<1x144x36xf32>)
+// CHECK: %[[MATMUL:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[IMG2COL]], %[[CS_FILTER]] : tensor<1x144x36xf32>, tensor<36x16xf32>)
+// CHECK-SAME: outs(%[[CS_RESULT]] : tensor<1x144x16xf32>)
+// CHECK: %[[CS_FINAL:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0], [1, 2], [3]] output_shape [1, 12, 12, 16] : tensor<1x144x16xf32> into tensor<1x12x12x16xf32>
+// CHECK: return %[[CS_FINAL]]
+func.func @conv_nhwc_hwcf_dilated(%arg0: tensor<1x16x16x4xf32>, %arg1: tensor<3x3x4x16xf32>, %arg2: tensor<1x12x12x16xf32>) -> tensor<1x12x12x16xf32> {
+ %0 = linalg.conv_2d_nhwc_hwcf
+ {dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
+ ins(%arg0, %arg1: tensor<1x16x16x4xf32>, tensor<3x3x4x16xf32>)
+ outs(%arg2: tensor<1x12x12x16xf32>) -> tensor<1x12x12x16xf32>
+ return %0 : tensor<1x12x12x16xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1:2 = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Dilated NCHW-FCHW: h index is d2 floordiv 12 + ((d1 mod 9) floordiv 3) * 2.
+
+// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 floordiv 9, d2 floordiv 12 + ((d1 mod 9) floordiv 3) * 2, d2 mod 12 + (d1 mod 3) * 2)>
+// CHECK-DAG: #[[MAPI2C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+// CHECK: func.func @conv_nchw_fchw_dilated
+// CHECK-SAME: (%[[INPUT:.+]]: tensor<8x4x16x16xf32>, %[[FILTER:.+]]: tensor<16x4x3x3xf32>, %[[INIT:.+]]: tensor<8x16x12x12xf32>)
+// CHECK-DAG: %[[CS_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0], [1, 2, 3]] : tensor<16x4x3x3xf32> into tensor<16x36xf32>
+// CHECK-DAG: %[[CS_RESULT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1], [2, 3]] : tensor<8x16x12x12xf32> into tensor<8x16x144xf32>
+// CHECK: %[[IT:.+]] = tensor.empty() : tensor<8x36x144xf32>
+// CHECK: %[[IMG2COL:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAPI2C]]]
+// CHECK-SAME: ins(%[[INPUT]] : tensor<8x4x16x16xf32>)
+// CHECK-SAME: outs(%[[IT]] : tensor<8x36x144xf32>)
+// CHECK: %[[MATMUL:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[CS_FILTER]], %[[IMG2COL]] : tensor<16x36xf32>, tensor<8x36x144xf32>)
+// CHECK-SAME: outs(%[[CS_RESULT]] : tensor<8x16x144xf32>)
+// CHECK: %[[CS_FINAL:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0], [1], [2, 3]] output_shape [8, 16, 12, 12] : tensor<8x16x144xf32> into tensor<8x16x12x12xf32>
+// CHECK: return %[[CS_FINAL]]
+func.func @conv_nchw_fchw_dilated(%arg0: tensor<8x4x16x16xf32>, %arg1: tensor<16x4x3x3xf32>, %arg2: tensor<8x16x12x12xf32>) -> tensor<8x16x12x12xf32> {
+ %0 = linalg.conv_2d_nchw_fchw
+ {dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
+ ins(%arg0, %arg1: tensor<8x4x16x16xf32>, tensor<16x4x3x3xf32>)
+ outs(%arg2: tensor<8x16x12x12xf32>) -> tensor<8x16x12x12xf32>
+ return %0 : tensor<8x16x12x12xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.conv_2d_nchw_fchw"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1:2 = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Dilated NHWC-FHWC: same dilated input gather map as NHWC-HWCF.
+
+// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 floordiv 12 + (d2 floordiv 12) * 2, d1 mod 12 + ((d2 mod 12) floordiv 4) * 2, d2 mod 4)>
+// CHECK-DAG: #[[MAPI2C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+// CHECK: func.func @conv_nhwc_fhwc_dilated
+// CHECK-SAME: (%[[INPUT:.+]]: tensor<1x16x16x4xf32>, %[[FILTER:.+]]: tensor<16x3x3x4xf32>, %[[INIT:.+]]: tensor<1x12x12x16xf32>)
+// CHECK-DAG: %[[CS_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0], [1, 2, 3]] : tensor<16x3x3x4xf32> into tensor<16x36xf32>
+// CHECK-DAG: %[[CS_RESULT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1, 2], [3]] : tensor<1x12x12x16xf32> into tensor<1x144x16xf32>
+// CHECK: %[[IT:.+]] = tensor.empty() : tensor<1x144x36xf32>
+// CHECK: %[[IMG2COL:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAPI2C]]]
+// CHECK-SAME: ins(%[[INPUT]] : tensor<1x16x16x4xf32>)
+// CHECK-SAME: outs(%[[IT]] : tensor<1x144x36xf32>)
+// CHECK: %[[MATMUL:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[IMG2COL]], %[[CS_FILTER]] : tensor<1x144x36xf32>, tensor<16x36xf32>)
+// CHECK-SAME: outs(%[[CS_RESULT]] : tensor<1x144x16xf32>)
+// CHECK: %[[CS_FINAL:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0], [1, 2], [3]] output_shape [1, 12, 12, 16] : tensor<1x144x16xf32> into tensor<1x12x12x16xf32>
+// CHECK: return %[[CS_FINAL]]
+func.func @conv_nhwc_fhwc_dilated(%arg0: tensor<1x16x16x4xf32>, %arg1: tensor<16x3x3x4xf32>, %arg2: tensor<1x12x12x16xf32>) -> tensor<1x12x12x16xf32> {
+ %0 = linalg.conv_2d_nhwc_fhwc
+ {dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
+ ins(%arg0, %arg1: tensor<1x16x16x4xf32>, tensor<16x3x3x4xf32>)
+ outs(%arg2: tensor<1x12x12x16xf32>) -> tensor<1x12x12x16xf32>
+ return %0 : tensor<1x12x12x16xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_fhwc"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1:2 = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/208424
More information about the Mlir-commits
mailing list