[mlir][Arith] Add arith.is_nan and arith.is_inf predicates

[krzysz00]

Summary

Both LLVM and SPIR-V have some form of "is this float a NaN/Inf"
operation (though LLVM's uses the rather opaque "is.fpclass"
intrinsic), which is not exposed in MLIR.

This has lead to awkward workarounds in -arith-expands-ops where a NaN
test was performed by comparing an operation to itself. This commit
resolves that issue.

[kuhar]

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

991

nit: I wonder if we should call these 'is_nan` and is_inf in the name of readability. We already have a few airth ops with underscore in the name, e.g., index_cast or addi_extended.

mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp

58

Would it be possible to use enums here instead of hardcoding the bitpatterns?

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

1131–1132

nit: is there an overload that doesn't require the 'success' parameter that we could use here?

[krzysz00]

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

991

On the other hand, the relevant operations in C/C++/SPIR-V/... are without an underscore and this'll make it easier to find

mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp

58

There doesn't seem to be one, given that flang uses 0b constants for the same purpose

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

1131–1132

I don't think so, and I'm not sure it'd be worth adding one.

[krzysz00]

Update flang tests

[kuhar]

LGTM for SPIR-V but please get at least one more approval from someone familiar with the LLVM lowering path

[dcaballe]

LGTM % the comments. Thanks for the contribution!

mlir/include/mlir/Conversion/ArithCommon/AttrToLLVMConverter.h

58

-> ///

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

86

Is the comment correct?

991

+1 to using underscore, for consistency within the Arith dialect

mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp

58

Yeah, an enum would be great or at least a comment explaining what each bitfield means.

[krzysz00]

Address review comments

[krzysz00]

Address comments

mlir/include/mlir/Conversion/ArithCommon/AttrToLLVMConverter.h

if (arithFMFAttr) {
       llvm::StringRef targetAttrName = TargetOp::getFastmathAttrName();
       convertedAttr.set(targetAttrName,
                         convertArithFastMathAttrToLLVM(arithFMFAttr));
     }
   }
 
   ArrayRef<NamedAttribute> getAttrs() const { return convertedAttr.getAttrs(); }
 
-private:
+protected:
   NamedAttrList convertedAttr;
 };
+
+// Wrapper around AttrConvertFastMathToLLVM that also sets the "kinds" attribute
+// to the bitmask specified in `Kinds`, which is used for converting operations
+// that lower to llvm.is.fpclass.
+template <unsigned Kinds, typename SourceOp, typename TargetOp>
+class AttrConvertAddFpclassKinds
+    : public AttrConvertFastMathToLLVM<SourceOp, TargetOp> {
+public:
+  AttrConvertAddFpclassKinds(SourceOp op)
+      : AttrConvertFastMathToLLVM<SourceOp, TargetOp>(op) {
+    convertedAttr.set(
+        "kinds",
+        IntegerAttr::get(IntegerType::get(op.getContext(), 32), Kinds));
+  }
+
+  ArrayRef<NamedAttribute> getAttrs() const { return convertedAttr.getAttrs(); }
+
+protected:
+  using AttrConvertFastMathToLLVM<SourceOp, TargetOp>::convertedAttr;
+};
 } // namespace arith
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_ARITHCOMMON_ATTRTOLLVMCONVERTER_H

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

class Arith_FloatBinaryOp<string mnemonic, list<Trait> traits = []> :
     Arguments<(ins FloatLike:$lhs, FloatLike:$rhs,
       DefaultValuedAttr<
         Arith_FastMathAttr, "::mlir::arith::FastMathFlags::none">:$fastmath)>,
     Results<(outs FloatLike:$result)> {
   let assemblyFormat = [{ $lhs `,` $rhs (`fastmath` `` $fastmath^)?
                           attr-dict `:` type($result) }];
 }
 
