[Mlir-commits] [mlir] [mlir][linalg] Support quantized NCHW/FCHW convolution in img2col (PR #225308)
Dmitry Manjiev
llvmlistbot at llvm.org
Mon Sep 28 22:48:21 PDT 2026
https://github.com/Dima365 updated https://github.com/llvm/llvm-project/pull/225308
>From def0eeabce50b54ba965491db9df96d011fc0179 Mon Sep 17 00:00:00 2001
From: Dmitry Manjiev <mandzievdmitrij at gmail.com>
Date: Tue, 22 Sep 2026 07:54:51 +0300
Subject: [PATCH 1/2] [mlir][linalg] Support quantized NCHW/FCHW convolution in
img2col
Share the NCHW/FCHW img2col rewrite between regular and quantized
convolutions. Preserve scalar zero-point operands and reuse the original
convolution body in the generated linalg.generic.
Keep the original filter element type when collapsing its shape and
register the quantized rewrite with the Transform dialect and rewrite
patterns.
Add FileCheck coverage and an integration test comparing the original
and transformed convolutions against expected results. Update existing
NCHW/FCHW checks to account for the preserved input operand order.
---
.../Linalg/TransformOps/LinalgTransformOps.td | 2 +-
.../Dialect/Linalg/Transforms/Transforms.h | 4 +
.../TransformOps/LinalgTransformOps.cpp | 3 +
.../Transforms/ConvertConv2DToImg2Col.cpp | 71 +++++++++++++----
.../Linalg/convert-conv2d-to-img2col.mlir | 74 +++++++++++++++++-
.../CPU/test-conv-2d-nchw-fchw-q-img2col.mlir | 76 +++++++++++++++++++
6 files changed, 209 insertions(+), 21 deletions(-)
create mode 100644 mlir/test/Integration/Dialect/Linalg/CPU/test-conv-2d-nchw-fchw-q-img2col.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 5df9e14b8a3f06..6914cdc12270c0 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -2754,7 +2754,7 @@ def ConvertConv2DToImg2ColOp : Op<Transform_Dialect,
ReportTrackingListenerFailuresOpTrait]> {
let description = [{
Convert linalg.conv_2d_xxx into linalg.generic (for img2col packing)
- and linalg.matmul.
+ followed by matrix multiplication or batched matrix-vector multiplication.
A convolution operation can be written as a matrix-matrix multiplication by
unfolding the cross-correlation between input and filter and explicitly copy
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 575b8c20b790cc..3e608b6b4fd5ce 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1553,6 +1553,10 @@ rewriteInIm2Col(RewriterBase &rewriter,
FailureOr<std::pair<Operation *, Operation *>>
rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp);
+/// Quantized variant of rewriteInIm2Col for linalg::Conv2DNchwFchwOp.
+FailureOr<std::pair<Operation *, Operation *>>
+rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwQOp convOp);
+
/// Convert linalg.conv_2d_nhwc_fhwc(_q) to linalg.conv_2d_nhwc_hwcf(_q) by
/// materializing transpose.
FailureOr<Operation *> transposeConv2D(RewriterBase &rewriter,
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 16b03ab5426a0d..be5bdd1f1aeb86 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -4408,6 +4408,9 @@ DiagnosedSilenceableFailure transform::ConvertConv2DToImg2ColOp::applyToOne(
.Case([&](linalg::Conv2DNchwFchwOp op) {
return rewriteInIm2Col(rewriter, op);
})
+ .Case([&](linalg::Conv2DNchwFchwQOp op) {
+ return rewriteInIm2Col(rewriter, op);
+ })
.Default([&](Operation *op) {
return rewriter.notifyMatchFailure(op, "not supported");
});
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
index e2dd0a6d79a733..bc1cf5f65cd759 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
@@ -410,8 +410,9 @@ rewriteInIm2Col(RewriterBase &rewriter,
transposedResult.getDefiningOp());
}
-FailureOr<std::pair<Operation *, Operation *>>
-rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp) {
+template <typename ConvOp>
+static FailureOr<std::pair<Operation *, Operation *>>
+rewriteNchwFchwInIm2col(RewriterBase &rewriter, ConvOp convOp) {
auto inputType = cast<ShapedType>(convOp.getInputs()[0].getType());
auto filterType = cast<ShapedType>(convOp.getInputs()[1].getType());
auto outputType = cast<ShapedType>(convOp.getOutputs()[0].getType());
@@ -452,7 +453,7 @@ rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp) {
SmallVector<ReassociationIndices> filterReassocIndices = {{0}, {1, 2, 3}};
auto reshapedFilterType =
- RankedTensorType::get({oc, ic * fh * fw}, inputType.getElementType(),
+ RankedTensorType::get({oc, ic * fh * fw}, filterType.getElementType(),
tensorFilterType.getEncoding());
Value reshapedFilter = tensor::CollapseShapeOp::create(
rewriter, loc, reshapedFilterType, filter, filterReassocIndices);
@@ -488,8 +489,10 @@ rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp) {
i2cToOperExprs.ohIndex = mIndicesExprs[0];
i2cToOperExprs.owIndex = mIndicesExprs[1];
Im2ColToInputDimsExprs inExprs = getIm2ColInputExpressions(
- i2cToOperExprs, llvm::to_vector(convOp.getStrides().getValues<int64_t>()),
- llvm::to_vector(convOp.getDilations().getValues<int64_t>()), rewriter);
+ i2cToOperExprs,
+ llvm::to_vector(convOp.getStrides().template getValues<int64_t>()),
+ llvm::to_vector(convOp.getDilations().template getValues<int64_t>()),
+ rewriter);
auto inMap =
AffineMap::inferFromExprList({ArrayRef{inExprs.bIndex, inExprs.cIndex,
inExprs.hIndex, inExprs.wIndex}},
@@ -506,6 +509,10 @@ rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp) {
linalg::YieldOp::create(nestedBuilder, nestedLoc, args[0]);
});
+ SmallVector<Value> preparedInputs(convOp.getInputs());
+ preparedInputs[0] = img2ColTensor.getResult(0);
+ preparedInputs[1] = reshapedFilter;
+
// Because the filter does not share the same batch dimension,
// the batch dimension is only used in indexing the input and output. Thus
// we cannot use existing linalg named ops like linalg.batch_matmul.
@@ -515,20 +522,28 @@ rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp) {
auto lhsMap = AffineMap::get(4, 0, {mDim, kDim}, context);
auto rhsMap = AffineMap::get(4, 0, {bDim, kDim, nDim}, context);
auto resultMap = AffineMap::get(4, 0, {bDim, mDim, nDim}, context);
+
+ SmallVector<AffineMap> map = {rhsMap, lhsMap};
+ for (unsigned i = 2; i < preparedInputs.size(); ++i) {
+ map.push_back(AffineMap::get(
+ /*dimCount=*/4,
+ /*symbolCount=*/0,
+ /*results=*/ArrayRef<AffineExpr>{}, context));
+ }
+ map.push_back(resultMap);
+
SmallVector<utils::IteratorType> genericIterators = {parallel, parallel,
parallel, reduction};
+
auto genericOp = linalg::GenericOp::create(
rewriter, loc, reshapedOutputType,
- /*inputs=*/ValueRange{reshapedFilter, img2ColTensor.getResult(0)},
- /*outputs=*/ValueRange{reshapedOutput},
- ArrayRef<AffineMap>{lhsMap, rhsMap, resultMap}, genericIterators,
- [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
- Value mul =
- createMul(loc, args[0], args[1], args[2].getType(), nestedBuilder);
- Value add = createAdd(loc, mul, args[2], nestedBuilder);
- linalg::YieldOp::create(nestedBuilder, nestedLoc, add);
- });
- Value result = genericOp.getResults().front();
+ /*inputs=*/preparedInputs,
+ /*outputs=*/ValueRange{reshapedOutput}, map, genericIterators);
+
+ rewriter.inlineRegionBefore(convOp->getRegion(0), genericOp.getRegion(),
+ genericOp.getRegion().begin());
+
+ Value result = genericOp.getResult(0);
auto reshapedResult = tensor::ExpandShapeOp::create(
rewriter, loc, outputType, result, outputReassocIndices);
@@ -539,6 +554,16 @@ rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp) {
reshapedResult.getOperation());
}
+FailureOr<std::pair<Operation *, Operation *>>
+rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp) {
+ return rewriteNchwFchwInIm2col(rewriter, convOp);
+}
+
+FailureOr<std::pair<Operation *, Operation *>>
+rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwQOp convOp) {
+ return rewriteNchwFchwInIm2col(rewriter, convOp);
+}
+
FailureOr<std::pair<Operation *, Operation *>>
rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcFhwcOp convOp) {
auto inputType = cast<ShapedType>(convOp.getInputs()[0].getType());
@@ -712,6 +737,19 @@ class ConvertConv2DNchwFchw final
}
};
+class ConvertConv2DNchwFchwQ final
+ : public OpRewritePattern<linalg::Conv2DNchwFchwQOp> {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(linalg::Conv2DNchwFchwQOp convOp,
+ PatternRewriter &rewriter) const override {
+ if (failed(rewriteInIm2Col(rewriter, convOp)))
+ return failure();
+ return success();
+ }
+};
+
class ConvertConv2DNhwcFhwc final
: public OpRewritePattern<linalg::Conv2DNhwcFhwcOp> {
public:
@@ -729,7 +767,8 @@ class ConvertConv2DNhwcFhwc final
void populateConvertConv2DToImg2ColPatterns(RewritePatternSet &patterns) {
MLIRContext *context = patterns.getContext();
patterns.insert<ConvertConv2DNhwcHwcf, ConvertDepthwiseConv2DNhwcHwc,
- ConvertConv2DNchwFchw, ConvertConv2DNhwcFhwc>(context);
+ ConvertConv2DNchwFchw, ConvertConv2DNchwFchwQ,
+ ConvertConv2DNhwcFhwc>(context);
}
} // end namespace linalg
} // end namespace mlir
diff --git a/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir b/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir
index 19b108014c7a19..1022cfbc635aca 100644
--- a/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir
+++ b/mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir
@@ -266,6 +266,16 @@ module attributes {transform.with_named_sequence} {
// CHECK-SAME: outs(%[[IT]] : tensor<8x36x196xf32>)
// CHECK: ^bb0(%[[IN:.+]]: f32, %out: f32):
// CHECK: linalg.yield %[[IN]] : f32
+// CHECK: } -> tensor<8x36x196xf32>
+// CHECK: %[[MATMUL:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[RHSMAP]], #[[LHSMAP]], #[[RESMAP]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"]
+// CHECK-SAME: ins(%[[IMG2COL]], %[[CS_FILTER]] : tensor<8x36x196xf32>, tensor<16x36xf32>)
+// CHECK-SAME: outs(%[[CS_RESULT]] : tensor<8x16x196xf32>)
+// CHECK: ^bb0(%[[PIXEL:.+]]: f32, %[[WEIGHT:.+]]: f32, %[[ACC:.+]]: f32):
+// CHECK: %[[PRODUCT:.+]] = arith.mulf %[[PIXEL]], %[[WEIGHT]] : f32
+// CHECK: %[[SUM:.+]] = arith.addf %[[ACC]], %[[PRODUCT]] : f32
+// CHECK: linalg.yield %[[SUM]] : f32
// CHECK: } -> tensor<8x16x196xf32>
// CHECK: %[[CS_FINAL:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0], [1], [2, 3]] output_shape [8, 16, 14, 14] : tensor<8x16x196xf32> into tensor<8x16x14x14xf32>
// CHECK: return %[[CS_FINAL]]
@@ -287,13 +297,69 @@ module attributes {transform.with_named_sequence} {
// -----
-// Check that the encoding on the filter (weights) tensor is propagated when applying the transform.
+// Im2col maps
+// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 floordiv 9, d2 floordiv 14 + (d1 mod 9) floordiv 3, d2 mod 14 + d1 mod 3)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+// CHECK-DAG: #[[LHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
+// CHECK-DAG: #[[RHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[ZPMAP:.+]] = affine_map<(d0, d1, d2, d3) -> ()>
+// CHECK-DAG: #[[RESMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+
+// CHECK: func.func @batch_nchw_conv_q
+// CHECK-SAME: (%[[INPUT:.+]]: tensor<8x4x16x16xi8>, %[[FILTER:.+]]: tensor<16x4x3x3xi8>, %[[INIT:.+]]: tensor<8x16x14x14xi32>, %[[INPUT_ZP:.+]]: i32, %[[FILTER_ZP:.+]]: i32)
+// CHECK-DAG: %[[CS_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0], [1, 2, 3]] : tensor<16x4x3x3xi8> into tensor<16x36xi8>
+// CHECK-DAG: %[[CS_RESULT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1], [2, 3]] : tensor<8x16x14x14xi32> into tensor<8x16x196xi32>
+// CHECK: %[[IT:.+]] = tensor.empty() : tensor<8x36x196xi8>
+// CHECK: %[[IMG2COL:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP1]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]
+// CHECK-SAME: ins(%[[INPUT]] : tensor<8x4x16x16xi8>)
+// CHECK-SAME: outs(%[[IT]] : tensor<8x36x196xi8>)
+// CHECK: ^bb0(%[[IN:.+]]: i8, %out: i8):
+// CHECK: linalg.yield %[[IN]] : i8
+// CHECK: } -> tensor<8x36x196xi8>
+// CHECK: %[[MATMUL:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[RHSMAP]], #[[LHSMAP]], #[[ZPMAP]], #[[ZPMAP]], #[[RESMAP]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"]
+// CHECK-SAME: ins(%[[IMG2COL]], %[[CS_FILTER]], %[[INPUT_ZP]], %[[FILTER_ZP]] : tensor<8x36x196xi8>, tensor<16x36xi8>, i32, i32)
+// CHECK-SAME: outs(%[[CS_RESULT]] : tensor<8x16x196xi32>)
+// CHECK: ^bb0(%[[PIXEL:.+]]: i8, %[[WEIGHT:.+]]: i8, %[[PIXEL_ZP:.+]]: i32, %[[WEIGHT_ZP:.+]]: i32, %[[ACC:.+]]: i32):
+// CHECK: %[[PIXEL_EXT:.+]] = arith.extsi %[[PIXEL]] : i8 to i32
+// CHECK: %[[PIXEL_CENTERED:.+]] = arith.subi %[[PIXEL_EXT]], %[[PIXEL_ZP]] : i32
+// CHECK: %[[WEIGHT_EXT:.+]] = arith.extsi %[[WEIGHT]] : i8 to i32
+// CHECK: %[[WEIGHT_CENTERED:.+]] = arith.subi %[[WEIGHT_EXT]], %[[WEIGHT_ZP]] : i32
+// CHECK: %[[PRODUCT:.+]] = arith.muli %[[PIXEL_CENTERED]], %[[WEIGHT_CENTERED]] : i32
+// CHECK: %[[SUM:.+]] = arith.addi %[[ACC]], %[[PRODUCT]] : i32
+// CHECK: linalg.yield %[[SUM]] : i32
+// CHECK: } -> tensor<8x16x196xi32>
+// CHECK: %[[CS_FINAL:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0], [1], [2, 3]] output_shape [8, 16, 14, 14] : tensor<8x16x196xi32> into tensor<8x16x14x14xi32>
+// CHECK: return %[[CS_FINAL]]
+func.func @batch_nchw_conv_q(%arg0: tensor<8x4x16x16xi8>, %arg1: tensor<16x4x3x3xi8>, %arg2: tensor<8x16x14x14xi32>, %input_zp: i32, %filter_zp: i32) -> tensor<8x16x14x14xi32> {
+ %0 = linalg.conv_2d_nchw_fchw_q
+ {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
+ ins(%arg0, %arg1, %input_zp, %filter_zp: tensor<8x4x16x16xi8>, tensor<16x4x3x3xi8>, i32, i32)
+ outs(%arg2: tensor<8x16x14x14xi32>) -> tensor<8x16x14x14xi32>
+ return %0 : tensor<8x16x14x14xi32>
+}
+
+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_q"]} 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
+ }
+}
+
+// -----
+
+// Check that the encoding on the filter (weights) tensor is propagated when applying the transform.
// CHECK: func.func @batch_nchw_conv_with_filter_encoding(%[[INPUT:.+]]: tensor<8x4x16x16xf32>, %[[FILTER:.*]]: tensor<16x4x3x3xf32, 42 : i32>, %[[OUTPUT:.*]]: tensor<8x16x14x14xf32>)
// CHECK-DAG: %[[COLLAPSED_FILTER:.+]] = tensor.collapse_shape %[[FILTER]]
// CHECK-SAME{LITERAL}: [[0], [1, 2, 3]] : tensor<16x4x3x3xf32, 42 : i32> into tensor<16x36xf32, 42 : i32>
// CHECK: %[[COL_TENSOR:.+]] = linalg.generic {{.*}} ins(%[[INPUT]] : tensor<8x4x16x16xf32>)
-// CHECK: %[[MATMUL_RESULT:.+]] = linalg.generic {{.*}} ins(%[[COLLAPSED_FILTER]], %[[COL_TENSOR]] : tensor<16x36xf32, 42 : i32>, tensor<8x36x196xf32>)
+// CHECK: %[[MATMUL_RESULT:.+]] = linalg.generic {{.*}} ins(%[[COL_TENSOR]], %[[COLLAPSED_FILTER]] : tensor<8x36x196xf32>, tensor<16x36xf32, 42 : i32>)
func.func @batch_nchw_conv_with_filter_encoding(%arg0: tensor<8x4x16x16xf32>, %arg1: tensor<16x4x3x3xf32, 42 : i32>, %arg2: tensor<8x16x14x14xf32>) -> tensor<8x16x14x14xf32> {
%0 = linalg.conv_2d_nchw_fchw
{dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
@@ -372,7 +438,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// Check that the encoding on the filter (weights) tensor is propagated when applying the transform.
+// Check that the encoding on the filter (weights) tensor is propagated when applying the transform.
// CHECK: func.func @conv_2d_nhwc_fhwc_with_filter_encoding(%[[INPUT:.+]]: tensor<1x16x16x4xf32>, %[[FILTER:.*]]: tensor<16x3x3x4xf32, 42 : i32>, %[[OUTPUT:.*]]: tensor<1x14x14x16xf32>)
// CHECK-DAG: %[[COLLAPSED_FILTER:.+]] = tensor.collapse_shape %[[FILTER]]
@@ -614,7 +680,7 @@ module attributes {transform.with_named_sequence} {
// CHECK-SAME: ins(%[[INPUT]] : tensor<8x4x16x16xf32>)
// CHECK-SAME: outs(%[[IT]] : tensor<8x36x120xf32>)
// CHECK: %[[MATMUL:.+]] = linalg.generic
-// CHECK-SAME: ins(%[[CS_FILTER]], %[[IMG2COL]] : tensor<16x36xf32>, tensor<8x36x120xf32>)
+// CHECK-SAME: ins(%[[IMG2COL]], %[[CS_FILTER]] : tensor<8x36x120xf32>, tensor<16x36xf32>)
// CHECK-SAME: outs(%[[CS_RESULT]] : tensor<8x16x120xf32>)
// CHECK: %[[CS_FINAL:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0], [1], [2, 3]] output_shape [8, 16, 10, 12] : tensor<8x16x120xf32> into tensor<8x16x10x12xf32>
// CHECK: return %[[CS_FINAL]]
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/test-conv-2d-nchw-fchw-q-img2col.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/test-conv-2d-nchw-fchw-q-img2col.mlir
new file mode 100644
index 00000000000000..ae5287502c786a
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/test-conv-2d-nchw-fchw-q-img2col.mlir
@@ -0,0 +1,76 @@
+// UNSUPPORTED: asan
+// DEFINE: %{lower} = -test-transform-dialect-erase-schedule \
+// DEFINE: -one-shot-bufferize="bufferize-function-boundaries" \
+// DEFINE: -buffer-deallocation-pipeline -convert-bufferization-to-memref \
+// DEFINE: -convert-linalg-to-loops -convert-scf-to-cf \
+// DEFINE: -expand-strided-metadata -lower-affine -convert-arith-to-llvm \
+// DEFINE: -convert-scf-to-cf -finalize-memref-to-llvm -convert-func-to-llvm \
+// DEFINE: -convert-cf-to-llvm -reconcile-unrealized-casts
+// DEFINE: %{run} = mlir-runner -e main -entry-point-result=void \
+// DEFINE: -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
+
+// Check the reference convolution
+// RUN: mlir-opt %s %{lower} | %{run} | FileCheck %s
+
+// RUN: mlir-opt %s -transform-interpreter -test-transform-dialect-erase-schedule -o %t
+// RUN: FileCheck %s --check-prefix=IM2COL --implicit-check-not=linalg.conv_2d_nchw_fchw_q < %t
+// RUN: mlir-opt %t %{lower} | %{run} | FileCheck %s
+
+func.func private @printMemrefI32(tensor<*xi32>)
+
+// IM2COL-LABEL: func.func @conv(
+// IM2COL: linalg.generic
+// IM2COL: linalg.generic
+// IM2COL: tensor.expand_shape
+// IM2COL: return
+func.func @conv(%input: tensor<2x2x3x3xi8>,
+ %filter: tensor<3x2x2x2xi8>,
+ %init: tensor<2x3x2x2xi32>,
+ %input_zp: i32, %filter_zp: i32) -> tensor<2x3x2x2xi32> {
+ %result = linalg.conv_2d_nchw_fchw_q
+ {strides = dense<1> : tensor<2xi64>, dilations = dense<1> : tensor<2xi64>}
+ ins(%input, %filter, %input_zp, %filter_zp :
+ tensor<2x2x3x3xi8>, tensor<3x2x2x2xi8>, i32, i32)
+ outs(%init : tensor<2x3x2x2xi32>) -> tensor<2x3x2x2xi32>
+ return %result : tensor<2x3x2x2xi32>
+}
+
+func.func @main() {
+ %input = arith.constant dense<[
+ [[[-4, -3, -2],[-1, 0, 1], [2, 3, 4]],
+ [[5, 4, 3], [2, 1, 0], [-1, -2, -3]]],
+ [[[6, -5, 4], [-3, 2, -1], [0, 1, -2]],
+ [[-7, 6, -5], [4, -3, 2], [-1, 0, 1]]]
+ ]> : tensor<2x2x3x3xi8>
+ %filter = arith.constant dense<[
+ [[[1, -2], [3, 0]], [[-1, 2], [-3, 4]]],
+ [[[-4, 1], [2, -1]], [[3, -2], [0, 1]]],
+ [[[2, 0], [-1, -3]], [[-2, 4], [1, -4]]]
+ ]> : tensor<3x2x2x2xi8>
+ %init = arith.constant dense<7> : tensor<2x3x2x2xi32>
+ %input_zp = arith.constant 3 : i32
+ %filter_zp = arith.constant -2 : i32
+ %result = func.call @conv(%input, %filter, %init, %input_zp, %filter_zp) :
+ (tensor<2x2x3x3xi8>, tensor<3x2x2x2xi8>, tensor<2x3x2x2xi32>, i32, i32)
+ -> tensor<2x3x2x2xi32>
+
+ %flat = tensor.collapse_shape %result [[0], [1, 2, 3]] :
+ tensor<2x3x2x2xi32> into tensor<2x12xi32>
+ %unranked = tensor.cast %flat : tensor<2x12xi32> to tensor<*xi32>
+ func.call @printMemrefI32(%unranked) : (tensor<*xi32>) -> ()
+ return
+}
+
+// CHECK: Unranked Memref
+// CHECK-SAME: rank = 2 offset = 0 sizes = [2, 12] strides = [12, 1] data =
+// CHECK-NEXT: [-45, -45, -45, -45, -14, -18, -26, -30, -27, -28, -30, -31]
+// CHECK-NEXT: [-51, -59, -67, -35, -114, 18, -10, -58, 31, -84, -62, -13]
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %conv = transform.structured.match ops{["linalg.conv_2d_nchw_fchw_q"]}
+ in %root : (!transform.any_op) -> !transform.any_op
+ %col, %result = transform.structured.convert_conv2d_to_img2col %conv :
+ (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
>From 7116574b8d30c1394b7471af136223be901b80d6 Mon Sep 17 00:00:00 2001
From: Dmitry Manjiev <mandzievdmitrij at gmail.com>
Date: Tue, 29 Sep 2026 08:47:44 +0300
Subject: [PATCH 2/2] [mlir][linalg] Address review comments
---
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h | 2 +-
.../Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp | 7 ++++++-
2 files changed, 7 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 3e608b6b4fd5ce..d5699bf8d36e33 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1553,7 +1553,7 @@ rewriteInIm2Col(RewriterBase &rewriter,
FailureOr<std::pair<Operation *, Operation *>>
rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwOp convOp);
-/// Quantized variant of rewriteInIm2Col for linalg::Conv2DNchwFchwOp.
+/// Quantized variant of rewriteInIm2Col for linalg::Conv2DNchwFchwQOp.
FailureOr<std::pair<Operation *, Operation *>>
rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNchwFchwQOp convOp);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
index bc1cf5f65cd759..ba6dc514cdee3c 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
@@ -23,6 +23,7 @@
#include "llvm/ADT/SmallVectorExtras.h"
#include <cassert>
#include <utility>
+#include <type_traits>
namespace mlir {
namespace linalg {
@@ -412,7 +413,11 @@ rewriteInIm2Col(RewriterBase &rewriter,
template <typename ConvOp>
static FailureOr<std::pair<Operation *, Operation *>>
-rewriteNchwFchwInIm2col(RewriterBase &rewriter, ConvOp convOp) {
+rewriteNchwFchwInIm2Col(RewriterBase &rewriter, ConvOp convOp) {
+ static_assert(std::is_same_v<ConvOp, linalg::Conv2DNchwFchwOp> ||
+ std::is_same_v<ConvOp, linalg::Conv2DNchwFchwQOp>,
+ "expected Conv2DNchwFchwOp or Conv2DNchwFchwQOp");
+
auto inputType = cast<ShapedType>(convOp.getInputs()[0].getType());
auto filterType = cast<ShapedType>(convOp.getInputs()[1].getType());
auto outputType = cast<ShapedType>(convOp.getOutputs()[0].getType());
More information about the Mlir-commits
mailing list