[Mlir-commits] [mlir] 8edba1a - [mlir][arith] Add subui_extended op (#197376)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu May 14 08:52:20 PDT 2026
Author: Arseniy Obolenskiy
Date: 2026-05-14T17:52:14+02:00
New Revision: 8edba1a3128ed3c6ab104470ef3f3567bf73c516
URL: https://github.com/llvm/llvm-project/commit/8edba1a3128ed3c6ab104470ef3f3567bf73c516
DIFF: https://github.com/llvm/llvm-project/commit/8edba1a3128ed3c6ab104470ef3f3567bf73c516.diff
LOG: [mlir][arith] Add subui_extended op (#197376)
Unsigned extended subtraction mirroring addui_extended. Gives a
primitive for unsigned underflow that maps cleanly to
llvm.usub.with.overflow and spirv.ISub
Added:
Modified:
mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
mlir/lib/Dialect/Arith/IR/ArithOps.cpp
mlir/test/Dialect/Arith/canonicalize.mlir
mlir/test/Dialect/Arith/invalid.mlir
mlir/test/Dialect/Arith/ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
index f359070b6842f..fa85b840e2707 100644
--- a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
+++ b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
@@ -359,6 +359,59 @@ def Arith_AddUIExtendedOp : Arith_Op<"addui_extended", [Pure, Commutative,
}];
}
+//===----------------------------------------------------------------------===//
+// SubUIExtendedOp
+//===----------------------------------------------------------------------===//
+
+def Arith_SubUIExtendedOp : Arith_Op<"subui_extended", [Pure,
+ AllTypesMatch<["lhs", "rhs", "
diff "]>]> {
+ let summary = [{
+ extended unsigned integer subtraction operation returning
diff erence and
+ borrow bit
+ }];
+
+ let description = [{
+ Performs (N+1)-bit subtraction on zero-extended operands. Returns two
+ results: the N-bit
diff erence (same type as both operands), and the borrow
+ bit (boolean-like), where `1` indicates unsigned subtraction underflow
+ (i.e. `lhs < rhs` when interpreted as unsigned), while `0` indicates no
+ underflow.
+
+ Example:
+
+ ```mlir
+ // Scalar subtraction.
+ %
diff , %borrow = arith.subui_extended %b, %c : i64, i1
+
+ // Vector element-wise subtraction.
+ %d:2 = arith.subui_extended %e, %f : vector<4xi32>, vector<4xi1>
+
+ // Tensor element-wise subtraction.
+ %x:2 = arith.subui_extended %y, %z : tensor<4x?xi8>, tensor<4x?xi1>
+ ```
+ }];
+
+ let arguments = (ins Arith_SignlessIntegerOrIndexLike:$lhs, Arith_SignlessIntegerOrIndexLike:$rhs);
+ let results = (outs Arith_SignlessIntegerOrIndexLike:$
diff , BoolLike:$borrow);
+ let assemblyFormat = [{
+ $lhs `,` $rhs attr-dict `:` type($
diff ) `,` type($borrow)
+ }];
+
+ let builders = [
+ OpBuilder<(ins "Value":$lhs, "Value":$rhs), [{
+ build($_builder, $_state, lhs.getType(), ::getI1SameShape(lhs.getType()),
+ lhs, rhs);
+ }]>
+ ];
+
+ let hasFolder = 1;
+ let hasCanonicalizer = 1;
+
+ let extraClassDeclaration = [{
+ std::optional<SmallVector<int64_t, 4>> getShapeForUnroll();
+ }];
+}
+
//===----------------------------------------------------------------------===//
// SubIOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
index 66f4ace265201..cfd25f28b3e65 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
+++ b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
@@ -105,6 +105,17 @@ def AddUIExtendedToAddI:
[(Arith_AddIOp $x, $y, DefOverflow), (replaceWithValue $x)],
[(Constraint<CPred<"$0.getUses().empty()">> $res__1)]>;
+//===----------------------------------------------------------------------===//
+// SubUIExtendedOp
+//===----------------------------------------------------------------------===//
+
+// subui_extended(x, y) -> [subi(x, y), x], when the `borrow` result has no
+// uses. Since the 'borrow' result is unused, any replacement value will do.
+def SubUIExtendedToSubI:
+ Pattern<(Arith_SubUIExtendedOp:$res $x, $y),
+ [(Arith_SubIOp $x, $y, DefOverflow), (replaceWithValue $x)],
+ [(Constraint<CPred<"$0.getUses().empty()">> $res__1)]>;
+
//===----------------------------------------------------------------------===//
// SubIOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index bcdab9ee8e978..0d62f1f3c07d3 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -491,6 +491,80 @@ void arith::AddUIExtendedOp::getCanonicalizationPatterns(
patterns.add<AddUIExtendedToAddI>(context);
}
+//===----------------------------------------------------------------------===//
+// SubUIExtendedOp
+//===----------------------------------------------------------------------===//
+
+std::optional<SmallVector<int64_t, 4>>
+arith::SubUIExtendedOp::getShapeForUnroll() {
+ if (auto vt = dyn_cast<VectorType>(getType(0)))
+ return llvm::to_vector<4>(vt.getShape());
+ return std::nullopt;
+}
+
+// Returns the borrow bit, assuming `lhs` and `rhs` are operands of an unsigned
+// subtraction whose mathematical result underflows iff `lhs < rhs`.
+static APInt calculateUnsignedBorrow(const APInt &lhs, const APInt &rhs) {
+ return lhs.ult(rhs) ? APInt::getAllOnes(1) : APInt::getZero(1);
+}
+
+LogicalResult
+arith::SubUIExtendedOp::fold(FoldAdaptor adaptor,
+ SmallVectorImpl<OpFoldResult> &results) {
+ Type borrowTy = getBorrow().getType();
+ // subui_extended(x, 0) -> x, false
+ if (matchPattern(getRhs(), m_Zero())) {
+ Builder builder(getContext());
+ auto falseValue = builder.getZeroAttr(borrowTy);
+
+ results.push_back(getLhs());
+ results.push_back(falseValue);
+ return success();
+ }
+
+ // subui_extended(x, x) -> 0, false
+ if (getLhs() == getRhs()) {
+ Builder builder(getContext());
+ auto zeroDiff = builder.getZeroAttr(getDiff().getType());
+ auto falseValue = builder.getZeroAttr(borrowTy);
+ if (!zeroDiff)
+ return failure();
+
+ results.push_back(zeroDiff);
+ results.push_back(falseValue);
+ return success();
+ }
+
+ // subui_extended(constant_a, constant_b) -> constant_
diff , constant_borrow
+ if (Attribute
diff Attr = constFoldBinaryOp<IntegerAttr>(
+ adaptor.getOperands(),
+ [](APInt a, const APInt &b) { return std::move(a) - b; })) {
+ // If any operand is poison, propagate poison to both results.
+ if (matchPattern(
diff Attr, ub::m_Poison())) {
+ results.push_back(
diff Attr);
+ results.push_back(
diff Attr);
+ return success();
+ }
+ Attribute borrowAttr = constFoldBinaryOp<IntegerAttr>(
+ adaptor.getOperands(),
+ getI1SameShape(llvm::cast<TypedAttr>(
diff Attr).getType()),
+ calculateUnsignedBorrow);
+ if (!borrowAttr)
+ return failure();
+
+ results.push_back(
diff Attr);
+ results.push_back(borrowAttr);
+ return success();
+ }
+
+ return failure();
+}
+
+void arith::SubUIExtendedOp::getCanonicalizationPatterns(
+ RewritePatternSet &patterns, MLIRContext *context) {
+ patterns.add<SubUIExtendedToSubI>(context);
+}
+
//===----------------------------------------------------------------------===//
// SubIOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 02626bef856d3..5a52d2e6aa8c1 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -1759,6 +1759,115 @@ func.func @adduiExtendedPoisonRhs() -> (i32, i1) {
return %sum, %overflow : i32, i1
}
+// CHECK-LABEL: @subuiExtendedZeroRhs
+// CHECK-NEXT: %[[false:.+]] = arith.constant false
+// CHECK-NEXT: return %arg0, %[[false]]
+func.func @subuiExtendedZeroRhs(%arg0: i32) -> (i32, i1) {
+ %zero = arith.constant 0 : i32
+ %
diff , %borrow = arith.subui_extended %arg0, %zero: i32, i1
+ return %
diff , %borrow : i32, i1
+}
+
+// CHECK-LABEL: @subuiExtendedZeroRhsSplat
+// CHECK-NEXT: %[[false:.+]] = arith.constant dense<false> : vector<4xi1>
+// CHECK-NEXT: return %arg0, %[[false]]
+func.func @subuiExtendedZeroRhsSplat(%arg0: vector<4xi32>) -> (vector<4xi32>, vector<4xi1>) {
+ %zero = arith.constant dense<0> : vector<4xi32>
+ %
diff , %borrow = arith.subui_extended %arg0, %zero: vector<4xi32>, vector<4xi1>
+ return %
diff , %borrow : vector<4xi32>, vector<4xi1>
+}
+
+// CHECK-LABEL: @subuiExtendedSameOperand
+// CHECK-DAG: %[[zero:.+]] = arith.constant 0 : i32
+// CHECK-DAG: %[[false:.+]] = arith.constant false
+// CHECK-NEXT: return %[[zero]], %[[false]]
+func.func @subuiExtendedSameOperand(%arg0: i32) -> (i32, i1) {
+ %
diff , %borrow = arith.subui_extended %arg0, %arg0: i32, i1
+ return %
diff , %borrow : i32, i1
+}
+
+// CHECK-LABEL: @subuiExtendedSameOperandVector
+// CHECK-DAG: %[[zero:.+]] = arith.constant dense<0> : vector<4xi32>
+// CHECK-DAG: %[[false:.+]] = arith.constant dense<false> : vector<4xi1>
+// CHECK-NEXT: return %[[zero]], %[[false]]
+func.func @subuiExtendedSameOperandVector(%arg0: vector<4xi32>) -> (vector<4xi32>, vector<4xi1>) {
+ %
diff , %borrow = arith.subui_extended %arg0, %arg0: vector<4xi32>, vector<4xi1>
+ return %
diff , %borrow : vector<4xi32>, vector<4xi1>
+}
+
+// CHECK-LABEL: @subuiExtendedUnusedBorrowScalar
+// CHECK-SAME: (%[[LHS:.+]]: i32, %[[RHS:.+]]: i32) -> i32
+// CHECK-NEXT: %[[RES:.+]] = arith.subi %[[LHS]], %[[RHS]] : i32
+// CHECK-NEXT: return %[[RES]] : i32
+func.func @subuiExtendedUnusedBorrowScalar(%arg0: i32, %arg1: i32) -> i32 {
+ %
diff , %borrow = arith.subui_extended %arg0, %arg1: i32, i1
+ return %
diff : i32
+}
+
+// CHECK-LABEL: @subuiExtendedUnusedBorrowVector
+// CHECK-SAME: (%[[LHS:.+]]: vector<3xi32>, %[[RHS:.+]]: vector<3xi32>) -> vector<3xi32>
+// CHECK-NEXT: %[[RES:.+]] = arith.subi %[[LHS]], %[[RHS]] : vector<3xi32>
+// CHECK-NEXT: return %[[RES]] : vector<3xi32>
+func.func @subuiExtendedUnusedBorrowVector(%arg0: vector<3xi32>, %arg1: vector<3xi32>) -> vector<3xi32> {
+ %
diff , %borrow = arith.subui_extended %arg0, %arg1: vector<3xi32>, vector<3xi1>
+ return %
diff : vector<3xi32>
+}
+
+// CHECK-LABEL: @subuiExtendedConstants
+// CHECK-DAG: %[[false:.+]] = arith.constant false
+// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : i32
+// CHECK-NEXT: return %[[c2]], %[[false]]
+func.func @subuiExtendedConstants() -> (i32, i1) {
+ %c5 = arith.constant 5 : i32
+ %c3 = arith.constant 3 : i32
+ %
diff , %borrow = arith.subui_extended %c5, %c3: i32, i1
+ return %
diff , %borrow : i32, i1
+}
+
+// CHECK-LABEL: @subuiExtendedConstantsBorrow
+// CHECK-DAG: %[[true:.+]] = arith.constant true
+// CHECK-DAG: %[[c_2:.+]] = arith.constant -2 : i32
+// CHECK-NEXT: return %[[c_2]], %[[true]]
+func.func @subuiExtendedConstantsBorrow() -> (i32, i1) {
+ %c3 = arith.constant 3 : i32
+ %c5 = arith.constant 5 : i32
+ %
diff , %borrow = arith.subui_extended %c3, %c5: i32, i1
+ return %
diff , %borrow : i32, i1
+}
+
+// CHECK-LABEL: @subuiExtendedConstantsBorrowVector
+// CHECK-DAG: %[[
diff :.+]] = arith.constant dense<[1, 0, -1, 0]> : vector<4xi32>
+// CHECK-DAG: %[[borrow:.+]] = arith.constant dense<[false, false, true, false]> : vector<4xi1>
+// CHECK-NEXT: return %[[
diff ]], %[[borrow]]
+func.func @subuiExtendedConstantsBorrowVector() -> (vector<4xi32>, vector<4xi1>) {
+ %v1 = arith.constant dense<[1, 3, 3, 7]> : vector<4xi32>
+ %v2 = arith.constant dense<[0, 3, 4, 7]> : vector<4xi32>
+ %
diff , %borrow = arith.subui_extended %v1, %v2 : vector<4xi32>, vector<4xi1>
+ return %
diff , %borrow : vector<4xi32>, vector<4xi1>
+}
+
+// CHECK-LABEL: @subuiExtendedPoisonLhs
+// CHECK-NEXT: %[[P0:.+]] = ub.poison : i32
+// CHECK-NEXT: %[[P1:.+]] = ub.poison : i1
+// CHECK-NEXT: return %[[P0]], %[[P1]]
+func.func @subuiExtendedPoisonLhs() -> (i32, i1) {
+ %poison = ub.poison : i32
+ %c5 = arith.constant 5 : i32
+ %
diff , %borrow = arith.subui_extended %poison, %c5 : i32, i1
+ return %
diff , %borrow : i32, i1
+}
+
+// CHECK-LABEL: @subuiExtendedPoisonRhs
+// CHECK-NEXT: %[[P0:.+]] = ub.poison : i32
+// CHECK-NEXT: %[[P1:.+]] = ub.poison : i1
+// CHECK-NEXT: return %[[P0]], %[[P1]]
+func.func @subuiExtendedPoisonRhs() -> (i32, i1) {
+ %c5 = arith.constant 5 : i32
+ %poison = ub.poison : i32
+ %
diff , %borrow = arith.subui_extended %c5, %poison : i32, i1
+ return %
diff , %borrow : i32, i1
+}
+
// CHECK-LABEL: @mulsiExtendedZeroRhs
// CHECK-NEXT: %[[zero:.+]] = arith.constant 0 : i32
// CHECK-NEXT: return %[[zero]], %[[zero]]
diff --git a/mlir/test/Dialect/Arith/invalid.mlir b/mlir/test/Dialect/Arith/invalid.mlir
index 96013a4fadde5..421dac9cfee15 100644
--- a/mlir/test/Dialect/Arith/invalid.mlir
+++ b/mlir/test/Dialect/Arith/invalid.mlir
@@ -168,6 +168,38 @@ func.func @func_with_ops(%a: vector<8xi32>) {
// -----
+func.func @func_with_ops(%a: f32) {
+ // expected-error at +1 {{'arith.subui_extended' op operand #0 must be signless-non-zero-bitwidth-integer-like}}
+ %r:2 = arith.subui_extended %a, %a : f32, i32
+ return
+}
+
+// -----
+
+func.func @func_with_ops(%a: i32) {
+ // expected-error at +1 {{'arith.subui_extended' op result #1 must be bool-like}}
+ %r:2 = arith.subui_extended %a, %a : i32, i32
+ return
+}
+
+// -----
+
+func.func @func_with_ops(%a: vector<8xi32>) {
+ // expected-error at +1 {{'arith.subui_extended' op if an operand is non-scalar, then all results must be non-scalar}}
+ %r:2 = arith.subui_extended %a, %a : vector<8xi32>, i1
+ return
+}
+
+// -----
+
+func.func @func_with_ops(%a: vector<8xi32>) {
+ // expected-error at +1 {{'arith.subui_extended' op all non-scalar operands/results must have the same shape and base type}}
+ %r:2 = arith.subui_extended %a, %a : vector<8xi32>, tensor<8xi1>
+ return
+}
+
+// -----
+
func.func @func_with_ops(i32) {
^bb0(%a : i32):
%sf = arith.addf %a, %a : i32 // expected-error {{'arith.addf' op operand #0 must be floating-point-like}}
diff --git a/mlir/test/Dialect/Arith/ops.mlir b/mlir/test/Dialect/Arith/ops.mlir
index 059e35c384dac..76aef203a64fb 100644
--- a/mlir/test/Dialect/Arith/ops.mlir
+++ b/mlir/test/Dialect/Arith/ops.mlir
@@ -49,6 +49,30 @@ func.func @test_addui_extended_scalable_vector(%arg0 : vector<[8]xi64>, %arg1 :
return %0#0 : vector<[8]xi64>
}
+// CHECK-LABEL: test_subui_extended
+func.func @test_subui_extended(%arg0 : i64, %arg1 : i64) -> i64 {
+ %
diff , %borrow = arith.subui_extended %arg0, %arg1 : i64, i1
+ return %
diff : i64
+}
+
+// CHECK-LABEL: test_subui_extended_tensor
+func.func @test_subui_extended_tensor(%arg0 : tensor<8x8xi64>, %arg1 : tensor<8x8xi64>) -> tensor<8x8xi64> {
+ %
diff , %borrow = arith.subui_extended %arg0, %arg1 : tensor<8x8xi64>, tensor<8x8xi1>
+ return %
diff : tensor<8x8xi64>
+}
+
+// CHECK-LABEL: test_subui_extended_vector
+func.func @test_subui_extended_vector(%arg0 : vector<8xi64>, %arg1 : vector<8xi64>) -> vector<8xi64> {
+ %0:2 = arith.subui_extended %arg0, %arg1 : vector<8xi64>, vector<8xi1>
+ return %0#0 : vector<8xi64>
+}
+
+// CHECK-LABEL: test_subui_extended_scalable_vector
+func.func @test_subui_extended_scalable_vector(%arg0 : vector<[8]xi64>, %arg1 : vector<[8]xi64>) -> vector<[8]xi64> {
+ %0:2 = arith.subui_extended %arg0, %arg1 : vector<[8]xi64>, vector<[8]xi1>
+ return %0#0 : vector<[8]xi64>
+}
+
// CHECK-LABEL: test_subi
func.func @test_subi(%arg0 : i64, %arg1 : i64) -> i64 {
%0 = arith.subi %arg0, %arg1 : i64
More information about the Mlir-commits
mailing list