+// Base class for floating point unary operations.
+class Arith_FloatPredicateOp<string mnemonic, list<Trait> traits = []> :
+    Arith_Op<mnemonic,
+      !listconcat([DeclareOpInterfaceMethods<ArithFastMathInterface>, Pure,
+        TypesMatchWith<"result type has i1 element type and same shape as operands",
+        "operand", "result", "::getI1SameShape($_self)">], traits)>,
+    Arguments<(ins FloatLike:$operand,
+      DefaultValuedAttr<
+        Arith_FastMathAttr, "::mlir::arith::FastMathFlags::none">:$fastmath)>,
+    Results<(outs BoolLike:$result)> {
+  let assemblyFormat = [{ $operand (`fastmath` `` $fastmath^)?
+                          attr-dict `:` type($operand) }];
+}
+
 // Base class for arithmetic cast operations. Requires a single operand and
 // result. If either is a shaped type, then the other must be of the same shape.
 class Arith_CastOp<string mnemonic, TypeConstraint From, TypeConstraint To,
                    list<Trait> traits = []> :
     Arith_Op<mnemonic, traits # [Pure, SameOperandsAndResultShape,
       DeclareOpInterfaceMethods<CastOpInterface>]>,
     Arguments<(ins From:$in)>,
     Results<(outs To:$out)> {

 //===----------------------------------------------------------------------===//
 
 def Arith_RemFOp : Arith_FloatBinaryOp<"remf"> {
   let summary = "floating point division remainder operation";
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// IsNanOp
+//===----------------------------------------------------------------------===//
+def Arith_IsNanOp : Arith_FloatPredicateOp<"isnan"> {
+  let summary = "Returns true for IEEE NaN inputs";
+  let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// IsInfOp
+//===----------------------------------------------------------------------===//
+def Arith_IsInfOp : Arith_FloatPredicateOp<"isinf"> {
+  let summary = "Returns true for infinite float inputs";
+  let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // ExtUIOp
 //===----------------------------------------------------------------------===//
 
 def Arith_ExtUIOp : Arith_IToICastOp<"extui"> {
   let summary = "integer zero extension operation";
   let description = [{
     The integer zero extension operation takes an integer input of
     width M and an integer destination type of width N. The destination

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

string llvmBuilder = [{
     auto *inst = createIntrinsicCall(
       builder, llvm::Intrinsic::}] # llvmEnumName # [{,
       {$in, builder.getInt1(op.getIsIntMinPoison())}, }]
       # declTypes # [{);
     $res = inst;
   }];
 }
 
-def LLVM_IsFPClass : LLVM_OneResultIntrOp<"is.fpclass", [], [0], [Pure]> {
-  let arguments = (ins LLVM_ScalarOrVectorOf<LLVM_AnyFloat>:$in, I32Attr:$bit);
+def LLVM_IsFPClass : LLVM_OneResultIntrOp<"is.fpclass", [], [0], [Pure],
+                                          /*requiresFastmath=*/1> {
+  let arguments = (ins LLVM_ScalarOrVectorOf<LLVM_AnyFloat>:$in,
+                       I32Attr:$kinds,
+                       DefaultValuedAttr<LLVM_FastmathFlagsAttr, "{}">:$fastmathFlags);
   string mlirBuilder = [{
     auto op = $_builder.create<$_qualCppClassName>($_location,
-      $_resultType, $in, $_int_attr($bit));
-      $res = op;
+      $_resultType, $in, $_int_attr($kinds), nullptr);
+    moduleImport.setFastmathFlagsAttr(inst, op);
   }];
   string llvmBuilder = [{
     auto *inst = createIntrinsicCall(
       builder, llvm::Intrinsic::}] # llvmEnumName # [{,
-      {$in, builder.getInt32(op.getBit())},
+      {$in, builder.getInt32(op.getKinds())},
       }] # declTypes # [{);
     $res = inst;
   }];
 }
 
 def LLVM_CopySignOp : LLVM_BinarySameArgsIntrOpF<"copysign">;
 def LLVM_CosOp : LLVM_UnaryIntrOpF<"cos">;
 def LLVM_ExpOp : LLVM_UnaryIntrOpF<"exp">;

mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp

 using ExtSIOpLowering =
     VectorConvertToLLVMPattern<arith::ExtSIOp, LLVM::SExtOp>;
 using ExtUIOpLowering =
     VectorConvertToLLVMPattern<arith::ExtUIOp, LLVM::ZExtOp>;
 using FPToSIOpLowering =
     VectorConvertToLLVMPattern<arith::FPToSIOp, LLVM::FPToSIOp>;
 using FPToUIOpLowering =
     VectorConvertToLLVMPattern<arith::FPToUIOp, LLVM::FPToUIOp>;
