[Mlir-commits] [mlir] [mlir][arith] Add `exact` to `index_cast{, ui}` (PR #183395)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Feb 26 07:02:42 PST 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-arith
Author: Erick Ochoa Lopez (amd-eochoalo)
<details>
<summary>Changes</summary>
The `exact` flag with the following semantics
> If the `exact` attribute is present, it is assumed that the index type width
> is such that the conversion does not lose information. When this assumption
> is violated, the result is poison.
can be added to index_cast and index_castui operations. This unlocks
the following lowerings:
* index_cast (signed) exact -> trunc nsw
* index_castui (unsigned) exact -> trunc nuw
* index_castui nneg exact -> trunc nuw nsw
Changes:
* Adds ArithExactFlagInterface.
* Updates Arith_IntBinaryOpWithExactFlag to use ArithExactFlagInterface
* Update IndexCastOp and IndexCastUIOp to declare `ArithExactFlagInterface`
* Update canonicalization patterns
* Update roundtrip, lowering, and canonicalization tests.
---
Full diff: https://github.com/llvm/llvm-project/pull/183395.diff
7 Files Affected:
- (modified) mlir/include/mlir/Dialect/Arith/IR/ArithOps.td (+31-5)
- (modified) mlir/include/mlir/Dialect/Arith/IR/ArithOpsInterfaces.td (+47)
- (modified) mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp (+28-8)
- (modified) mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td (+11-8)
- (modified) mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir (+52)
- (modified) mlir/test/Dialect/Arith/canonicalize.mlir (+61)
- (modified) mlir/test/Dialect/Arith/ops.mlir (+28)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
index 3d8517c56e784..65cbed612f6b3 100644
--- a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
+++ b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
@@ -160,7 +160,8 @@ class Arith_IntBinaryOpWithOverflowFlags<string mnemonic, list<Trait> traits = [
class Arith_IntBinaryOpWithExactFlag<string mnemonic, list<Trait> traits = []> :
Arith_BinaryOp<mnemonic, traits #
- [DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>]>,
+ [DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>,
+ DeclareOpInterfaceMethods<ArithExactFlagInterface>]>,
Arguments<(ins SignlessIntegerOrIndexLike:$lhs,
SignlessIntegerOrIndexLike:$rhs,
UnitAttr:$isExact)>,
@@ -1579,15 +1580,32 @@ def IndexCastTypeConstraint : TypeConstraint<Or<[
def Arith_IndexCastOp
: Arith_CastOp<"index_cast", IndexCastTypeConstraint, IndexCastTypeConstraint,
- [DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>]> {
+ [DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>,
+ DeclareOpInterfaceMethods<ArithExactFlagInterface>]> {
let summary = "cast between index and integer types";
let description = [{
Casts between scalar or vector integers and corresponding 'index' scalar or
vectors. Index is an integer of platform-specific bit width. If casting to
a wider integer, the value is sign-extended. If casting to a narrower
integer, the value is truncated.
+
+ If the `exact` attribute is present, it is assumed that the index type width
+ is such that the conversion does not lose information. When this assumption
+ is violated, the result is poison.
+
+ Example:
+
+ ```mlir
+ %0 = arith.index_cast %a : index to i64
+ %1 = arith.index_cast %a exact : index to i64
+ ```
}];
+ let arguments = (ins IndexCastTypeConstraint:$in, UnitAttr:$isExact);
+ let results = (outs IndexCastTypeConstraint:$out);
+ let assemblyFormat = [{
+ $in (`exact` $isExact^)? attr-dict `:` type($in) `to` type($out)
+ }];
let hasFolder = 1;
let hasCanonicalizer = 1;
}
@@ -1599,7 +1617,8 @@ def Arith_IndexCastOp
def Arith_IndexCastUIOp
: Arith_CastOp<"index_castui", IndexCastTypeConstraint, IndexCastTypeConstraint,
[DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>,
- DeclareOpInterfaceMethods<ArithNonNegFlagInterface>]> {
+ DeclareOpInterfaceMethods<ArithNonNegFlagInterface>,
+ DeclareOpInterfaceMethods<ArithExactFlagInterface>]> {
let summary = "unsigned cast between index and integer types";
let description = [{
Casts between scalar or vector integers and corresponding 'index' scalar or
@@ -1612,19 +1631,26 @@ def Arith_IndexCastUIOp
is equivalent to sign extension. When this assumption is violated,
the result is poison.
+ If the `exact` attribute is present, it is assumed that the index type width
+ is such that the conversion does not lose information. When this assumption
+ is violated, the result is poison.
+
Example:
```mlir
%0 = arith.index_castui %a : i32 to index
%1 = arith.index_castui %a nneg : i32 to index
%2 = arith.index_castui %b nneg : index to i64
+ %3 = arith.index_castui %a nneg exact : i64 to index
```
}];
- let arguments = (ins IndexCastTypeConstraint:$in, UnitAttr:$nonNeg);
+ let arguments = (ins IndexCastTypeConstraint:$in, UnitAttr:$nonNeg,
+ UnitAttr:$isExact);
let results = (outs IndexCastTypeConstraint:$out);
let assemblyFormat = [{
- $in (`nneg` $nonNeg^)? attr-dict `:` type($in) `to` type($out)
+ $in oilist(`exact` $isExact | `nneg` $nonNeg) attr-dict
+ `:` type($in) `to` type($out)
}];
let hasFolder = 1;
let hasCanonicalizer = 1;
diff --git a/mlir/include/mlir/Dialect/Arith/IR/ArithOpsInterfaces.td b/mlir/include/mlir/Dialect/Arith/IR/ArithOpsInterfaces.td
index d1b8e250cdb59..e8287ac2d6bcc 100644
--- a/mlir/include/mlir/Dialect/Arith/IR/ArithOpsInterfaces.td
+++ b/mlir/include/mlir/Dialect/Arith/IR/ArithOpsInterfaces.td
@@ -153,6 +153,53 @@ def ArithNonNegFlagInterface : OpInterface<"ArithNonNegFlagInterface"> {
];
}
+def ArithExactFlagInterface : OpInterface<"ArithExactFlagInterface"> {
+ let description = [{
+ Access to op exact flag.
+ }];
+
+ let cppNamespace = "::mlir::arith";
+
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/ "Returns whether the operation has the exact flag set",
+ /*returnType=*/ "bool",
+ /*methodName=*/ "getExact",
+ /*args=*/ (ins),
+ /*methodBody=*/ [{}],
+ /*defaultImpl=*/ [{
+ auto op = cast<ConcreteOp>(this->getOperation());
+ return op.getIsExactAttr() != nullptr;
+ }]
+ >,
+ InterfaceMethod<
+ /*desc=*/ "Set the exact flag for the operation",
+ /*returnType=*/ "void",
+ /*methodName=*/ "setExact",
+ /*args=*/ (ins "bool":$isExact),
+ /*methodBody=*/ [{}],
+ /*defaultImpl=*/ [{
+ auto op = cast<ConcreteOp>(this->getOperation());
+ if (isExact)
+ op.setIsExactAttr(UnitAttr::get(op->getContext()));
+ else
+ op.removeIsExactAttr();
+ }]
+ >,
+ StaticInterfaceMethod<
+ /*desc=*/ [{Returns the name of the exact flag attribute for
+ the operation}],
+ /*returnType=*/ "StringRef",
+ /*methodName=*/ "getExactFlagAttrName",
+ /*args=*/ (ins),
+ /*methodBody=*/ [{}],
+ /*defaultImpl=*/ [{
+ return "isExact";
+ }]
+ >
+ ];
+}
+
def ArithRoundingModeInterface : OpInterface<"ArithRoundingModeInterface"> {
let description = [{
Access to op rounding mode.
diff --git a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp b/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
index e7f561e8a4d67..e0e1be35e4e1d 100644
--- a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
+++ b/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
@@ -311,13 +311,32 @@ LogicalResult IndexCastOpLowering<OpTy, ExtCastTy>::matchAndRewrite(
if constexpr (std::is_same_v<ExtCastTy, LLVM::ZExtOp>)
isNonNeg = op.getNonNeg();
+ bool isExact = op.getExact();
+
+ // Map exact to the appropriate overflow flag(s) for truncation:
+ // index_cast (signed) exact -> trunc nsw
+ // index_castui (unsigned) exact -> trunc nuw
+ // index_castui nneg exact -> trunc nuw nsw
+ LLVM::IntegerOverflowFlags truncOverflow = LLVM::IntegerOverflowFlags::none;
+ if (isExact) {
+ if constexpr (std::is_same_v<ExtCastTy, LLVM::SExtOp>) {
+ truncOverflow = LLVM::IntegerOverflowFlags::nsw;
+ } else {
+ truncOverflow = LLVM::IntegerOverflowFlags::nuw;
+ if (isNonNeg)
+ truncOverflow |= LLVM::IntegerOverflowFlags::nsw;
+ }
+ }
+
// Handle the scalar and 1D vector cases.
Type operandType = adaptor.getIn().getType();
if (!isa<LLVM::LLVMArrayType>(operandType)) {
Type targetType = this->typeConverter->convertType(resultType);
if (targetBits < sourceBits) {
- rewriter.replaceOpWithNewOp<LLVM::TruncOp>(op, targetType,
- adaptor.getIn());
+ auto truncOp = rewriter.replaceOpWithNewOp<LLVM::TruncOp>(
+ op, targetType, adaptor.getIn());
+ if (isExact)
+ truncOp.setOverflowFlags(truncOverflow);
} else {
auto extOp = rewriter.replaceOpWithNewOp<ExtCastTy>(op, targetType,
adaptor.getIn());
@@ -335,15 +354,16 @@ LogicalResult IndexCastOpLowering<OpTy, ExtCastTy>::matchAndRewrite(
[&](Type llvm1DVectorTy, ValueRange operands) -> Value {
typename OpTy::Adaptor adaptor(operands);
if (targetBits < sourceBits) {
- return LLVM::TruncOp::create(rewriter, op.getLoc(), llvm1DVectorTy,
- adaptor.getIn());
+ auto truncOp = LLVM::TruncOp::create(rewriter, op.getLoc(),
+ llvm1DVectorTy, adaptor.getIn());
+ if (isExact)
+ truncOp.setOverflowFlags(truncOverflow);
+ return truncOp;
}
auto extOp = ExtCastTy::create(rewriter, op.getLoc(), llvm1DVectorTy,
adaptor.getIn());
- if constexpr (std::is_same_v<ExtCastTy, LLVM::ZExtOp>) {
- if (isNonNeg)
- extOp.setNonNeg(true);
- }
+ if constexpr (std::is_same_v<ExtCastTy, LLVM::ZExtOp>)
+ extOp.setNonNeg(isNonNeg);
return extOp;
},
rewriter);
diff --git a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
index fb9c16db91431..4a311b13277c2 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
+++ b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
@@ -290,29 +290,32 @@ def SelectI1ToNot :
// index_cast(index_cast(x)) -> x, if dstType == srcType.
def IndexCastOfIndexCast :
- Pat<(Arith_IndexCastOp:$res (Arith_IndexCastOp $x)),
+ Pat<(Arith_IndexCastOp:$res (Arith_IndexCastOp $x, $exact1), $exact2),
(replaceWithValue $x),
[(Constraint<CPred<"$0.getType() == $1.getType()">> $res, $x)]>;
// index_cast(extsi(x)) -> index_cast(x)
def IndexCastOfExtSI :
- Pat<(Arith_IndexCastOp (Arith_ExtSIOp $x)), (Arith_IndexCastOp $x)>;
+ Pat<(Arith_IndexCastOp (Arith_ExtSIOp $x), $exact),
+ (Arith_IndexCastOp $x, $exact)>;
//===----------------------------------------------------------------------===//
// IndexCastUIOp
//===----------------------------------------------------------------------===//
-// index_castui(index_castui(x)) -> x, if dstType == srcType.
+// index_castui(index_castui(x)) -> x, if dstType == srcType and at least one
+// exact flag is set (guaranteeing no information loss in either cast).
def IndexCastUIOfIndexCastUI :
- Pat<(Arith_IndexCastUIOp:$res (Arith_IndexCastUIOp $x, $nneg1), $nneg2),
+ Pat<(Arith_IndexCastUIOp:$res
+ (Arith_IndexCastUIOp $x, $nneg1, $exact1), $nneg2, $exact2),
(replaceWithValue $x),
- [(Constraint<CPred<"$0.getType() == $1.getType()">> $res, $x)]>;
+ [(Constraint<CPred<"$0.getType() == $1.getType()">> $res, $x),
+ (Constraint<CPred<"$0 || $1">> $exact1, $exact2)]>;
// index_castui(extui(x)) -> index_castui(x)
def IndexCastUIOfExtUI :
- Pat<(Arith_IndexCastUIOp (Arith_ExtUIOp $x, $nneg1), $nneg2),
- (Arith_IndexCastUIOp $x, $nneg1)>;
-
+ Pat<(Arith_IndexCastUIOp (Arith_ExtUIOp $x, $nneg1), $nneg2, $exact),
+ (Arith_IndexCastUIOp $x, $nneg1, $exact)>;
//===----------------------------------------------------------------------===//
// BitcastOp
diff --git a/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir b/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir
index 47069906fa110..2845df23293d5 100644
--- a/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir
+++ b/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir
@@ -160,6 +160,58 @@ func.func @index_castui_nneg_not_set(%arg0: i1) {
// -----
+// index_cast exact on truncation lowers to trunc nsw (signed semantics).
+// CHECK-LABEL: @index_cast_exact_trunc
+func.func @index_cast_exact_trunc(%arg0: index) {
+// CHECK: llvm.trunc %{{.*}} overflow<nsw> : i{{.*}} to i1
+ %0 = arith.index_cast %arg0 exact : index to i1
+ return
+}
+
+// -----
+
+// index_cast exact on widening: exact is vacuously true, sext has no flag.
+// CHECK-LABEL: @index_cast_exact_ext
+func.func @index_cast_exact_ext(%arg0: i1) {
+// CHECK: llvm.sext %{{.*}} : i1 to i{{.*}}
+// CHECK-NOT: nsw
+ %0 = arith.index_cast %arg0 exact : i1 to index
+ return
+}
+
+// -----
+
+// index_castui exact on truncation lowers to trunc nuw (unsigned semantics).
+// CHECK-LABEL: @index_castui_exact_trunc
+func.func @index_castui_exact_trunc(%arg0: index) {
+// CHECK: llvm.trunc %{{.*}} overflow<nuw> : i{{.*}} to i1
+ %0 = arith.index_castui %arg0 exact : index to i1
+ return
+}
+
+// -----
+
+// index_castui nneg exact on truncation lowers to trunc nuw nsw.
+// CHECK-LABEL: @index_castui_nneg_exact_trunc
+func.func @index_castui_nneg_exact_trunc(%arg0: index) {
+// CHECK: llvm.trunc %{{.*}} overflow<nsw, nuw> : i{{.*}} to i1
+ %0 = arith.index_castui %arg0 nneg exact : index to i1
+ return
+}
+
+// -----
+
+// index_castui exact on widening: exact is vacuously true, zext has no flag.
+// CHECK-LABEL: @index_castui_exact_ext
+func.func @index_castui_exact_ext(%arg0: i1) {
+// CHECK: llvm.zext %{{.*}} : i1 to i{{.*}}
+// CHECK-NOT: nuw
+ %0 = arith.index_castui %arg0 exact : i1 to index
+ return
+}
+
+// -----
+
// Checking conversion of signed integer types to floating point.
// CHECK-LABEL: @sitofp
func.func @sitofp(%arg0 : i32, %arg1 : i64) {
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 4dc29897cec26..d43af670d2793 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -588,6 +588,15 @@ func.func @indexCastOfSignExtend(%arg0: i8) -> index {
return %idx : index
}
+// CHECK-LABEL: @indexCastOfSignExtend_exact
+// CHECK: %[[res:.+]] = arith.index_cast %arg0 exact : i8 to index
+// CHECK: return %[[res]]
+func.func @indexCastOfSignExtend_exact(%arg0: i8) -> index {
+ %ext = arith.extsi %arg0 : i8 to i16
+ %idx = arith.index_cast %ext exact : i16 to index
+ return %idx : index
+}
+
// CHECK-LABEL: @indexCastUIOfUnsignedExtend
// CHECK: %[[res:.+]] = arith.index_castui %arg0 : i8 to index
// CHECK: return %[[res]]
@@ -616,6 +625,58 @@ func.func @indexCastUIOfUnsignedExtend_nneg_on_castui(%arg0: i8) -> index {
return %idx : index
}
+// CHECK-LABEL: @indexCastUIOfUnsignedExtend_exact
+// CHECK: %[[res:.+]] = arith.index_castui %arg0 exact : i8 to index
+// CHECK: return %[[res]]
+func.func @indexCastUIOfUnsignedExtend_exact(%arg0: i8) -> index {
+ %ext = arith.extui %arg0 : i8 to i16
+ %idx = arith.index_castui %ext exact : i16 to index
+ return %idx : index
+}
+
+// CHECK-LABEL: @indexCastUIOfUnsignedExtend_nneg_exact
+// CHECK: %[[res:.+]] = arith.index_castui %arg0 exact nneg : i8 to index
+// CHECK: return %[[res]]
+func.func @indexCastUIOfUnsignedExtend_nneg_exact(%arg0: i8) -> index {
+ %ext = arith.extui %arg0 nneg : i8 to i16
+ %idx = arith.index_castui %ext exact : i16 to index
+ return %idx : index
+}
+
+// index_castui(index_castui(x)) -> x only when at least one exact is set.
+// CHECK-LABEL: @indexCastUIOfIndexCastUI_no_exact
+// CHECK: arith.index_castui
+// CHECK: arith.index_castui
+func.func @indexCastUIOfIndexCastUI_no_exact(%arg0: i32) -> i32 {
+ %idx = arith.index_castui %arg0 : i32 to index
+ %res = arith.index_castui %idx : index to i32
+ return %res : i32
+}
+
+// CHECK-LABEL: @indexCastUIOfIndexCastUI_exact_inner
+// CHECK: return %arg0 : i32
+func.func @indexCastUIOfIndexCastUI_exact_inner(%arg0: i32) -> i32 {
+ %idx = arith.index_castui %arg0 exact : i32 to index
+ %res = arith.index_castui %idx : index to i32
+ return %res : i32
+}
+
+// CHECK-LABEL: @indexCastUIOfIndexCastUI_exact_outer
+// CHECK: return %arg0 : i32
+func.func @indexCastUIOfIndexCastUI_exact_outer(%arg0: i32) -> i32 {
+ %idx = arith.index_castui %arg0 : i32 to index
+ %res = arith.index_castui %idx exact : index to i32
+ return %res : i32
+}
+
+// CHECK-LABEL: @indexCastUIOfIndexCastUI_exact_both
+// CHECK: return %arg0 : i32
+func.func @indexCastUIOfIndexCastUI_exact_both(%arg0: i32) -> i32 {
+ %idx = arith.index_castui %arg0 exact : i32 to index
+ %res = arith.index_castui %idx exact : index to i32
+ return %res : i32
+}
+
// CHECK-LABEL: @indexCastFold
// CHECK: %[[res:.*]] = arith.constant -2 : index
// CHECK: return %[[res]]
diff --git a/mlir/test/Dialect/Arith/ops.mlir b/mlir/test/Dialect/Arith/ops.mlir
index 9765db69d6dd5..a9eabe97ebfcd 100644
--- a/mlir/test/Dialect/Arith/ops.mlir
+++ b/mlir/test/Dialect/Arith/ops.mlir
@@ -909,6 +909,20 @@ func.func @test_index_cast_scalable_vector1(%arg0 : vector<[8]xindex>) -> vector
return %0 : vector<[8]xi64>
}
+// CHECK-LABEL: test_index_cast_exact
+// CHECK: arith.index_cast %{{.*}} exact : i32 to index
+func.func @test_index_cast_exact(%arg0 : i32) -> index {
+ %0 = arith.index_cast %arg0 exact : i32 to index
+ return %0 : index
+}
+
+// CHECK-LABEL: test_index_cast_exact_vector
+// CHECK: arith.index_cast %{{.*}} exact : vector<8xi32> to vector<8xindex>
+func.func @test_index_cast_exact_vector(%arg0 : vector<8xi32>) -> vector<8xindex> {
+ %0 = arith.index_cast %arg0 exact : vector<8xi32> to vector<8xindex>
+ return %0 : vector<8xindex>
+}
+
// CHECK-LABEL: test_index_castui0
func.func @test_index_castui0(%arg0 : i32) -> index {
%0 = arith.index_castui %arg0 : i32 to index
@@ -971,6 +985,20 @@ func.func @test_index_castui_nneg_vector(%arg0 : vector<8xi32>) -> vector<8xinde
return %0 : vector<8xindex>
}
+// CHECK-LABEL: test_index_castui_exact
+// CHECK: arith.index_castui %{{.*}} exact : i32 to index
+func.func @test_index_castui_exact(%arg0 : i32) -> index {
+ %0 = arith.index_castui %arg0 exact : i32 to index
+ return %0 : index
+}
+
+// CHECK-LABEL: test_index_castui_nneg_exact
+// CHECK: arith.index_castui %{{.*}} exact nneg : i32 to index
+func.func @test_index_castui_nneg_exact(%arg0 : i32) -> index {
+ %0 = arith.index_castui %arg0 nneg exact : i32 to index
+ return %0 : index
+}
+
// CHECK-LABEL: test_bitcast0
func.func @test_bitcast0(%arg0 : i64) -> f64 {
%0 = arith.bitcast %arg0 : i64 to f64
``````````
</details>
https://github.com/llvm/llvm-project/pull/183395
More information about the Mlir-commits
mailing list