Implement extension SPV_KHR_float_controls2

include/LLVMSPIRVExtensions.inc

@@ -89,3 +89,4 @@ EXT(SPV_INTEL_predicated_io)
 EXT(SPV_INTEL_sigmoid)
 EXT(SPV_INTEL_float4)
 EXT(SPV_INTEL_fp_conversions)
+EXT(SPV_KHR_float_controls2)

lib/SPIRV/SPIRVMDWalker.h

@@ -91,7 +91,9 @@ class SPIRVMDWalker {
 
     bool atEnd() const { return !(M && I < E); }
 
-    template <typename T> MDWrapper &get(T &V) {
+    template <typename IntegerType,
+              typename = std::enable_if_t<std::is_integral_v<IntegerType>>>
+    MDWrapper &get(IntegerType &V) {
       if (!Q)
         assert(I < E && "out of bound");
       if (atEnd())
@@ -115,12 +117,14 @@ class SPIRVMDWalker {
       return *this;
     }
 
-    MDWrapper &get(Function *&F) {
+    template <typename ValueType,
+              typename = std::enable_if_t<std::is_base_of_v<Value, ValueType>>>
+    MDWrapper &get(ValueType *&V) {
       if (!Q)
         assert(I < E && "out of bound");
       if (atEnd())
         return *this;
-      F = mdconst::dyn_extract<Function>(M->getOperand(I++));
+      V = mdconst::dyn_extract<ValueType>(M->getOperand(I++));
       return *this;
     }
 

lib/SPIRV/SPIRVReader.cpp

@@ -1218,26 +1218,24 @@ static void applyNoIntegerWrapDecorations(const SPIRVValue *BV,
   }
 }
 
-static void applyFPFastMathModeDecorations(const SPIRVValue *BV,
-                                           Instruction *Inst) {
-  SPIRVWord V;
-  FastMathFlags FMF;
+void SPIRVToLLVM::applyFPFastMathModeDecorations(const SPIRVValue *BV,
+                                                 Instruction *Inst) {
+  if (!isa<FPMathOperator>(Inst))
+    return;
+
+  SPIRVWord V{0};
   if (BV->hasDecorate(DecorationFPFastMathMode, 0, &V)) {
-    if (V & FPFastMathModeNotNaNMask)
-      FMF.setNoNaNs();
-    if (V & FPFastMathModeNotInfMask)
-      FMF.setNoInfs();
-    if (V & FPFastMathModeNSZMask)
-      FMF.setNoSignedZeros();
-    if (V & FPFastMathModeAllowRecipMask)
-      FMF.setAllowReciprocal();
-    if (V & FPFastMathModeAllowContractFastINTELMask)
-      FMF.setAllowContract();
-    if (V & FPFastMathModeAllowReassocINTELMask)
-      FMF.setAllowReassoc();
-    if (V & FPFastMathModeFastMask)
-      FMF.setFast();
+    FastMathFlags FMF = translateFastMathFlags(V);
     Inst->setFastMathFlags(FMF);
+    return;
+  }
+
+  // Get the scalar type to handle vector operands. And get the first operand
+  // type (instead of the result) due to fcmp instructions.
+  Type *FloatType = Inst->getOperand(0)->getType()->getScalarType();
+  auto Func2FMF = FuncToFastMathFlags.find({Inst->getFunction(), FloatType});
+  if (Func2FMF != FuncToFastMathFlags.end()) {
+    Inst->setFastMathFlags(Func2FMF->second);
   }
 }
 
@@ -3443,6 +3441,73 @@ static void validatePhiPredecessors(Function *F) {
 }
 } // namespace
 
+FastMathFlags SPIRVToLLVM::translateFastMathFlags(SPIRVWord V) const {
+  FastMathFlags FMF;
+  if (V & FPFastMathModeNotNaNMask)
+    FMF.setNoNaNs();
+  if (V & FPFastMathModeNotInfMask)
+    FMF.setNoInfs();
+  if (V & FPFastMathModeNSZMask)
+    FMF.setNoSignedZeros();
+  if (V & FPFastMathModeAllowRecipMask)
+    FMF.setAllowReciprocal();
+  static_assert(FPFastMathModeAllowContractFastINTELMask ==
+                FPFastMathModeAllowContractMask);
+  if (V & FPFastMathModeAllowContractFastINTELMask)
+    FMF.setAllowContract();
+  static_assert(FPFastMathModeAllowReassocINTELMask ==
+                FPFastMathModeAllowReassocMask);
+  if (V & FPFastMathModeAllowReassocINTELMask)
+    FMF.setAllowReassoc();
+  // There is no FPFastMathMode flag that represents LLVM approximate functions
+  // flag `afn`. Even the FPFastMathMode Fast flag should not imply it, but to
+  // avoid changing the previous behaviour we make it equivalent to LLVM's.
+  if (V & FPFastMathModeFastMask)
+    FMF.setFast();
+  if (V & FPFastMathModeAllowTransformMask) {
+    // AllowTransform requires the AllowContract and AllowReassoc bits to be
+    // set.
+    assert(FMF.allowContract() && FMF.allowReassoc() &&
+           "The FPFastMathMode AllowTransform requires AllowContract and "
+           "AllowReassoc to be set");
+  }
+
+  return FMF;
+}
+
+void SPIRVToLLVM::parseFloatControls2ExecutionModeId(SPIRVFunction *BF,
+                                                     Function *F) {
+
+  auto [Begin, End] =
+      BF->getExecutionModeRange(spv::ExecutionModeFPFastMathDefault);
+  if (Begin == End)
+    return;
+
+  LLVMContext &C = F->getContext();
+  NamedMDNode *ExecModeMD =
+      M->getOrInsertNamedMetadata(kSPIRVMD::ExecutionMode);
+
+  Metadata *FPFastMathMode[4] = {ConstantAsMetadata::get(F),
+                                 ConstantAsMetadata::get(getUInt32(
+                                     M, spv::ExecutionModeFPFastMathDefault)),
+                                 nullptr, nullptr};
+
+  for (auto [_, EM] : make_range(Begin, End)) {
+    const auto &Literals = EM->getLiterals();
+    assert(Literals.size() == 2);
+    SPIRVWord FloatTyId = Literals[0];
+    SPIRVType *FloatSPIRVType = BM->get<SPIRVType>(FloatTyId);
+    Type *FloatType = transFPType(FloatSPIRVType);
+    SPIRVWord Flags = *transIdAsConstant(Literals[1]);
+    FuncToFastMathFlags.try_emplace({F, FloatType},
+                                    translateFastMathFlags(Flags));
+
+    FPFastMathMode[2] = ConstantAsMetadata::get(PoisonValue::get(FloatType));
+    FPFastMathMode[3] = ConstantAsMetadata::get(getUInt32(M, Flags));
+    ExecModeMD->addOperand(MDNode::get(C, FPFastMathMode));
+  }
+}
+
 Function *SPIRVToLLVM::transFunction(SPIRVFunction *BF, unsigned AS) {
   auto Loc = FuncMap.find(BF);
   if (Loc != FuncMap.end())
@@ -3520,6 +3585,8 @@ Function *SPIRVToLLVM::transFunction(SPIRVFunction *BF, unsigned AS) {
                              : CallingConv::SPIR_FUNC);
   transFunctionAttrs(BF, F);
 
+  parseFloatControls2ExecutionModeId(BF, F);
+
   // Creating all basic blocks before creating instructions.
   for (size_t I = 0, E = BF->getNumBasicBlock(); I != E; ++I) {
     transValue(BF->getBasicBlock(I), F, nullptr);

lib/SPIRV/SPIRVReader.h

@@ -274,6 +274,14 @@ class SPIRVToLLVM : private BuiltinCallHelper {
   void
   transFunctionPointerCallArgumentAttributes(SPIRVValue *BV, CallInst *CI,
                                              SPIRVTypeFunction *CalledFnTy);
+
+  using FunctionAndTypeIdPair = std::pair<Function *, Type *>;
+  using FunctionToFastMathFlagsMap =
+      DenseMap<FunctionAndTypeIdPair, FastMathFlags>;
+  FunctionToFastMathFlagsMap FuncToFastMathFlags;
+  FastMathFlags translateFastMathFlags(SPIRVWord V) const;
+  void parseFloatControls2ExecutionModeId(SPIRVFunction *BF, Function *F);
+  void applyFPFastMathModeDecorations(const SPIRVValue *BV, Instruction *Inst);
 }; // class SPIRVToLLVM
 
 } // namespace SPIRV