[Mlir-commits] [mlir] 5368b81 - [MLIR][Arith] Add canonicalization rules for int-to-float of integer extension (#185386)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Mar 11 05:50:04 PDT 2026
Author: Adam Paszke
Date: 2026-03-11T14:49:59+02:00
New Revision: 5368b8163d50b576490c076df9c9b8e62d7a6ec6
URL: https://github.com/llvm/llvm-project/commit/5368b8163d50b576490c076df9c9b8e62d7a6ec6
DIFF: https://github.com/llvm/llvm-project/commit/5368b8163d50b576490c076df9c9b8e62d7a6ec6.diff
LOG: [MLIR][Arith] Add canonicalization rules for int-to-float of integer extension (#185386)
Three patterns are valid but were missing:
1. `sitofp(extsi(x)) → sitofp(x)`: extsi preserves the sign and value,
so it represents the same signed integer as x.
2. `uitofp(extui(x)) → uitofp(x)`: same reasoning as above, but for
unsigned extension.
3. `sitofp(extui(x)) → uitofp(x)` extui zero-extends, so the extended
value is always non-negative. For non-negative integers, sitofp and
uitofp produce the same result, meaning we could replace the left
expression by `uitofp(extui(x))`. At this point rule 2. above can be
used to simplify further to `uitofp(x)`.
All three rewrites have been verified with Alive2.
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/emulate-wide-int-canonicalization.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
index 27d3d0d9b6e71..4b830c05bf585 100644
--- a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
+++ b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
@@ -1530,6 +1530,7 @@ def Arith_UIToFPOp :
$in (`nneg` $nonNeg^)? attr-dict `:` type($in) `to` type($out)
}];
let hasFolder = 1;
+ let hasCanonicalizer = 1;
}
//===----------------------------------------------------------------------===//
@@ -1545,6 +1546,7 @@ def Arith_SIToFPOp : Arith_IToFCastOp<"sitofp"> {
elementwise.
}];
let hasFolder = 1;
+ let hasCanonicalizer = 1;
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
index fb9c16db91431..e22fc1d478e4f 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
+++ b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
@@ -421,6 +421,29 @@ def TruncFUIToFPToUIToFP :
(Arith_UIToFPOp $x, $nneg),
[(Constraint<CPred<"$0 == nullptr">, "default rounding mode"> $rmf)]>;
+//===----------------------------------------------------------------------===//
+// SIToFPOp
+//===----------------------------------------------------------------------===//
+
+// sitofp(extsi(x)) -> sitofp(x)
+def SIToFPOfExtSI :
+ Pat<(Arith_SIToFPOp (Arith_ExtSIOp $x)),
+ (Arith_SIToFPOp $x)>;
+
+// sitofp(extui(x)) -> uitofp(x)
+def SIToFPOfExtUI :
+ Pat<(Arith_SIToFPOp (Arith_ExtUIOp $x, $nneg)),
+ (Arith_UIToFPOp $x, $nneg)>;
+
+//===----------------------------------------------------------------------===//
+// UIToFPOp
+//===----------------------------------------------------------------------===//
+
+// uitofp(extui(x)) -> uitofp(x)
+def UIToFPOfExtUI :
+ Pat<(Arith_UIToFPOp (Arith_ExtUIOp $x, $nneg1), $nneg2),
+ (Arith_UIToFPOp $x, $nneg1)>;
+
//===----------------------------------------------------------------------===//
// MulFOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 73fa0df80bcba..6f368604df65a 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -1756,6 +1756,11 @@ OpFoldResult arith::UIToFPOp::fold(FoldAdaptor adaptor) {
});
}
+void arith::UIToFPOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+ MLIRContext *context) {
+ patterns.add<UIToFPOfExtUI>(context);
+}
+
//===----------------------------------------------------------------------===//
// SIToFPOp
//===----------------------------------------------------------------------===//
@@ -1778,6 +1783,11 @@ OpFoldResult arith::SIToFPOp::fold(FoldAdaptor adaptor) {
});
}
+void arith::SIToFPOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+ MLIRContext *context) {
+ patterns.add<SIToFPOfExtSI, SIToFPOfExtUI>(context);
+}
+
//===----------------------------------------------------------------------===//
// FPToUIOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 035c10e78bf9b..26b04e4209a43 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -875,7 +875,42 @@ func.func @truncUitofp_nneg(%arg0: i32) -> f32 {
return %trunc : f32
}
-// TODO: We should also add a test for not folding arith.extf on information loss.
+// CHECK-LABEL: @sitofpExtsi
+// CHECK: %[[SITOFP:.*]] = arith.sitofp %[[ARG0:.*]] : i8 to bf16
+// CHECK: return %[[SITOFP]]
+func.func @sitofpExtsi(%arg0: i8) -> bf16 {
+ %extsi = arith.extsi %arg0 : i8 to i32
+ %sitofp = arith.sitofp %extsi : i32 to bf16
+ return %sitofp : bf16
+}
+
+// CHECK-LABEL: @sitofpExtui
+// CHECK: %[[UITOFP:.*]] = arith.uitofp %[[ARG0:.*]] : i4 to bf16
+// CHECK-NOT: sitofp
+// CHECK: return %[[UITOFP]]
+func.func @sitofpExtui(%arg0: i4) -> bf16 {
+ %extui = arith.extui %arg0 : i4 to i8
+ %sitofp = arith.sitofp %extui : i8 to bf16
+ return %sitofp : bf16
+}
+
+// CHECK-LABEL: @uitofpExtui
+// CHECK: %[[UITOFP:.*]] = arith.uitofp %[[ARG0:.*]] : i8 to bf16
+// CHECK: return %[[UITOFP]]
+func.func @uitofpExtui(%arg0: i8) -> bf16 {
+ %extui = arith.extui %arg0 : i8 to i32
+ %uitofp = arith.uitofp %extui : i32 to bf16
+ return %uitofp : bf16
+}
+
+// CHECK-LABEL: @sitofpExtui_nneg
+// CHECK: %[[UITOFP:.*]] = arith.uitofp %[[ARG0:.*]] nneg : i4 to bf16
+// CHECK: return %[[UITOFP]]
+func.func @sitofpExtui_nneg(%arg0: i4) -> bf16 {
+ %extui = arith.extui %arg0 nneg : i4 to i8
+ %sitofp = arith.sitofp %extui : i8 to bf16
+ return %sitofp : bf16
+}
// This may happen when extending f8E5M2FNUZ to f16.
// CHECK-LABEL: @truncConstant
diff --git a/mlir/test/Dialect/Arith/emulate-wide-int-canonicalization.mlir b/mlir/test/Dialect/Arith/emulate-wide-int-canonicalization.mlir
index 0c95ab8284afa..75fd9eda08782 100644
--- a/mlir/test/Dialect/Arith/emulate-wide-int-canonicalization.mlir
+++ b/mlir/test/Dialect/Arith/emulate-wide-int-canonicalization.mlir
@@ -1,14 +1,17 @@
// RUN: mlir-opt --arith-emulate-wide-int="widest-int-supported=32" --canonicalize %s | FileCheck %s
-// Check that we can fold away the 'hi' part calculation when it is know to be zero.
+// Check that we can fold away the 'hi' part calculation when it is known to be zero.
//
// CHECK-LABEL: func @uitofp_i16_ext_f64
// CHECK-SAME: ([[ARG:%.+]]: i16) -> f64
-// CHECK-NEXT: [[EXT:%.+]] = arith.extui [[ARG]] : i16 to i32
-// CHECK-NEXT: [[FP:%.+]] = arith.uitofp [[EXT]] : i32 to f64
+// CHECK: [[EXT:%.+]] = arith.extui [[ARG]] : i16 to i32
+// CHECK: [[ORI:%.+]] = arith.ori [[EXT]], {{.*}} : i32
+// CHECK: [[FP:%.+]] = arith.uitofp [[ORI]] : i32 to f64
// CHECK-NEXT: return [[FP]] : f64
func.func @uitofp_i16_ext_f64(%a : i16) -> f64 {
%ext = arith.extui %a : i16 to i64
- %r = arith.uitofp %ext : i64 to f64
+ %c = arith.constant 1 : i64
+ %or = arith.ori %ext, %c : i64
+ %r = arith.uitofp %or : i64 to f64
return %r : f64
}
More information about the Mlir-commits
mailing list