[SPIRV] Add FloatControl2 capability (#144371)

Add handling for FPFastMathMode in SPIR-V shaders. This is a first pass
that
simply does a direct translation when the proper extension is available.
This will unblock work for HLSL. However, it is not a full solution.

The default math mode for spir-v is determined by the API. When
targeting Vulkan many of the fast math options are assumed. We should do
something particular when targeting Vulkan.

We will also need to handle the hlsl "precise" keyword correctly when
FPFastMathMode is not available.

Unblockes https://github.com/llvm/llvm-project/issues/140739, but we are
keeing it open to track the remaining issues mentioned above.

llvm/docs/SPIRVUsage.rst

@@ -217,6 +217,8 @@ list of supported SPIR-V extensions, sorted alphabetically by their extension na
      - Adds an instruction to compute the matrix product of an M x K matrix with a K x N matrix and then add an M x N matrix. 
    * - ``SPV_INTEL_int4``
      - Adds support for 4-bit integer type, and allow this type to be used in cooperative matrices.
+   * - ``SPV_KHR_float_controls2``
+     - Adds ability to specify the floating-point environment in shaders. It can be used on whole modules and individual instructions.
 
 To enable multiple extensions, list them separated by comma. For example, to enable support for atomic operations on floating-point numbers and arbitrary precision integers, use:
 

llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp

@@ -100,7 +100,9 @@ static const std::map<std::string, SPIRV::Extension::Extension, std::less<>>
          SPIRV::Extension::Extension::SPV_INTEL_ternary_bitwise_function},
         {"SPV_INTEL_2d_block_io",
          SPIRV::Extension::Extension::SPV_INTEL_2d_block_io},
