[mlir][spirv] Add support for fast math mode

[antiagainst]

Summary

This commit introduces a new option to SPIRVConversionOptions
to allow enabling fast math mode. With it, various patterns
would assume no NaN/infinity for floating point values and
avoid guards to check them. This is particularly useful for
CodeGen towards WebGPU environment, where fast math is assumed.

Along the way, fixed the conversion for arith.minf/maxf to
handle the NaN cases properly for Shader cases.

Part of https://github.com/llvm/llvm-project/issues/57584.

[ThomasRaoux]

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

934

nit: that's not really a cmp? Is there a better name?

[antiagainst]

Address comments

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

934

Good point. Changed.

mlir/include/mlir/Conversion/Passes.td

 def ConvertArithmeticToSPIRV : Pass<"convert-arith-to-spirv"> {
   let summary = "Convert Arithmetic dialect to SPIR-V dialect";
   let constructor = "mlir::arith::createConvertArithmeticToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];
   let options = [
     Option<"emulateNon32BitScalarTypes", "emulate-non-32-bit-scalar-types",
            "bool", /*default=*/"true",
            "Emulate non-32-bit scalar types with 32-bit ones if "
-           "missing native support">
+           "missing native support">,
+    Option<"enableFastMath", "enable-fast-math",
+           "bool", /*default=*/"false",
+           "Enable fast math mode (assuming no NaN and infinity for floating "
+           "point values) when performing conversion">
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // ArmNeon2dToIntr
 //===----------------------------------------------------------------------===//
 
 def ConvertArmNeon2dToIntr : Pass<"arm-neon-2d-to-intr"> {

mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h

 
 namespace mlir {
 
 //===----------------------------------------------------------------------===//
 // Type Converter
 //===----------------------------------------------------------------------===//
 
 struct SPIRVConversionOptions {
+  /// The number of bits to store a boolean value.
+  unsigned boolNumBits{8};
+
   /// Whether to emulate non-32-bit scalar types with 32-bit scalar types if
   /// no native support.
   ///
   /// Non-32-bit scalar types require special hardware support that may not
   /// exist on all GPUs. This is reflected in SPIR-V as that non-32-bit scalar
   /// types require special capabilities or extensions. This option controls
   /// whether to use 32-bit types to emulate, if a scalar type of a certain
   /// bitwidth is not supported in the target environment. This requires the
   /// runtime to also feed in data with a matched bitwidth and layout for
   /// interface types. The runtime can do that by inspecting the SPIR-V
   /// module.
   ///
   /// If the original scalar type has less than 32-bit, a multiple of its
   /// values will be packed into one 32-bit value to be memory efficient.
   bool emulateNon32BitScalarTypes{true};
 
   /// Use 64-bit integers to convert index types.
   bool use64bitIndex{false};
 
-  /// The number of bits to store a boolean value. It is eight bits by
-  /// default.
-  unsigned boolNumBits{8};
+  /// Whether to enable fast math mode during conversion. If true, various
+  /// patterns would assume no NaN/infinity numbers as inputs, and thus there
+  /// will be no special guards emitted to check and handle such cases.
+  bool enableFastMathMode{false};
 };
 
 /// Type conversion from builtin types to SPIR-V types for shader interface.
 ///
 /// For memref types, this converter additionally performs type wrapping to
 /// satisfy shader interface requirements: shader interface types must be
 /// pointers to structs.
 class SPIRVTypeConverter : public TypeConverter {

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

 class SelectOpPattern final : public OpConversionPattern<arith::SelectOp> {
 public:
   using OpConversionPattern<arith::SelectOp>::OpConversionPattern;
   LogicalResult
   matchAndRewrite(arith::SelectOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Converts arith.maxf to spv.GL.FMax.
+template <typename Op, typename SPIRVOp>
+class MinMaxFOpPattern final : public OpConversionPattern<Op> {
+public:
+  using OpConversionPattern<Op>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(Op op, typename Op::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Conversion Helpers
 //===----------------------------------------------------------------------===//
 
 /// Converts the given `srcAttr` into a boolean attribute if it holds an
 /// integral value. Returns null attribute if conversion fails.

 LogicalResult CmpFOpNanNonePattern::matchAndRewrite(
     arith::CmpFOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   if (op.getPredicate() != arith::CmpFPredicate::ORD &&
       op.getPredicate() != arith::CmpFPredicate::UNO)
     return failure();
 
   Location loc = op.getLoc();
+  auto *converter = getTypeConverter<SPIRVTypeConverter>();
 
+  Value replace;
+  if (converter->getOptions().enableFastMathMode) {
+    if (op.getPredicate() == arith::CmpFPredicate::ORD) {
+      // Ordered comparsion checks if neither operand is NaN.
+      replace = spirv::ConstantOp::getOne(op.getType(), loc, rewriter);
+    } else {
+      // Unordered comparsion checks if either operand is NaN.
+      replace = spirv::ConstantOp::getZero(op.getType(), loc, rewriter);
+    }
+  } else {
     Value lhsIsNan = rewriter.create<spirv::IsNanOp>(loc, adaptor.getLhs());
     Value rhsIsNan = rewriter.create<spirv::IsNanOp>(loc, adaptor.getRhs());
 
-  Value replace = rewriter.create<spirv::LogicalOrOp>(loc, lhsIsNan, rhsIsNan);
+    replace = rewriter.create<spirv::LogicalOrOp>(loc, lhsIsNan, rhsIsNan);
     if (op.getPredicate() == arith::CmpFPredicate::ORD)
       replace = rewriter.create<spirv::LogicalNotOp>(loc, replace);
+  }
 
   rewriter.replaceOp(op, replace);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // AddICarryOpPattern
 //===----------------------------------------------------------------------===//
SelectOpPattern::matchAndRewrite(arith::SelectOp op, OpAdaptor adaptor,
                                  ConversionPatternRewriter &rewriter) const {
   rewriter.replaceOpWithNewOp<spirv::SelectOp>(op, adaptor.getCondition(),
                                                adaptor.getTrueValue(),
                                                adaptor.getFalseValue());
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// MaxFOpPattern
+//===----------------------------------------------------------------------===//
+
+template <typename Op, typename SPIRVOp>
+LogicalResult MinMaxFOpPattern<Op, SPIRVOp>::matchAndRewrite(
+    Op op, typename Op::Adaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  auto *converter = this->template getTypeConverter<SPIRVTypeConverter>();
+  auto dstType = converter->convertType(op.getType());
+  if (!dstType)
+    return failure();
+
+  // arith.maxf/minf:
+  //   "if one of the arguments is NaN, then the result is also NaN."
+  // spv.GL.FMax/FMin:
+  //   "which operand is the result is undefined if one of the operands
+  //   is a NaN."
+
+  Location loc = op.getLoc();
+  Value spirvOp = rewriter.create<SPIRVOp>(loc, dstType, adaptor.getOperands());
+
+  if (converter->getOptions().enableFastMathMode) {
+    rewriter.replaceOp(op, spirvOp);
+    return success();
+  }
+
+  Value lhsIsNan = rewriter.create<spirv::IsNanOp>(loc, adaptor.getLhs());
+  Value rhsIsNan = rewriter.create<spirv::IsNanOp>(loc, adaptor.getRhs());
+
+  Value select1 = rewriter.create<spirv::SelectOp>(loc, dstType, lhsIsNan,
+                                                   adaptor.getLhs(), spirvOp);
+  Value select2 = rewriter.create<spirv::SelectOp>(loc, dstType, rhsIsNan,
+                                                   adaptor.getRhs(), select1);
+
+  rewriter.replaceOp(op, select2);
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // Pattern Population
 //===----------------------------------------------------------------------===//
 
 void mlir::arith::populateArithmeticToSPIRVPatterns(
     SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
     ConstantCompositeOpPattern,
patterns.add<
     TypeCastingOpPattern<arith::SIToFPOp, spirv::ConvertSToFOp>,
     TypeCastingOpPattern<arith::FPToSIOp, spirv::ConvertFToSOp>,
     TypeCastingOpPattern<arith::IndexCastOp, spirv::SConvertOp>,
     TypeCastingOpPattern<arith::BitcastOp, spirv::BitcastOp>,
     CmpIOpBooleanPattern, CmpIOpPattern,
     CmpFOpNanNonePattern, CmpFOpPattern,
     AddICarryOpPattern, SelectOpPattern,
 
-    spirv::ElementwiseOpPattern<arith::MaxFOp, spirv::GLFMaxOp>,
+    MinMaxFOpPattern<arith::MaxFOp, spirv::GLFMaxOp>,
+    MinMaxFOpPattern<arith::MinFOp, spirv::GLFMinOp>,
     spirv::ElementwiseOpPattern<arith::MaxSIOp, spirv::GLSMaxOp>,
     spirv::ElementwiseOpPattern<arith::MaxUIOp, spirv::GLUMaxOp>,
-    spirv::ElementwiseOpPattern<arith::MinFOp, spirv::GLFMinOp>,
     spirv::ElementwiseOpPattern<arith::MinSIOp, spirv::GLSMinOp>,
     spirv::ElementwiseOpPattern<arith::MinUIOp, spirv::GLUMinOp>
   >(typeConverter, patterns.getContext());
   // clang-format on
 
   // Give CmpFOpNanKernelPattern a higher benefit so it can prevail when Kernel
   // capability is available.
   patterns.add<CmpFOpNanKernelPattern>(typeConverter, patterns.getContext(),
struct ConvertArithmeticToSPIRVPass
     : public impl::ConvertArithmeticToSPIRVBase<ConvertArithmeticToSPIRVPass> {
   void runOnOperation() override {
     Operation *op = getOperation();
     auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
     auto target = SPIRVConversionTarget::get(targetAttr);
 
     SPIRVConversionOptions options;
     options.emulateNon32BitScalarTypes = this->emulateNon32BitScalarTypes;
+    options.enableFastMathMode = this->enableFastMath;
     SPIRVTypeConverter typeConverter(targetAttr, options);
 
     // Use UnrealizedConversionCast as the bridge so that we don't need to pull
     // in patterns for other dialects.
     auto addUnrealizedCast = [](OpBuilder &builder, Type type,
                                 ValueRange inputs, Location loc) {
       auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
       return Optional<Value>(cast.getResult(0));

mlir/test/Conversion/ArithmeticToSPIRV/arithmetic-to-spirv.mlir

func.func @float32_binary_scalar(%lhs: f32, %rhs: f32) {
   // CHECK: spv.FSub %{{.*}}, %{{.*}}: f32
   %1 = arith.subf %lhs, %rhs: f32
   // CHECK: spv.FMul %{{.*}}, %{{.*}}: f32
   %2 = arith.mulf %lhs, %rhs: f32
   // CHECK: spv.FDiv %{{.*}}, %{{.*}}: f32
   %3 = arith.divf %lhs, %rhs: f32
   // CHECK: spv.FRem %{{.*}}, %{{.*}}: f32
   %4 = arith.remf %lhs, %rhs: f32
-  // CHECK: spv.GL.FMax %{{.*}}, %{{.*}}: f32
-  %5 = arith.maxf %lhs, %rhs: f32
-  // CHECK: spv.GL.FMin %{{.*}}, %{{.*}}: f32
-  %6 = arith.minf %lhs, %rhs: f32
   return
 }
 
+// CHECK-LABEL: @float32_minf_scalar
+// CHECK-SAME: %[[LHS:.+]]: f32, %[[RHS:.+]]: f32
+func.func @float32_minf_scalar(%arg0 : f32, %arg1 : f32) -> f32 {
+  // CHECK: %[[MIN:.+]] = spv.GL.FMin %arg0, %arg1 : f32
+  // CHECK: %[[LHS_NAN:.+]] = spv.IsNan %[[LHS]] : f32
+  // CHECK: %[[RHS_NAN:.+]] = spv.IsNan %[[RHS]] : f32
+  // CHECK: %[[SELECT1:.+]] = spv.Select %[[LHS_NAN]], %[[LHS]], %[[MIN]]
+  // CHECK: %[[SELECT2:.+]] = spv.Select %[[RHS_NAN]], %[[RHS]], %[[SELECT1]]
+  %0 = arith.minf %arg0, %arg1 : f32
+  // CHECK: return %[[SELECT2]]
+  return %0: f32
+}
+
+// CHECK-LABEL: @float32_maxf_scalar
+// CHECK-SAME: %[[LHS:.+]]: vector<2xf32>, %[[RHS:.+]]: vector<2xf32>
+func.func @float32_maxf_scalar(%arg0 : vector<2xf32>, %arg1 : vector<2xf32>) -> vector<2xf32> {
+  // CHECK: %[[MAX:.+]] = spv.GL.FMax %arg0, %arg1 : vector<2xf32>
+  // CHECK: %[[LHS_NAN:.+]] = spv.IsNan %[[LHS]] : vector<2xf32>
+  // CHECK: %[[RHS_NAN:.+]] = spv.IsNan %[[RHS]] : vector<2xf32>
+  // CHECK: %[[SELECT1:.+]] = spv.Select %[[LHS_NAN]], %[[LHS]], %[[MAX]]
+  // CHECK: %[[SELECT2:.+]] = spv.Select %[[RHS_NAN]], %[[RHS]], %[[SELECT1]]
+  %0 = arith.maxf %arg0, %arg1 : vector<2xf32>
+  // CHECK: return %[[SELECT2]]
+  return %0: vector<2xf32>
+}
+
 // Check int vector types.
 // CHECK-LABEL: @int_vector234
 func.func @int_vector234(%arg0: vector<2xi8>, %arg1: vector<4xi64>) {
   // CHECK: spv.SDiv %{{.*}}, %{{.*}}: vector<2xi8>
   %0 = arith.divsi %arg0, %arg0: vector<2xi8>
   // CHECK: spv.UDiv %{{.*}}, %{{.*}}: vector<4xi64>
   %1 = arith.divui %arg1, %arg1: vector<4xi64>
   return

mlir/test/Conversion/ArithmeticToSPIRV/fast-math.mlir

+// RUN: mlir-opt -split-input-file -convert-arith-to-spirv=enable-fast-math -verify-diagnostics %s | FileCheck %s
+
+module attributes {
+  spv.target_env = #spv.target_env<#spv.vce<v1.0, [], []>, #spv.resource_limits<>>
+} {
+
+// CHECK-LABEL: @cmpf_ordered
+// CHECK-SAME: %[[LHS:.+]]: f32, %[[RHS:.+]]: f32
+func.func @cmpf_ordered(%arg0 : f32, %arg1 : f32) -> i1 {
+  // CHECK: %[[T:.+]] = spv.Constant true
+  %0 = arith.cmpf ord, %arg0, %arg1 : f32
+  // CHECK: return %[[T]]
+  return %0: i1
+}
+
+// CHECK-LABEL: @cmpf_unordered
+// CHECK-SAME: %[[LHS:.+]]: vector<4xf32>, %[[RHS:.+]]: vector<4xf32>
+func.func @cmpf_unordered(%arg0 : vector<4xf32>, %arg1 : vector<4xf32>) -> vector<4xi1> {
+  // CHECK: %[[F:.+]] = spv.Constant dense<false>
+  %0 = arith.cmpf uno, %arg0, %arg1 : vector<4xf32>
+  // CHECK: return %[[F]]
+  return %0: vector<4xi1>
+}
+
+} // end module
+
+// -----
+
+module attributes {
+  spv.target_env = #spv.target_env<#spv.vce<v1.0, [Shader], []>, #spv.resource_limits<>>
+} {
+
+// CHECK-LABEL: @minf
+// CHECK-SAME: %[[LHS:.+]]: f32, %[[RHS:.+]]: f32
+func.func @minf(%arg0 : f32, %arg1 : f32) -> f32 {
+  // CHECK: %[[F:.+]] = spv.GL.FMin %[[LHS]], %[[RHS]]
+  %0 = arith.minf %arg0, %arg1 : f32
+  // CHECK: return %[[F]]
+  return %0: f32
+}
+
+// CHECK-LABEL: @maxf
+// CHECK-SAME: %[[LHS:.+]]: vector<4xf32>, %[[RHS:.+]]: vector<4xf32>
+func.func @maxf(%arg0 : vector<4xf32>, %arg1 : vector<4xf32>) -> vector<4xf32> {
+  // CHECK: %[[F:.+]] = spv.GL.FMax %[[LHS]], %[[RHS]]
+  %0 = arith.maxf %arg0, %arg1 : vector<4xf32>
+  // CHECK: return %[[F]]
+  return %0: vector<4xf32>
+}
+
+} // end module