+template <typename SourceOp, typename TargetOp>
+using InfKinds =
+    arith::AttrConvertAddFpclassKinds<(1 << 2) | (1 << 9), SourceOp, TargetOp>;
+using IsInfLowering =
+    VectorConvertToLLVMPattern<arith::IsInfOp, LLVM::IsFPClass, InfKinds>;
+template <typename SourceOp, typename TargetOp>
+using NanKinds =
+    arith::AttrConvertAddFpclassKinds<(1 << 0) | (1 << 1), SourceOp, TargetOp>;
+using IsNanLowering =
+    VectorConvertToLLVMPattern<arith::IsNanOp, LLVM::IsFPClass, NanKinds>;
 using MaxFOpLowering =
     VectorConvertToLLVMPattern<arith::MaxFOp, LLVM::MaxNumOp,
                                arith::AttrConvertFastMathToLLVM>;
 using MaxSIOpLowering =
     VectorConvertToLLVMPattern<arith::MaxSIOp, LLVM::SMaxOp>;
 using MaxUIOpLowering =
     VectorConvertToLLVMPattern<arith::MaxUIOp, LLVM::UMaxOp>;
 using MinFOpLowering =
patterns.add<
     DivUIOpLowering,
     ExtFOpLowering,
     ExtSIOpLowering,
     ExtUIOpLowering,
     FPToSIOpLowering,
     FPToUIOpLowering,
     IndexCastOpSILowering,
     IndexCastOpUILowering,
+    IsInfLowering,
+    IsNanLowering,
     MaxFOpLowering,
     MaxSIOpLowering,
     MaxUIOpLowering,
     MinFOpLowering,
     MinSIOpLowering,
     MinUIOpLowering,
     MulFOpLowering,
     MulIOpLowering,

mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp

patterns.add<
     spirv::ElementwiseOpPattern<arith::ShRUIOp, spirv::ShiftRightLogicalOp>,
     spirv::ElementwiseOpPattern<arith::ShRSIOp, spirv::ShiftRightArithmeticOp>,
     spirv::ElementwiseOpPattern<arith::NegFOp, spirv::FNegateOp>,
     spirv::ElementwiseOpPattern<arith::AddFOp, spirv::FAddOp>,
     spirv::ElementwiseOpPattern<arith::SubFOp, spirv::FSubOp>,
     spirv::ElementwiseOpPattern<arith::MulFOp, spirv::FMulOp>,
     spirv::ElementwiseOpPattern<arith::DivFOp, spirv::FDivOp>,
     spirv::ElementwiseOpPattern<arith::RemFOp, spirv::FRemOp>,
+    spirv::ElementwiseOpPattern<arith::IsNanOp, spirv::IsNanOp>,
+    spirv::ElementwiseOpPattern<arith::IsInfOp, spirv::IsInfOp>,
     TypeCastingOpPattern<arith::ExtUIOp, spirv::UConvertOp>, ExtUII1Pattern,
     TypeCastingOpPattern<arith::ExtSIOp, spirv::SConvertOp>, ExtSII1Pattern,
     TypeCastingOpPattern<arith::ExtFOp, spirv::FConvertOp>,
     TypeCastingOpPattern<arith::TruncIOp, spirv::SConvertOp>, TruncII1Pattern,
     TypeCastingOpPattern<arith::TruncFOp, spirv::FConvertOp>,
     TypeCastingOpPattern<arith::UIToFPOp, spirv::ConvertUToFOp>, UIToFPI1Pattern,
     TypeCastingOpPattern<arith::SIToFPOp, spirv::ConvertSToFOp>,
     TypeCastingOpPattern<arith::FPToUIOp, spirv::ConvertFToUOp>,

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

OpFoldResult arith::RemFOp::fold(FoldAdaptor adaptor) {
   return constFoldBinaryOp<FloatAttr>(adaptor.getOperands(),
                                       [](const APFloat &a, const APFloat &b) {
                                         APFloat result(a);
                                         (void)result.remainder(b);
                                         return result;
                                       });
 }
 
