[MLIR][Math] Add lowering for isnan and isfinite

Author: wsmoses

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,

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) {