-        {"SPV_INTEL_int4", SPIRV::Extension::Extension::SPV_INTEL_int4}};
+        {"SPV_INTEL_int4", SPIRV::Extension::Extension::SPV_INTEL_int4},
+        {"SPV_KHR_float_controls2",
+         SPIRV::Extension::Extension::SPV_KHR_float_controls2}};
 
 bool SPIRVExtensionsParser::parse(cl::Option &O, StringRef ArgName,
                                   StringRef ArgValue,

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

@@ -70,6 +70,8 @@ getSymbolicOperandRequirements(SPIRV::OperandCategory::OperandCategory Category,
   AvoidCapabilitiesSet AvoidCaps;
   if (!ST.isShader())
     AvoidCaps.S.insert(SPIRV::Capability::Shader);
+  else
+    AvoidCaps.S.insert(SPIRV::Capability::Kernel);
 
   VersionTuple ReqMinVer = getSymbolicOperandMinVersion(Category, i);
   VersionTuple ReqMaxVer = getSymbolicOperandMaxVersion(Category, i);
@@ -88,8 +90,11 @@ getSymbolicOperandRequirements(SPIRV::OperandCategory::OperandCategory Category,
   } else if (MinVerOK && MaxVerOK) {
     if (ReqCaps.size() == 1) {
       auto Cap = ReqCaps[0];
-      if (Reqs.isCapabilityAvailable(Cap))
+      if (Reqs.isCapabilityAvailable(Cap)) {
+        ReqExts.append(getSymbolicOperandExtensions(
+            SPIRV::OperandCategory::CapabilityOperand, Cap));
         return {true, {Cap}, ReqExts, ReqMinVer, ReqMaxVer};
+      }
     } else {
       // By SPIR-V specification: "If an instruction, enumerant, or other
       // feature specifies multiple enabling capabilities, only one such
@@ -103,8 +108,11 @@ getSymbolicOperandRequirements(SPIRV::OperandCategory::OperandCategory Category,
           UseCaps.push_back(Cap);
       for (size_t i = 0, Sz = UseCaps.size(); i < Sz; ++i) {
         auto Cap = UseCaps[i];
-        if (i == Sz - 1 || !AvoidCaps.S.contains(Cap))
+        if (i == Sz - 1 || !AvoidCaps.S.contains(Cap)) {
+          ReqExts.append(getSymbolicOperandExtensions(
+              SPIRV::OperandCategory::CapabilityOperand, Cap));
           return {true, {Cap}, ReqExts, ReqMinVer, ReqMaxVer};
+        }
       }
     }
   }
@@ -1975,6 +1983,14 @@ static unsigned getFastMathFlags(const MachineInstr &I) {
   return Flags;
 }
 
+static bool isFastMathMathModeAvailable(const SPIRVSubtarget &ST) {
+  if (ST.isKernel())
+    return true;
+  if (ST.getSPIRVVersion() < VersionTuple(1, 2))
+    return false;
+  return ST.canUseExtension(SPIRV::Extension::SPV_KHR_float_controls2);
+}
+
 static void handleMIFlagDecoration(MachineInstr &I, const SPIRVSubtarget &ST,
                                    const SPIRVInstrInfo &TII,
                                    SPIRV::RequirementHandler &Reqs) {
@@ -1998,8 +2014,12 @@ static void handleMIFlagDecoration(MachineInstr &I, const SPIRVSubtarget &ST,
   unsigned FMFlags = getFastMathFlags(I);
   if (FMFlags == SPIRV::FPFastMathMode::None)
     return;
-  Register DstReg = I.getOperand(0).getReg();
-  buildOpDecorate(DstReg, I, TII, SPIRV::Decoration::FPFastMathMode, {FMFlags});
+
+  if (isFastMathMathModeAvailable(ST)) {
+    Register DstReg = I.getOperand(0).getReg();
+    buildOpDecorate(DstReg, I, TII, SPIRV::Decoration::FPFastMathMode,
+                    {FMFlags});
+  }
 }
 
 // Walk all functions and add decorations related to MI flags.

llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td

@@ -319,6 +319,7 @@ defm SPV_INTEL_ternary_bitwise_function : ExtensionOperand<120>;
 defm SPV_INTEL_subgroup_matrix_multiply_accumulate : ExtensionOperand<121>;
 defm SPV_INTEL_2d_block_io : ExtensionOperand<122>;
 defm SPV_INTEL_int4 : ExtensionOperand<123>;
+defm SPV_KHR_float_controls2 : ExtensionOperand<124>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define Capabilities enum values and at the same time
@@ -489,6 +490,8 @@ defm DotProductInput4x8Bit : CapabilityOperand<6017, 0x10600, 0, [SPV_KHR_intege
 defm DotProductInput4x8BitPacked : CapabilityOperand<6018, 0x10600, 0, [SPV_KHR_integer_dot_product], []>;
 defm DotProduct : CapabilityOperand<6019, 0x10600, 0, [SPV_KHR_integer_dot_product], []>;
 defm GroupNonUniformRotateKHR : CapabilityOperand<6026, 0, 0, [SPV_KHR_subgroup_rotate], [GroupNonUniform]>;
+defm FloatControls2
+    : CapabilityOperand<6029, 0x10200, 0, [SPV_KHR_float_controls2], []>;
 defm AtomicFloat32AddEXT : CapabilityOperand<6033, 0, 0, [SPV_EXT_shader_atomic_float_add], []>;
 defm AtomicFloat64AddEXT : CapabilityOperand<6034, 0, 0, [SPV_EXT_shader_atomic_float_add], []>;
 defm AtomicFloat16AddEXT : CapabilityOperand<6095, 0, 0, [SPV_EXT_shader_atomic_float16_add], []>;
@@ -1239,7 +1242,7 @@ defm XfbBuffer : DecorationOperand<36, 0, 0, [], [TransformFeedback]>;
 defm XfbStride : DecorationOperand<37, 0, 0, [], [TransformFeedback]>;
 defm FuncParamAttr : DecorationOperand<38, 0, 0, [], [Kernel]>;
 defm FPRoundingMode : DecorationOperand<39, 0, 0, [], []>;
-defm FPFastMathMode : DecorationOperand<40, 0, 0, [], [Kernel]>;
+defm FPFastMathMode : DecorationOperand<40, 0, 0, [], [Kernel, FloatControls2]>;
 defm LinkageAttributes : DecorationOperand<41, 0, 0, [], [Linkage]>;
 defm NoContraction : DecorationOperand<42, 0, 0, [], [Shader]>;
 defm InputAttachmentIndex : DecorationOperand<43, 0, 0, [], [InputAttachment]>;

llvm/test/CodeGen/SPIRV/capability-FloatControl2.ll

@@ -0,0 +1,22 @@
+; RUN: llc -O0 -mtriple=spirv1.6-vulkan1.3-compute %s -o - | FileCheck %s --check-prefix=CHECK-NOEXT
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv1.6-vulkan1.3-compute %s -o - -filetype=obj | spirv-val --target-env vulkan1.3 %}
+
+; RUN: llc -O0 -mtriple=spirv1.6-vulkan1.3-compute -spirv-ext=+SPV_KHR_float_controls2 %s -o - | FileCheck %s --check-prefix=CHECK-EXT
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv1.6-vulkan1.3-compute -spirv-ext=+SPV_KHR_float_controls2 %s -o - -filetype=obj | spirv-val --target-env vulkan1.3 %}
+
+; CHECK-NOEXT-NOT: OpDecorate FPFastMathMode
+
+; CHECK-EXT: OpCapability FloatControls2
+; CHECK-EXT: OpExtension "SPV_KHR_float_controls2"
+; CHECK-EXT: OpDecorate {{%[0-9]+}} FPFastMathMode NotNaN|NotInf|NSZ|AllowRecip|Fast
+
+define hidden spir_func float @foo(float  %0) local_unnamed_addr {
+  %2 = fmul reassoc nnan ninf nsz arcp afn float %0, 2.000000e+00
+  ret float %2
+}
+
+define void @main() local_unnamed_addr #1 {
+  ret void
+}
+
+attributes #1 = { "hlsl.numthreads"="8,1,1" "hlsl.shader"="compute" }