[Mlir-commits] [mlir] [mlir][tosa] Don't broadcast non-constant mul shift and negate zero points (PR #233121)
Maher Ahmed Maher Shola
llvmlistbot at llvm.org
Sat Oct 17 10:56:11 PDT 2026
https://github.com/Maher-star-eng created https://github.com/llvm/llvm-project/pull/233121
When the `tosa.mul` shift or the `tosa.negate` zero points are not constants,
the TosaToLinalg lowering passed them to the elementwise broadcasting helpers,
which assume every operand has the rank of the input. These operands are always
`tensor<1xT>`, so the lowering asserted for rank-0 inputs, for dynamic shapes,
and for inputs with a size-1 dimension after the first (e.g. `3x1`).
Separately, when only one negate zero point was non-constant, the body read both
zero points as block arguments although neither was passed to the
`linalg.generic`.
This patch keeps these single-element operands out of broadcasting and passes
them to the `linalg.generic` with an indexing map that always reads element 0.
Negate passes both zero points whenever either is non-constant, matching what
the body expects. Constant shifts and zero points are handled as before.
The new lit tests cover rank-0, size-1 dimension, dynamic and mixed
constant/non-constant zero points, based on the reproducers from the issues.
AI disclosure: I found this bug with a testing harness built with Claude
(Claude Code). The fix, the tests are built tested my self
Fixes #231418
Fixes #231834
Fixes #231835
>From 029bfcd53e334cd6ccdcf40ed443b30befd0bed9 Mon Sep 17 00:00:00 2001
From: Maher Ahmed Maher Shola <maher.shola.officer at gmail.com>
Date: Sat, 17 Oct 2026 20:06:33 +0300
Subject: [PATCH] [mlir][tosa] Don't broadcast non-constant mul shift and
negate zero points
Fixes #231418
Fixes #231834
Fixes #231835
Assisted-by: Claude
---
.../Conversion/TosaToLinalg/TosaToLinalg.cpp | 57 ++++++++++++-------
.../TosaToLinalg/tosa-to-linalg.mlir | 55 ++++++++++++++++++
2 files changed, 93 insertions(+), 19 deletions(-)
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 25b4b360723b22..4647d60987b350 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -992,11 +992,10 @@ broadcastDynamicDimensions(PatternRewriter &rewriter, Location loc,
});
}
-static LogicalResult
-emitElementwiseComputation(ConversionPatternRewriter &rewriter, Location loc,
- Operation *operation, ValueRange operands,
- ArrayRef<OpFoldResult> targetShape,
- const TypeConverter &converter) {
+static LogicalResult emitElementwiseComputation(
+ ConversionPatternRewriter &rewriter, Location loc, Operation *operation,
+ ValueRange operands, ValueRange scalarOperands,
+ ArrayRef<OpFoldResult> targetShape, const TypeConverter &converter) {
// Generate output tensor
auto resultType = cast_or_null<RankedTensorType>(
converter.convertType(operation->getResultTypes().front()));
@@ -1026,12 +1025,23 @@ emitElementwiseComputation(ConversionPatternRewriter &rewriter, Location loc,
}
return AffineMap::get(rank, 0, affineExprs, rewriter.getContext());
});
+ // Scalar operands are single-element tensors: every iteration reads their
+ // first element.
+ for (Value scalar : scalarOperands) {
+ auto scalarRank = cast<ShapedType>(scalar.getType()).getRank();
+ SmallVector<AffineExpr> zeros(scalarRank,
+ rewriter.getAffineConstantExpr(0));
+ affineMaps.push_back(AffineMap::get(rank, 0, zeros, rewriter.getContext()));
+ }
affineMaps.push_back(rewriter.getMultiDimIdentityMap(rank));
+ SmallVector<Value> inputs(operands);
+ llvm::append_range(inputs, scalarOperands);
+
// Emit 'linalg.generic' op
bool encounteredError = false;
auto linalgOp = linalg::GenericOp::create(
- rewriter, loc, outputTensor.getType(), operands, outputTensor, affineMaps,
+ rewriter, loc, outputTensor.getType(), inputs, outputTensor, affineMaps,
getNParallelLoopsAttrs(rank),
[&](OpBuilder &opBuilder, Location loc, ValueRange blockArgs) {
Value opResult = createLinalgBodyCalculationForElementwiseOp(
@@ -1054,26 +1064,34 @@ emitElementwiseComputation(ConversionPatternRewriter &rewriter, Location loc,
return success();
}
+// The shift of tosa.mul and the zero points of tosa.negate are single-element
+// tensors, so they never take part in broadcasting.
static ValueRange getBroadcastableOperands(Operation *operation,
ValueRange operands) {
- // Shift cannot broadcast
+ if (isa<tosa::MulOp>(operation))
+ return operands.take_front(2);
+ if (isa<tosa::NegateOp>(operation))
+ return operands.take_front(1);
+ return operands;
+}
+
+// Returns the single-element operands that the body reads as block arguments:
+// a non-constant tosa.mul shift, or both tosa.negate zero points when either of
+// them is non-constant. Constant ones are created in the body instead.
+static ValueRange getScalarOperands(Operation *operation, ValueRange operands) {
if (isa<tosa::MulOp>(operation)) {
DenseElementsAttr shiftElems;
- // Shift cannot broadcast when it is constant
if (matchPattern(operation->getOperand(2), m_Constant(&shiftElems)))
- return operands.take_front(2);
- else
- return operands.take_front(3);
+ return {};
+ return operands.drop_front(2);
}
if (auto negate = dyn_cast<tosa::NegateOp>(operation)) {
- FailureOr<int64_t> maybeInZp = negate.getInput1ZeroPoint();
- FailureOr<int64_t> maybeOutZp = negate.getOutputZeroPoint();
- if (failed(maybeOutZp) && failed(maybeInZp))
- return operands;
- // Input1_zp and output_zp cannot broadcast when they are constants.
- return operands.take_front(1);
+ if (succeeded(negate.getInput1ZeroPoint()) &&
+ succeeded(negate.getOutputZeroPoint()))
+ return {};
+ return operands.drop_front(1);
}
- return operands;
+ return {};
}
static LogicalResult
@@ -1093,13 +1111,14 @@ elementwiseMatchAndRewriteHelper(Operation *operation, ValueRange operands,
IndexPool indexPool;
auto loc = operation->getLoc();
auto operandsToBroadcast = getBroadcastableOperands(operation, operands);
+ auto scalarOperands = getScalarOperands(operation, operands);
auto [targetShape, masterOperands] =
computeTargetShape(rewriter, loc, indexPool, operandsToBroadcast);
auto broadcastOperands =
broadcastDynamicDimensions(rewriter, loc, indexPool, operandsToBroadcast,
targetShape, masterOperands);
return emitElementwiseComputation(rewriter, loc, operation, broadcastOperands,
- targetShape, converter);
+ scalarOperands, targetShape, converter);
}
// Returns the constant initial value for a given reduction operation. The
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index 67a51e1e2e80f9..8a29c69ef5c921 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -2781,6 +2781,61 @@ func.func @negate_i64_no_noop_cast(%arg0: tensor<4xi64>, %in_zp: tensor<1xi64>,
// -----
+// A non-constant mul shift or negate zero point is a single-element tensor: it
+// does not take part in broadcasting and every iteration reads its element 0.
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<() -> ()>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<() -> (0)>
+// CHECK-LABEL: @mul_non_const_shift_rank0
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<i32>, %[[SHIFT:[0-9a-zA-Z_]*]]: tensor<1xi8>
+func.func @mul_non_const_shift_rank0(%arg0: tensor<i32>, %shift: tensor<1xi8>) -> tensor<i32> {
+ // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP1]], #[[$MAP0]]], iterator_types = []} ins(%[[ARG0]], %[[ARG0]], %[[SHIFT]] : tensor<i32>, tensor<i32>, tensor<1xi8>)
+ // CHECK: tosa.apply_scale
+ %0 = tosa.mul %arg0, %arg0, %shift : (tensor<i32>, tensor<i32>, tensor<1xi8>) -> tensor<i32>
+ return %0 : tensor<i32>
+}
+
+// -----
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (0)>
+// CHECK-LABEL: @negate_non_const_zp_size1_dim
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<1x1xf32>, %[[ZP:[0-9a-zA-Z_]*]]: tensor<1xf32>
+func.func @negate_non_const_zp_size1_dim(%arg0: tensor<1x1xf32>, %zp: tensor<1xf32>) -> tensor<1x1xf32> {
+ // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins(%[[ARG0]], %[[ZP]], %[[ZP]] : tensor<1x1xf32>, tensor<1xf32>, tensor<1xf32>)
+ // CHECK: arith.negf
+ %0 = tosa.negate %arg0, %zp, %zp : (tensor<1x1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x1xf32>
+ return %0 : tensor<1x1xf32>
+}
+
+// -----
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (0)>
+// CHECK-LABEL: @negate_non_const_zp_dynamic
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<?xi8>, %[[IN_ZP:[0-9a-zA-Z_]*]]: tensor<1xi8>, %[[OUT_ZP:[0-9a-zA-Z_]*]]: tensor<1xi8>
+func.func @negate_non_const_zp_dynamic(%arg0: tensor<?xi8>, %in_zp: tensor<1xi8>, %out_zp: tensor<1xi8>) -> tensor<?xi8> {
+ // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%[[ARG0]], %[[IN_ZP]], %[[OUT_ZP]] : tensor<?xi8>, tensor<1xi8>, tensor<1xi8>)
+ %0 = tosa.negate %arg0, %in_zp, %out_zp : (tensor<?xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<?xi8>
+ return %0 : tensor<?xi8>
+}
+
+// -----
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (0)>
+// CHECK-LABEL: @negate_mixed_const_zp
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<1xi32>, %[[IN_ZP:[0-9a-zA-Z_]*]]: tensor<1xi32>
+func.func @negate_mixed_const_zp(%arg0: tensor<1xi32>, %in_zp: tensor<1xi32>) -> tensor<1xi32> {
+ // CHECK: %[[OUT_ZP:.*]] = tosa.const
+ // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%[[ARG0]], %[[IN_ZP]], %[[OUT_ZP]] : tensor<1xi32>, tensor<1xi32>, tensor<1xi32>)
+ %out_zp = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %0 = tosa.negate %arg0, %in_zp, %out_zp : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+ return %0 : tensor<1xi32>
+}
+
+// -----
+
// The TOSA identity for a float REDUCE_MIN is maximum_s<in_out_t>(), which
// Table 5 of the specification gives as +infinity for fp16_t/bf16_t/fp32_t.
// CHECK-LABEL: @reduce_min_f16
More information about the Mlir-commits
mailing list