+static Attribute getBoolAttribute(Type type, MLIRContext *ctx, bool value) {
+  auto boolAttr = BoolAttr::get(ctx, value);
+  ShapedType shapedType = llvm::dyn_cast_or_null<ShapedType>(type);
+  if (!shapedType)
+    return boolAttr;
+  return DenseElementsAttr::get(shapedType, boolAttr);
+}
+
+//===----------------------------------------------------------------------===//
+// IsNanOp
+//===----------------------------------------------------------------------===//
+OpFoldResult IsNanOp::fold(FoldAdaptor adaptor) {
+  if (bitEnumContainsAll(getFastmath(), FastMathFlags::nnan))
+    return getBoolAttribute(getType(), getContext(), false);
+  return constFoldCastOp<FloatAttr, IntegerAttr>(
+      adaptor.getOperands(), getType(),
+      [](const APFloat &x, bool &success) -> APInt {
+        success = true;
+        return APInt(1, x.isNaN());
+      });
+}
+
+//===----------------------------------------------------------------------===//
+// IsInfOp
+//===----------------------------------------------------------------------===//
+OpFoldResult IsInfOp::fold(FoldAdaptor adaptor) {
+  if (bitEnumContainsAll(getFastmath(), FastMathFlags::ninf))
+    return getBoolAttribute(getType(), getContext(), false);
+  return constFoldCastOp<FloatAttr, IntegerAttr>(
+      adaptor.getOperands(), getType(),
+      [](const APFloat &x, bool &success) -> APInt {
+        success = true;
+        return APInt(1, x.isInfinity());
+      });
+}
+
 //===----------------------------------------------------------------------===//
 // Utility functions for verifying cast ops
 //===----------------------------------------------------------------------===//
 
 template <typename... Types>
 using type_list = std::tuple<Types...> *;
 
 /// Returns a non-null type only if the provided type is one of the allowed
static bool applyCmpPredicateToEqualOperands(arith::CmpIPredicate predicate) {
   case arith::CmpIPredicate::sgt:
   case arith::CmpIPredicate::ult:
   case arith::CmpIPredicate::ugt:
     return false;
   }
   llvm_unreachable("unknown cmpi predicate kind");
 }
 
-static Attribute getBoolAttribute(Type type, MLIRContext *ctx, bool value) {
-  auto boolAttr = BoolAttr::get(ctx, value);
-  ShapedType shapedType = llvm::dyn_cast_or_null<ShapedType>(type);
-  if (!shapedType)
-    return boolAttr;
-  return DenseElementsAttr::get(shapedType, boolAttr);
-}
-
 static std::optional<int64_t> getIntegerWidth(Type t) {
   if (auto intType = llvm::dyn_cast<IntegerType>(t)) {
     return intType.getWidth();
   }
   if (auto vectorIntType = llvm::dyn_cast<VectorType>(t)) {
     return llvm::cast<IntegerType>(vectorIntType.getElementType()).getWidth();
   }
   return std::nullopt;

mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp

LogicalResult matchAndRewrite(OpTy op,
     // If any operand is NaN, 'cmp' will be true (and 'select' returns 'lhs').
     static_assert(pred == arith::CmpFPredicate::UGT ||
                       pred == arith::CmpFPredicate::ULT,
                   "pred must be either UGT or ULT");
     Value cmp = rewriter.create<arith::CmpFOp>(loc, pred, lhs, rhs);
     Value select = rewriter.create<arith::SelectOp>(loc, cmp, lhs, rhs);
 
     // Handle the case where rhs is NaN: 'isNaN(rhs) ? rhs : select'.
-    Value isNaN = rewriter.create<arith::CmpFOp>(loc, arith::CmpFPredicate::UNO,
-                                                 rhs, rhs);
+    Value isNaN = rewriter.create<arith::IsNanOp>(loc, rhs, op.getFastmath());
     rewriter.replaceOpWithNewOp<arith::SelectOp>(op, isNaN, rhs, select);
     return success();
   }
 };
 
 struct BFloat16ExtFOpConverter : public OpRewritePattern<arith::ExtFOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(arith::ExtFOp op,

mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir

// CHECK: = llvm.shl %arg2, %arg3 : i32
   %16 = arith.shli %arg2, %arg3 : i32
 // CHECK: = llvm.ashr %arg2, %arg3 : i32
   %17 = arith.shrsi %arg2, %arg3 : i32
 // CHECK: = llvm.lshr %arg2, %arg3 : i32
   %18 = arith.shrui %arg2, %arg3 : i32
   return %0, %4 : f32, i32
 }
 
