[Mlir-commits] [mlir] [MLIR][Math] Add floating point value folders (PR #127947)

[William Moses]

https://github.com/wsmoses updated https://github.com/llvm/llvm-project/pull/127947

>From d5262cb0060ed636286d717aff0948acb533aa43 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh at wsmoses.com>
Date: Wed, 19 Feb 2025 20:45:20 -0600
Subject: [PATCH] [MLIR][Math] Add floating point value folders

---
 mlir/include/mlir/Dialect/Math/IR/MathOps.td |  4 ++
 mlir/lib/Dialect/Math/IR/MathOps.cpp         | 64 +++++++++++++++++
 mlir/test/Dialect/Math/canonicalize.mlir     | 72 ++++++++++++++++++++
 3 files changed, 140 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Math/IR/MathOps.td b/mlir/include/mlir/Dialect/Math/IR/MathOps.td
index 16ce4e2366c76..56370388dea87 100644
--- a/mlir/include/mlir/Dialect/Math/IR/MathOps.td
+++ b/mlir/include/mlir/Dialect/Math/IR/MathOps.td
@@ -736,6 +736,7 @@ def Math_IsFiniteOp : Math_FloatClassificationOp<"isfinite"> {
     %f = math.isfinite %a : f32
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
@@ -754,6 +755,7 @@ def Math_IsInfOp : Math_FloatClassificationOp<"isinf"> {
     %f = math.isinf %a : f32
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
@@ -772,6 +774,7 @@ def Math_IsNaNOp : Math_FloatClassificationOp<"isnan"> {
     %f = math.isnan %a : f32
     ```
   }];
+  let hasFolder = 1;
 }
 
 
@@ -791,6 +794,7 @@ def Math_IsNormalOp : Math_FloatClassificationOp<"isnormal"> {
     %f = math.isnormal %a : f32
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Math/IR/MathOps.cpp b/mlir/lib/Dialect/Math/IR/MathOps.cpp
index 9c4d88e2191ce..26441a9d78658 100644
--- a/mlir/lib/Dialect/Math/IR/MathOps.cpp
+++ b/mlir/lib/Dialect/Math/IR/MathOps.cpp
@@ -579,6 +579,70 @@ OpFoldResult math::ExpM1Op::fold(FoldAdaptor adaptor) {
       });
 }
 
+//===----------------------------------------------------------------------===//
+// IsFiniteOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::IsFiniteOp::fold(FoldAdaptor adaptor) {
+  if (auto val = dyn_cast_or_null<FloatAttr>(adaptor.getOperand())) {
+    return BoolAttr::get(val.getContext(), val.getValue().isFinite());
+  }
+  if (auto splat = dyn_cast_or_null<SplatElementsAttr>(adaptor.getOperand())) {
+    return DenseElementsAttr::get(
+        cast<ShapedType>(getType()),
+        APInt(1, splat.getSplatValue<APFloat>().isFinite()));
+  }
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
+// IsInfOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::IsInfOp::fold(FoldAdaptor adaptor) {
+  if (auto val = dyn_cast_or_null<FloatAttr>(adaptor.getOperand())) {
+    return BoolAttr::get(val.getContext(), val.getValue().isInfinity());
+  }
+  if (auto splat = dyn_cast_or_null<SplatElementsAttr>(adaptor.getOperand())) {
+    return DenseElementsAttr::get(
+        cast<ShapedType>(getType()),
+        APInt(1, splat.getSplatValue<APFloat>().isInfinity()));
+  }
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
+// IsNaNOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::IsNaNOp::fold(FoldAdaptor adaptor) {
+  if (auto val = dyn_cast_or_null<FloatAttr>(adaptor.getOperand())) {
+    return BoolAttr::get(val.getContext(), val.getValue().isNaN());
+  }
+  if (auto splat = dyn_cast_or_null<SplatElementsAttr>(adaptor.getOperand())) {
+    return DenseElementsAttr::get(
+        cast<ShapedType>(getType()),
+        APInt(1, splat.getSplatValue<APFloat>().isNaN()));
+  }
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
+// IsNormalOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::IsNormalOp::fold(FoldAdaptor adaptor) {
+  if (auto val = dyn_cast_or_null<FloatAttr>(adaptor.getOperand())) {
+    return BoolAttr::get(val.getContext(), val.getValue().isNormal());
+  }
+  if (auto splat = dyn_cast_or_null<SplatElementsAttr>(adaptor.getOperand())) {
+    return DenseElementsAttr::get(
+        cast<ShapedType>(getType()),
+        APInt(1, splat.getSplatValue<APFloat>().isNormal()));
+  }
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // TanOp folder
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Math/canonicalize.mlir b/mlir/test/Dialect/Math/canonicalize.mlir
index d24f7649269fe..fff5f7a4390d9 100644
--- a/mlir/test/Dialect/Math/canonicalize.mlir
+++ b/mlir/test/Dialect/Math/canonicalize.mlir
@@ -492,3 +492,75 @@ func.func @abs_poison() -> f32 {
   %1 = math.absf %0 : f32
   return %1 : f32
 }
+
+// CHECK-LABEL: @isfinite_fold
+// CHECK: %[[cst:.+]] = arith.constant true : i1
+  // CHECK: return %[[cst]]
+func.func @isfinite_fold() -> i1 {
+  %c = arith.constant 2.0 : f32
+  %r = math.isfinite %c : f32
+  return %r : i1
+}
+
+// CHECK-LABEL: @isfinite_fold_vec
+// CHECK: %[[cst:.+]] = arith.constant dense<true> : vector<4xi1>
+// CHECK: return %[[cst]]
+func.func @isfinite_fold_vec() -> (vector<4xi1>) {
+  %v1 = arith.constant dense<2.0> : vector<4xf32>
+  %0 = math.isfinite %v1 : vector<4xf32>
+  return %0 : vector<4xi1>
+}
+
+// CHECK-LABEL: @isinf_fold
+// CHECK: %[[cst:.+]] = arith.constant false : i1
+  // CHECK: return %[[cst]]
+func.func @isinf_fold() -> i1 {
+  %c = arith.constant 2.0 : f32
+  %r = math.isinf %c : f32
+  return %r : i1
+}
+
+// CHECK-LABEL: @isinf_fold_vec
+// CHECK: %[[cst:.+]] = arith.constant dense<false> : vector<4xi1>
+// CHECK: return %[[cst]]
+func.func @isinf_fold_vec() -> (vector<4xi1>) {
+  %v1 = arith.constant dense<2.0> : vector<4xf32>
+  %0 = math.isinf %v1 : vector<4xf32>
+  return %0 : vector<4xi1>
+}
+
+// CHECK-LABEL: @isnan_fold
+// CHECK: %[[cst:.+]] = arith.constant false : i1
+// CHECK: return %[[cst]]
+func.func @isnan_fold() -> i1 {
+  %c = arith.constant 2.0 : f32
+  %r = math.isnan %c : f32
+  return %r : i1
+}
+
+// CHECK-LABEL: @isnan_fold_vec
+// CHECK: %[[cst:.+]] = arith.constant dense<false> : vector<4xi1>
+// CHECK: return %[[cst]]
+func.func @isnan_fold_vec() -> (vector<4xi1>) {
+  %v1 = arith.constant dense<2.0> : vector<4xf32>
+  %0 = math.isnan %v1 : vector<4xf32>
+  return %0 : vector<4xi1>
+}
+
+// CHECK-LABEL: @isnormal_fold
+// CHECK: %[[cst:.+]] = arith.constant true : i1
+  // CHECK: return %[[cst]]
+func.func @isnormal_fold() -> i1 {
+  %c = arith.constant 2.0 : f32
+  %r = math.isnormal %c : f32
+  return %r : i1
+}
+
+// CHECK-LABEL: @isnormal_fold_vec
+// CHECK: %[[cst:.+]] = arith.constant dense<true> : vector<4xi1>
+// CHECK: return %[[cst]]
+func.func @isnormal_fold_vec() -> (vector<4xi1>) {
+  %v1 = arith.constant dense<2.0> : vector<4xf32>
+  %0 = math.isnormal %v1 : vector<4xf32>
+  return %0 : vector<4xi1>
+}