[Mlir-commits] [mlir] [MLIR][Math] Add lowering for isnan and isfinite (PR #128125)

[William Moses]

https://github.com/wsmoses created https://github.com/llvm/llvm-project/pull/128125

None

>From d4c510097f5e054b0bece96373300ad9b6aa54ee Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh at wsmoses.com>
Date: Fri, 21 Feb 2025 00:02:02 -0500
Subject: [PATCH] [MLIR][Math] Add lowering for isnan and isfinite

---
 mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp | 36 +++++++++++++++++++
 .../Conversion/MathToLLVM/math-to-llvm.mlir   | 20 +++++++++++
 2 files changed, 56 insertions(+)

diff --git a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
index 85ec288268aeb..fe0d2a6ae9747 100644
--- a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
+++ b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
@@ -286,6 +286,40 @@ struct RsqrtOpLowering : public ConvertOpToLLVMPattern<math::RsqrtOp> {
   }
 };
 
+struct IsNaNOpLowering : public ConvertOpToLLVMPattern<math::IsNaNOp> {
+  using ConvertOpToLLVMPattern<math::IsNaNOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(math::IsNaNOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto operandType = adaptor.getOperand().getType();
+
+    if (!operandType || !LLVM::isCompatibleType(operandType))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<LLVM::IsFPClass>(op, op.getType(),
+                                                 adaptor.getOperand(), 3);
+    return success();
+  }
+};
+
+struct IsFiniteOpLowering : public ConvertOpToLLVMPattern<math::IsFiniteOp> {
+  using ConvertOpToLLVMPattern<math::IsFiniteOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(math::IsFiniteOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto operandType = adaptor.getOperand().getType();
+
+    if (!operandType || !LLVM::isCompatibleType(operandType))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<LLVM::IsFPClass>(op, op.getType(),
+                                                 adaptor.getOperand(), 504);
+    return success();
+  }
+};
+
 struct ConvertMathToLLVMPass
     : public impl::ConvertMathToLLVMPassBase<ConvertMathToLLVMPass> {
   using Base::Base;
@@ -307,6 +341,8 @@ void mlir::populateMathToLLVMConversionPatterns(
     bool approximateLog1p, PatternBenefit benefit) {
   if (approximateLog1p)
     patterns.add<Log1pOpLowering>(converter, benefit);
+  patterns.add<IsNaNOpLowering, IsFiniteOpLowering>(converter);
+
   // clang-format off
   patterns.add<
     AbsFOpLowering,
diff --git a/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir b/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
index 45a37af293890..974743a55932b 100644
--- a/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
+++ b/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
@@ -263,6 +263,26 @@ func.func @ctpop_scalable_vector(%arg0 : vector<[4]xi32>) -> vector<[4]xi32> {
 
 // -----
 
+// CHECK-LABEL: func @isnan_double(
+// CHECK-SAME: f64
+func.func @isnan_double(%arg0 : f64) {
+  // CHECK: "llvm.intr.is.fpclass"(%arg0) <{bit = 3 : i32}> : (f64) -> i1
+  %0 = math.isnan %arg0 : f64
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: func @isfinite_double(
+// CHECK-SAME: f64
+func.func @isfinite_double(%arg0 : f64) {
+  // CHECK: "llvm.intr.is.fpclass"(%arg0) <{bit = 504 : i32}> : (f64) -> i1
+  %0 = math.isfinite %arg0 : f64
+  func.return
+}
+
+// -----
+
 // CHECK-LABEL: func @rsqrt_double(
 // CHECK-SAME: f64
 func.func @rsqrt_double(%arg0 : f64) {