+// CHECK-LABEL: @float_pred_ops
+func.func @float_pred_ops(%arg0: f32) {
+  // CHECK: "llvm.intr.is.fpclass"(%{{.*}}) <{fastmathFlags = #llvm.fastmath<none>, kinds = 516 : i32}> : (f32) -> i1
+  arith.isinf %arg0 : f32
+  // CHECK-NEXT: "llvm.intr.is.fpclass"(%{{.*}}) <{fastmathFlags = #llvm.fastmath<nnan>, kinds = 516 : i32}> : (f32) -> i1
+  arith.isinf %arg0 fastmath <nnan> : f32
+  // CHECK-NEXT: "llvm.intr.is.fpclass"(%{{.*}}) <{fastmathFlags = #llvm.fastmath<none>, kinds = 3 : i32}> : (f32) -> i1
+  arith.isnan %arg0 : f32
+  // CHECK-NEXT: "llvm.intr.is.fpclass"(%{{.*}}) <{fastmathFlags = #llvm.fastmath<ninf>, kinds = 3 : i32}> : (f32) -> i1
+  arith.isnan %arg0 fastmath <ninf> : f32
+  return
+}
+
+// CHECK-LABEL: @vector_float_pred_ops
+func.func @vector_float_pred_ops(%arg0: vector<4xf32>) {
+  // CHECK: "llvm.intr.is.fpclass"(%{{.*}}) <{fastmathFlags = #llvm.fastmath<none>, kinds = 516 : i32}> : (vector<4xf32>) -> vector<4xi1>
+  arith.isinf %arg0 : vector<4xf32>
+  // CHECK: "llvm.intr.is.fpclass"(%{{.*}}) <{fastmathFlags = #llvm.fastmath<none>, kinds = 3 : i32}> : (vector<4xf32>) -> vector<4xi1>
+  arith.isnan %arg0 : vector<4xf32>
+  return
+}
+
 // Checking conversion of index types to integers using i1, assuming no target
 // system would have a 1-bit address space.  Otherwise, we would have had to
 // make this test dependent on the pointer size on the target system.
 // CHECK-LABEL: @index_cast
 func.func @index_cast(%arg0: index, %arg1: i1) {
 // CHECK: = llvm.trunc %0 : i{{.*}} to i1
   %0 = arith.index_cast %arg0: index to i1
 // CHECK-NEXT: = llvm.sext %arg1 : i1 to i{{.*}}

mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir

func.func @float32_binary_scalar(%lhs: f32, %rhs: f32) {
   %2 = arith.mulf %lhs, %rhs: f32
   // CHECK: spirv.FDiv %{{.*}}, %{{.*}}: f32
   %3 = arith.divf %lhs, %rhs: f32
   // CHECK: spirv.FRem %{{.*}}, %{{.*}}: f32
   %4 = arith.remf %lhs, %rhs: f32
   return
 }
 
+// Check float predicate operation conversions
+// CHECK-LABEL: @float32_predicate_scalar
+func.func @float32_predicate_scalar(%arg0 : f32) {
+  // CHECK: spirv.IsNan %{{.*}}: f32
+  %0 = arith.isnan %arg0 : f32
+  // CHECK: spirv.IsInf %{{.*}}: f32
+  %1 = arith.isinf %arg0 : f32
+  return
+}
+
 // Check int vector types.
 // CHECK-LABEL: @int_vector234
 func.func @int_vector234(%arg0: vector<2xi8>, %arg1: vector<4xi64>) {
   // CHECK: spirv.SDiv %{{.*}}, %{{.*}}: vector<2xi8>
   %0 = arith.divsi %arg0, %arg0: vector<2xi8>
   // CHECK: spirv.UDiv %{{.*}}, %{{.*}}: vector<4xi64>
   %1 = arith.divui %arg1, %arg1: vector<4xi64>
   return

mlir/test/Dialect/Arith/canonicalize.mlir

 //       CHECK:   %[[ONE:.*]] = arith.constant -1
 //       CHECK:   return %[[ONE]]
 func.func @foldOrXor6(%arg0: index) -> index {
   %0 = arith.constant -1 : index
   %1 = arith.xori %arg0, %0 : index
   %2 = arith.ori %arg0, %1 : index
   return %2 : index
 }
+
+// -----
+
+// CHECK-LABEL: @foldIsNanFastmath
+//  CHECK-SAME: (%[[ARG:.+]]: f32)
+//      CHECK:   %[[FALSE:.+]] = arith.constant false
+//      CHECK:   return %[[FALSE]]
+func.func @foldIsNanFastmath(%arg0: f32) -> i1 {
+  %0 = arith.isnan %arg0 fastmath <nnan> : f32
+  func.return %0 : i1
+}
+
+// CHECK-LABEL: @foldIsNan
+//      CHECK:   %[[TRUE:.+]] = arith.constant true
+//      CHECK:   return %[[TRUE]]
+func.func @foldIsNan() -> i1 {
+  %cNan = arith.constant 0x7FFFFFFF : f32
+  %0 = arith.isnan %cNan : f32
+  func.return %0 : i1
+}
+
+// CHECK-LABEL: @foldNanIsNotNanWithFastmath
+//      CHECK:   %[[FALSE:.+]] = arith.constant false
+//      CHECK:   return %[[FALSE]]
+func.func @foldNanIsNotNanWithFastmath() -> i1 {
+  %cNan = arith.constant 0x7FFFFFFF : f32
+  %0 = arith.isnan %cNan fastmath<nnan> : f32
+  func.return %0 : i1
+}
+
+
+// CHECK-LABEL: @foldIsNotNan
+//      CHECK:   %[[FALSE:.+]] = arith.constant false
+//      CHECK:   return %[[FALSE]]
+func.func @foldIsNotNan() -> i1 {
+  %cNan = arith.constant 1.0 : f32
+  %0 = arith.isnan %cNan : f32
+  func.return %0 : i1
+}
+
+// CHECK-LABEL: @foldIsInfFastmath
+//      CHECK:   %[[FALSE:.+]] = arith.constant false
+//      CHECK:   return %[[FALSE]]
+func.func @foldIsInfFastmath(%arg0: f32) -> i1 {
+  %0 = arith.isinf %arg0 fastmath <ninf> : f32
+  func.return %0 : i1
+}
+
+// CHECK-LABEL: @foldIsInf
+//      CHECK:   %[[TRUE:.+]] = arith.constant true
+//      CHECK:   return %[[TRUE]]
+func.func @foldIsInf() -> i1 {
+  %cInf = arith.constant 0x7F800000 : f32
+  %0 = arith.isinf %cInf : f32
+  func.return %0 : i1
+}
+
+// CHECK-LABEL: @foldInfIsNotInfWithFastmath
+//      CHECK:   %[[FALSE:.+]] = arith.constant false
+//      CHECK:   return %[[FALSE]]
+func.func @foldInfIsNotInfWithFastmath() -> i1 {
+  %cInf = arith.constant 0x7F800000 : f32
+  %0 = arith.isinf %cInf fastmath<ninf> : f32
+  func.return %0 : i1
+}
+
+
+// CHECK-LABEL: @foldIsNotInf
+//      CHECK:   %[[FALSE:.+]] = arith.constant false
+//      CHECK:   return %[[FALSE]]
+func.func @foldIsNotInf() -> i1 {
+  %cInf = arith.constant 1.0 : f32
+  %0 = arith.isinf %cInf : f32
+  func.return %0 : i1
+}

mlir/test/Dialect/Arith/expand-ops.mlir

 // CHECK-LABEL: func @maxf
 func.func @maxf(%a: f32, %b: f32) -> f32 {
   %result = arith.maxf %a, %b : f32
   return %result : f32
 }
 // CHECK-SAME: %[[LHS:.*]]: f32, %[[RHS:.*]]: f32)
 // CHECK-NEXT: %[[CMP:.*]] = arith.cmpf ugt, %[[LHS]], %[[RHS]] : f32
 // CHECK-NEXT: %[[SELECT:.*]] = arith.select %[[CMP]], %[[LHS]], %[[RHS]] : f32
-// CHECK-NEXT: %[[IS_NAN:.*]] = arith.cmpf uno, %[[RHS]], %[[RHS]] : f32
+// CHECK-NEXT: %[[IS_NAN:.*]] = arith.isnan %[[RHS]] : f32
 // CHECK-NEXT: %[[RESULT:.*]] = arith.select %[[IS_NAN]], %[[RHS]], %[[SELECT]] : f32
 // CHECK-NEXT: return %[[RESULT]] : f32
 
 // -----
 
 // CHECK-LABEL: func @maxf_vector
 func.func @maxf_vector(%a: vector<4xf16>, %b: vector<4xf16>) -> vector<4xf16> {
   %result = arith.maxf %a, %b : vector<4xf16>
   return %result : vector<4xf16>
 }
 // CHECK-SAME: %[[LHS:.*]]: vector<4xf16>, %[[RHS:.*]]: vector<4xf16>)
 // CHECK-NEXT: %[[CMP:.*]] = arith.cmpf ugt, %[[LHS]], %[[RHS]] : vector<4xf16>
 // CHECK-NEXT: %[[SELECT:.*]] = arith.select %[[CMP]], %[[LHS]], %[[RHS]]
-// CHECK-NEXT: %[[IS_NAN:.*]] = arith.cmpf uno, %[[RHS]], %[[RHS]] : vector<4xf16>
+// CHECK-NEXT: %[[IS_NAN:.*]] = arith.isnan  %[[RHS]] : vector<4xf16>
 // CHECK-NEXT: %[[RESULT:.*]] = arith.select %[[IS_NAN]], %[[RHS]], %[[SELECT]]
 // CHECK-NEXT: return %[[RESULT]] : vector<4xf16>
 
 // -----
 
 // CHECK-LABEL: func @minf
 func.func @minf(%a: f32, %b: f32) -> f32 {
   %result = arith.minf %a, %b : f32
   return %result : f32
 }
 
 // CHECK-SAME: %[[LHS:.*]]: f32, %[[RHS:.*]]: f32)
 // CHECK-NEXT: %[[CMP:.*]] = arith.cmpf ult, %[[LHS]], %[[RHS]] : f32
 // CHECK-NEXT: %[[SELECT:.*]] = arith.select %[[CMP]], %[[LHS]], %[[RHS]] : f32
-// CHECK-NEXT: %[[IS_NAN:.*]] = arith.cmpf uno, %[[RHS]], %[[RHS]] : f32
+// CHECK-NEXT: %[[IS_NAN:.*]] = arith.isnan %[[RHS]] : f32
 // CHECK-NEXT: %[[RESULT:.*]] = arith.select %[[IS_NAN]], %[[RHS]], %[[SELECT]] : f32
 // CHECK-NEXT: return %[[RESULT]] : f32
 
 // -----
 
 func.func @truncf_f32(%arg0 : f32) -> bf16 {
     %0 = arith.truncf %arg0 : f32 to bf16
     return %0 : bf16

mlir/test/Dialect/Arith/ops.mlir

 }
 
 // CHECK-LABEL: test_remf_scalable_vector
 func.func @test_remf_scalable_vector(%arg0 : vector<[8]xf64>, %arg1 : vector<[8]xf64>) -> vector<[8]xf64> {
   %0 = arith.remf %arg0, %arg1 : vector<[8]xf64>
   return %0 : vector<[8]xf64>
 }
 
+// CHECK-LABEL: test_isnan
+func.func @test_isnan(%arg0 : f32) -> i1 {
+  %0 = arith.isnan %arg0 : f32
+  func.return %0 : i1
+}
+
+// CHECK-LABEL: test_isnan_vector
+func.func @test_isnan_vector(%arg0 : vector<2x2xf32>) -> vector<2x2xi1> {
+  %0 = arith.isnan %arg0 : vector<2x2xf32>
+  func.return %0 : vector<2x2xi1>
+}
+
+// CHECK-LABEL: test_isinf
+func.func @test_isinf(%arg0 : f32) -> i1 {
+  %0 = arith.isinf %arg0 : f32
+  func.return %0 : i1
+}
+
+// CHECK-LABEL: test_isinf_vector
+func.func @test_isinf_vector(%arg0 : vector<2x2xf32>) -> vector<2x2xi1> {
+  %0 = arith.isinf %arg0 : vector<2x2xf32>
+  func.return %0 : vector<2x2xi1>
+}
+
 // CHECK-LABEL: test_extui
 func.func @test_extui(%arg0 : i32) -> i64 {
   %0 = arith.extui %arg0 : i32 to i64
   return %0 : i64
 }
 
 // CHECK-LABEL: test_extui_tensor
 func.func @test_extui_tensor(%arg0 : tensor<8x8xi32>) -> tensor<8x8xi64> {

mlir/test/Target/LLVMIR/Import/intrinsic.ll

define void @fmuladd_test(float %0, float %1, <8 x float> %2, ptr %3) {
   %8 = call <8 x float> @llvm.fma.v8f32(<8 x float> %2, <8 x float> %2, <8 x float> %2)
   ; CHECK: "llvm.intr.prefetch"(%{{.*}}) <{cache = 1 : i32, hint = 3 : i32, rw = 0 : i32}> : (!llvm.ptr) -> ()
   call void @llvm.prefetch.p0(ptr %3, i32 0, i32 3, i32 1)
   ret void
 }
 
 ; CHECK-LABEL:  llvm.func @fpclass_test
 define void @fpclass_test(float %0, <8 x float> %1) {
-  ; CHECK: "llvm.intr.is.fpclass"(%{{.*}}) <{bit = 0 : i32}> : (f32) -> i1
+  ; CHECK: "llvm.intr.is.fpclass"(%{{.*}}) <{fastmathFlags = #llvm.fastmath<none>, kinds = 0 : i32}> : (f32) -> i1
   %3 = call i1 @llvm.is.fpclass.f32(float %0, i32 0)
-  ; CHECK: "llvm.intr.is.fpclass"(%{{.*}}) <{bit = 1 : i32}> : (vector<8xf32>) -> vector<8xi1>
+  ; CHECK: "llvm.intr.is.fpclass"(%{{.*}}) <{fastmathFlags = #llvm.fastmath<none>, kinds = 1 : i32}> : (vector<8xf32>) -> vector<8xi1>
   %4 = call <8 x i1> @llvm.is.fpclass.v8f32(<8 x float> %1, i32 1)
   ret void
 }
 
 ; CHECK-LABEL:  llvm.func @exp_test
 define void @exp_test(float %0, <8 x float> %1) {
   ; CHECK: llvm.intr.exp(%{{.*}}) : (f32) -> f32
   %3 = call float @llvm.exp.f32(float %0)

mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir

llvm.func @intrinsics(%arg0: f32, %arg1: f32, %arg2: vector<8xf32>, %arg3: !llvm.ptr<i8>) {
   // CHECK: call void @llvm.prefetch.p0(ptr %3, i32 0, i32 3, i32 1)
   "llvm.intr.prefetch"(%arg3) <{cache = 1 : i32, hint = 3 : i32, rw = 0 : i32}> : (!llvm.ptr<i8>) -> ()
   llvm.return
 }
 
 // CHECK-LABEL: @fpclass_test
 llvm.func @fpclass_test(%arg0: f32) -> i1 {
   // CHECK: call i1 @llvm.is.fpclass
-  %0 = "llvm.intr.is.fpclass"(%arg0) <{bit = 0 : i32 }>: (f32) -> i1
+  %0 = "llvm.intr.is.fpclass"(%arg0) <{fastmathFlags = #llvm.fastmath<nnan>, kinds = 3 : i32 }>: (f32) -> i1
   llvm.return %0 : i1
 }
 
 // CHECK-LABEL: @exp_test
 llvm.func @exp_test(%arg0: f32, %arg1: vector<8xf32>) {
   // CHECK: call float @llvm.exp.f32
   "llvm.intr.exp"(%arg0) : (f32) -> f32
   // CHECK: call <8 x float> @llvm.exp.v8f32