[Mlir-commits] [mlir] fb826a1 - [mlir][SPIR-V] Add ISubBorrow canonicalization patterns (#198637)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed May 20 04:26:03 PDT 2026
Author: Arseniy Obolenskiy
Date: 2026-05-20T13:25:58+02:00
New Revision: fb826a150f2b33cb391f771a13f3f7e5d924f59b
URL: https://github.com/llvm/llvm-project/commit/fb826a150f2b33cb391f771a13f3f7e5d924f59b
DIFF: https://github.com/llvm/llvm-project/commit/fb826a150f2b33cb391f771a13f3f7e5d924f59b.diff
LOG: [mlir][SPIR-V] Add ISubBorrow canonicalization patterns (#198637)
Mirror the IAddCarry folder, rewrite isubborrow(x, 0) to <x, 0> via
CompositeConstruct, and fold the all-constant case into a single
spirv.Constant struct
Added:
Modified:
mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td
mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td
index e07b1ffd92f07..8a194bc288121 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td
@@ -439,6 +439,8 @@ def SPIRV_ISubBorrowOp : SPIRV_ArithmeticExtendedBinaryOp<"ISubBorrow",
%2 = spirv.ISubBorrow %0, %1 : !spirv.struct<(vector<2xi32>, vector<2xi32>)>
```
}];
+
+ let hasCanonicalizer = 1;
}
// -----
diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
index e4b48ed680d52..acfa6cdf85052 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
@@ -122,68 +122,68 @@ void spirv::AccessChainOp::getCanonicalizationPatterns(
}
//===----------------------------------------------------------------------===//
-// spirv.IAddCarry
+// spirv.IAddCarry / spirv.ISubBorrow
//===----------------------------------------------------------------------===//
-struct IAddCarryFold final : OpRewritePattern<spirv::IAddCarryOp> {
- using Base::Base;
+template <typename Op>
+struct ArithmeticExtendedBinaryFold final : OpRewritePattern<Op> {
+ using OpRewritePattern<Op>::OpRewritePattern;
+
+ static constexpr bool IsSub = std::is_same_v<Op, spirv::ISubBorrowOp>;
- LogicalResult matchAndRewrite(spirv::IAddCarryOp op,
+ LogicalResult matchAndRewrite(Op op,
PatternRewriter &rewriter) const override {
Value lhs = op.getOperand1();
Value rhs = op.getOperand2();
- // iaddcarry (x, 0) = <0, x>
+ // iaddcarry (x, 0) = <0, x>
+ // isubborrow (x, 0) = <x, 0>
if (matchPattern(rhs, m_Zero())) {
- Value constituents[2] = {rhs, lhs};
+ std::array<Value, 2> constituents =
+ IsSub ? std::array{lhs, rhs} : std::array{rhs, lhs};
rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(op, op.getType(),
constituents);
return success();
}
- // According to the SPIR-V spec:
- //
- // Result Type must be from OpTypeStruct. The struct must have two
- // members...
- //
- // Member 0 of the result gets the low-order bits (full component width) of
- // the addition.
- //
- // Member 1 of the result gets the high-order (carry) bit of the result of
- // the addition. That is, it gets the value 1 if the addition overflowed
- // the component width, and 0 otherwise.
Attribute lhsAttr;
Attribute rhsAttr;
if (!matchPattern(lhs, m_Constant(&lhsAttr)) ||
!matchPattern(rhs, m_Constant(&rhsAttr)))
return failure();
- auto adds = constFoldBinaryOp<IntegerAttr>(
+ auto lowBits = constFoldBinaryOp<IntegerAttr>(
{lhsAttr, rhsAttr},
- [](const APInt &a, const APInt &b) { return a + b; });
- if (!adds)
+ [](const APInt &a, const APInt &b) { return IsSub ? a - b : a + b; });
+ if (!lowBits)
return failure();
- auto carrys = constFoldBinaryOp<IntegerAttr>(
- ArrayRef{adds, lhsAttr}, [](const APInt &a, const APInt &b) {
- APInt zero = APInt::getZero(a.getBitWidth());
- return a.ult(b) ? (zero + 1) : zero;
+ auto wrapBit = constFoldBinaryOp<IntegerAttr>(
+ {lhsAttr, rhsAttr}, [](const APInt &a, const APInt &b) {
+ bool wrapped = IsSub ? a.ult(b) : (a + b).ult(a);
+ return APInt(a.getBitWidth(), wrapped ? 1 : 0);
});
-
- if (!carrys)
+ if (!wrapBit)
return failure();
rewriter.replaceOpWithNewOp<spirv::ConstantOp>(
- op, op.getType(), rewriter.getArrayAttr({adds, carrys}));
+ op, op.getType(), rewriter.getArrayAttr({lowBits, wrapBit}));
return success();
}
};
+using IAddCarryFold = ArithmeticExtendedBinaryFold<spirv::IAddCarryOp>;
void spirv::IAddCarryOp::getCanonicalizationPatterns(
RewritePatternSet &patterns, MLIRContext *context) {
patterns.add<IAddCarryFold>(context);
}
+using ISubBorrowFold = ArithmeticExtendedBinaryFold<spirv::ISubBorrowOp>;
+void spirv::ISubBorrowOp::getCanonicalizationPatterns(
+ RewritePatternSet &patterns, MLIRContext *context) {
+ patterns.add<ISubBorrowFold>(context);
+}
+
//===----------------------------------------------------------------------===//
// spirv.[S|U]MulExtended
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir b/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
index 235ca15c08d7b..e49372ae91aed 100644
--- a/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
@@ -379,6 +379,50 @@ func.func @const_fold_vector_iaddcarry() -> !spirv.struct<(vector<3xi32>, vector
// -----
+//===----------------------------------------------------------------------===//
+// spirv.ISubBorrow
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @isubborrow_x_0
+// CHECK-SAME: (%[[ARG:.*]]: i32)
+func.func @isubborrow_x_0(%arg0 : i32) -> !spirv.struct<(i32, i32)> {
+ // CHECK: %[[C0:.*]] = spirv.Constant 0
+ // CHECK: %[[RET:.*]] = spirv.CompositeConstruct %[[ARG]], %[[C0]]
+ %c0 = spirv.Constant 0 : i32
+ %0 = spirv.ISubBorrow %arg0, %c0 : !spirv.struct<(i32, i32)>
+
+ // CHECK: return %[[RET]]
+ return %0 : !spirv.struct<(i32, i32)>
+}
+
+// CHECK-LABEL: @const_fold_scalar_isubborrow
+func.func @const_fold_scalar_isubborrow() -> (!spirv.struct<(i32, i32)>, !spirv.struct<(i32, i32)>) {
+ %c5 = spirv.Constant 5 : i32
+ %c8 = spirv.Constant 8 : i32
+
+ // CHECK-DAG: %[[CST_CN3_C1:.*]] = spirv.Constant [-3 : i32, 1 : i32] : !spirv.struct<(i32, i32)>
+ // CHECK-DAG: %[[CST_C3_C0:.*]] = spirv.Constant [3 : i32, 0 : i32] : !spirv.struct<(i32, i32)>
+ %0 = spirv.ISubBorrow %c5, %c8 : !spirv.struct<(i32, i32)>
+ %1 = spirv.ISubBorrow %c8, %c5 : !spirv.struct<(i32, i32)>
+
+ // CHECK: return %[[CST_CN3_C1]], %[[CST_C3_C0]]
+ return %0, %1 : !spirv.struct<(i32, i32)>, !spirv.struct<(i32, i32)>
+}
+
+// CHECK-LABEL: @const_fold_vector_isubborrow
+func.func @const_fold_vector_isubborrow() -> !spirv.struct<(vector<3xi32>, vector<3xi32>)> {
+ %v0 = spirv.Constant dense<[5, 8, -1]> : vector<3xi32>
+ %v1 = spirv.Constant dense<[8, 5, 1]> : vector<3xi32>
+
+ // CHECK: %[[CST:.*]] = spirv.Constant [dense<[-3, 3, -2]> : vector<3xi32>, dense<[1, 0, 0]> : vector<3xi32>] : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
+ %0 = spirv.ISubBorrow %v0, %v1 : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
+
+ // CHECK: return %[[CST]]
+ return %0 : !spirv.struct<(vector<3xi32>, vector<3xi32>)>
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// spirv.IMul
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list