[Mlir-commits] [mlir] [MLIR][Arith] Add canonicalization rules for int-to-float of integer extension (PR #185386)
Adam Paszke
llvmlistbot at llvm.org
Tue Mar 10 04:13:17 PDT 2026
https://github.com/apaszke updated https://github.com/llvm/llvm-project/pull/185386
>From 7e22ea6764320cff996c75fa6ef9dee5501ac440 Mon Sep 17 00:00:00 2001
From: Adam Paszke <apaszke at google.com>
Date: Mon, 9 Mar 2026 10:13:50 +0000
Subject: [PATCH] [MLIR][Arith] Add canonicalization rules for int-to-float of
integer extension
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
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)`.
---
.../include/mlir/Dialect/Arith/IR/ArithOps.td | 2 +
.../Dialect/Arith/IR/ArithCanonicalization.td | 23 ++++++++++++
mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 10 +++++
mlir/test/Dialect/Arith/canonicalize.mlir | 37 ++++++++++++++++++-
.../emulate-wide-int-canonicalization.mlir | 11 ++++--
5 files changed, 78 insertions(+), 5 deletions(-)
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