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

[llvmlistbot at llvm.org]

Author: William Moses
Date: 2025-02-20T09:59:55-06:00
New Revision: 1b610e69d9621e00714b126e646bebb912b1e9a1

URL: https://github.com/llvm/llvm-project/commit/1b610e69d9621e00714b126e646bebb912b1e9a1
DIFF: https://github.com/llvm/llvm-project/commit/1b610e69d9621e00714b126e646bebb912b1e9a1.diff

LOG: [MLIR][Math] Add floating point value folders (#127947)

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Math/IR/MathOps.td
    mlir/lib/Dialect/Math/IR/MathOps.cpp
    mlir/test/Dialect/Math/canonicalize.mlir

Removed: 
    


################################################################################
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..3743768d901e3 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
+// 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
+// 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
+// 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
+// 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>
+}