[Mlir-commits] [mlir] [mlir][linalg] Allow extra DPS inputs in convolution dim inference (PR #198462)
Ahmad Tameem
llvmlistbot at llvm.org
Wed Jul 22 06:54:26 PDT 2026
https://github.com/Tameem-10xE updated https://github.com/llvm/llvm-project/pull/198462
>From 57a844c35d3bc82c5c159c5dd827df6b1255565b Mon Sep 17 00:00:00 2001
From: Tameem-10xE <ahmad.tameem at 10xengineers.ai>
Date: Wed, 22 Jul 2026 18:51:37 +0500
Subject: [PATCH] [mlir][linalg] Allow extra DPS inputs in convolution dim
inference
Signed-off-by: Tameem-10xE <ahmad.tameem at 10xengineers.ai>
---
.../Dialect/Linalg/IR/LinalgInterfaces.cpp | 2 +-
.../Linalg/InferConvolutionDimsTest.cpp | 130 +++++++++++++++++-
2 files changed, 130 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index e3f4988b3ab76..2f04ce66fa326 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -884,7 +884,7 @@ static FailureOr<ConvolutionDimensions> inferConvolutionDimsImpl(
/// Returns a failure if `output_image` (and implicitly `filter_loop`) is empty.
FailureOr<ConvolutionDimensions>
mlir::linalg::inferConvolutionDims(LinalgOp linalgOp) {
- if (linalgOp.getNumDpsInits() != 1 || linalgOp.getNumDpsInputs() != 2)
+ if (linalgOp.getNumDpsInits() != 1 || linalgOp.getNumDpsInputs() < 2)
return failure();
auto indexingMaps = linalgOp.getIndexingMapsArray();
diff --git a/mlir/unittests/Dialect/Linalg/InferConvolutionDimsTest.cpp b/mlir/unittests/Dialect/Linalg/InferConvolutionDimsTest.cpp
index cab17f2a6e9c1..1dd4e7cb94755 100644
--- a/mlir/unittests/Dialect/Linalg/InferConvolutionDimsTest.cpp
+++ b/mlir/unittests/Dialect/Linalg/InferConvolutionDimsTest.cpp
@@ -130,6 +130,88 @@ createConv2DWithSwappedFilterLoops(OpBuilder &builder,
});
}
+/// Creates a quantized 2D convolution as a linalg.generic with the same
+/// input/filter access pattern as linalg.conv_2d, plus extra scalar
+/// input/filter zero-point operands (mirroring the `*_q` named ops).
+/// The zero points are hard-coded constants since their values do not
+/// affect dimension inference.
+///
+/// Loop order:
+/// d0 = output height (oh), parallel
+/// d1 = output width (ow), parallel
+/// d2 = kernel height (kh), reduction
+/// d3 = kernel width (kw), reduction
+///
+/// Indexing maps:
+/// input: (d0 + d2, d1 + d3)
+/// filter: (d2, d3)
+/// input zp: scalar
+/// filter zp: scalar
+/// output: (d0, d1)
+///
+/// Semantic pairing: d0 <-> d2, d1 <-> d3
+static linalg::GenericOp createQConv2DOp(OpBuilder &builder, int64_t oh,
+ int64_t ow, int64_t kh, int64_t kw) {
+ Location loc = builder.getUnknownLoc();
+ MLIRContext *ctx = builder.getContext();
+
+ auto i8Type = builder.getI8Type();
+ auto i32Type = builder.getI32Type();
+
+ int64_t ih = oh + kh - 1;
+ int64_t iw = ow + kw - 1;
+
+ auto inputType = RankedTensorType::get({ih, iw}, i8Type);
+ auto filterType = RankedTensorType::get({kh, kw}, i8Type);
+ auto outputType = RankedTensorType::get({oh, ow}, i32Type);
+
+ Value input = tensor::EmptyOp::create(builder, loc, inputType.getShape(),
+ inputType.getElementType());
+ Value filter = tensor::EmptyOp::create(builder, loc, filterType.getShape(),
+ filterType.getElementType());
+
+ // Non-zero input and filter zero points (their values are irrelevant to
+ // dimension inference).
+ Value inputZeroPoint = arith::ConstantIntOp::create(builder, loc, 7, 32);
+ Value filterZeroPoint = arith::ConstantIntOp::create(builder, loc, 9, 32);
+
+ Value output = tensor::EmptyOp::create(builder, loc, outputType.getShape(),
+ outputType.getElementType());
+
+ AffineExpr d0, d1, d2, d3;
+ bindDims(ctx, d0, d1, d2, d3);
+
+ auto inputMap = AffineMap::get(4, 0, {d0 + d2, d1 + d3}, ctx);
+ auto filterMap = AffineMap::get(4, 0, {d2, d3}, ctx);
+ auto scalarMap = AffineMap::get(4, 0, ArrayRef<AffineExpr>{}, ctx);
+ auto outputMap = AffineMap::get(4, 0, {d0, d1}, ctx);
+
+ SmallVector<AffineMap> indexingMaps = {inputMap, filterMap, scalarMap,
+ scalarMap, outputMap};
+
+ SmallVector<utils::IteratorType> iterTypes = {
+ utils::IteratorType::parallel, utils::IteratorType::parallel,
+ utils::IteratorType::reduction, utils::IteratorType::reduction};
+
+ return linalg::GenericOp::create(
+ builder, loc, outputType,
+ ValueRange{input, filter, inputZeroPoint, filterZeroPoint},
+ ValueRange{output}, indexingMaps, iterTypes,
+ [&](OpBuilder &b, Location loc, ValueRange args) {
+ Value inputI32 =
+ arith::ExtSIOp::create(b, loc, b.getI32Type(), args[0]);
+ Value filterI32 =
+ arith::ExtSIOp::create(b, loc, b.getI32Type(), args[1]);
+
+ inputI32 = arith::SubIOp::create(b, loc, inputI32, args[2]);
+ filterI32 = arith::SubIOp::create(b, loc, filterI32, args[3]);
+
+ Value mul = arith::MulIOp::create(b, loc, inputI32, filterI32);
+ Value add = arith::AddIOp::create(b, loc, args[4], mul);
+ linalg::YieldOp::create(b, loc, add);
+ });
+}
+
TEST_F(InferConvolutionDimsTest, Conv2DPairing) {
// Use non-square kernel to ensure dimension swapping is tested properly.
const int64_t oh = 6, ow = 12, kh = 3, kw = 5;
@@ -191,7 +273,7 @@ static void expectConvDimsEq(const ConvolutionDimensions &lhs,
}
/// Verify that inferring convolution dimensions from indexing maps produces
-/// same result as inferring thrm directly from the convolution op.
+/// same result as inferring them directly from the convolution op.
TEST_F(InferConvolutionDimsTest, MapsOverloadMatchesOpOverload) {
OpBuilder builder(ctx.get());
OwningOpRef<ModuleOp> module = ModuleOp::create(builder.getUnknownLoc());
@@ -312,4 +394,50 @@ TEST_F(InferConvolutionDimsTest, MapsOverloadRejectsInvalidInput) {
}
}
+/// Extra scalar zero-point operands must not affect convolution dim inference.
+TEST_F(InferConvolutionDimsTest, QConv2DWithZeroPoints) {
+ // Use non-square kernel to ensure dimension swapping is tested properly.
+ const int64_t oh = 6, ow = 12, kh = 3, kw = 5;
+
+ // Create a module to own all test operations and ensure proper cleanup.
+ OpBuilder builder(ctx.get());
+ OwningOpRef<ModuleOp> module = ModuleOp::create(builder.getUnknownLoc());
+ builder.setInsertionPointToStart(module->getBody());
+
+ // Create Quantize ConvOp with two extra scalar zero-point inputs.
+ linalg::GenericOp qConvOp = createQConv2DOp(builder, oh, ow, kh, kw);
+
+ // The qconv op should have:
+ // input, filter, input zero point, filter zero point
+ ASSERT_EQ(qConvOp.getNumDpsInputs(), 4u);
+ ASSERT_EQ(qConvOp.getNumDpsInits(), 1u);
+
+ auto indexingMaps = qConvOp.getIndexingMapsArray();
+ ASSERT_EQ(indexingMaps.size(), static_cast<size_t>(qConvOp.getNumDpsInputs() +
+ qConvOp.getNumDpsInits()));
+
+ // The two extra quantized conv operands must be scalar inputs.
+ EXPECT_EQ(indexingMaps[2].getNumResults(), 0u);
+ EXPECT_EQ(indexingMaps[3].getNumResults(), 0u);
+ EXPECT_EQ(indexingMaps[2].getNumDims(), 4u);
+ EXPECT_EQ(indexingMaps[3].getNumDims(), 4u);
+
+ FailureOr<ConvolutionDimensions> qConvDims = inferConvolutionDims(qConvOp);
+ ASSERT_TRUE(succeeded(qConvDims));
+
+ // The scalar zero-point operands must not affect dimension inference: the
+ // result must match a plain conv_2d, with standard pairing d0 <-> d2
+ // (oh <-> kh) and d1 <-> d3 (ow <-> kw), unit strides/dilations, and no
+ // batch/channel/depth dims.
+ ConvolutionDimensions expected{/*batch=*/{},
+ /*outputImage=*/{0, 1},
+ /*outputChannel=*/{},
+ /*filterLoop=*/{2, 3},
+ /*inputChannel=*/{},
+ /*depth=*/{},
+ /*strides=*/{1, 1},
+ /*dilations=*/{1, 1}};
+ expectConvDimsEq(expected, *qConvDims);
+}
+
} // namespace
More information about the Mlir-commits
mailing list