Support rounding mode switching for function calls

spavloff · spavloff · commit cf6539667a9f · 2024-05-13T18:30:37.000+07:00
diff --git a/clang/include/clang/Basic/TargetInfo.h b/clang/include/clang/Basic/TargetInfo.h
@@ -1232,6 +1232,14 @@ class TargetInfo : public TransferrableTargetInfo,
     return true;
   }
 
+  /// Returns true, if an operations that depends on rounding mode can be
+  /// implemented without changing FP environment. In this case the rounding
+  /// mode is encoded in the bits of implementing instruction.
+  virtual bool hasStaticRounding() const {
+    // Most supported targets require setting hardware register to use
+    // particular rounding mode.
+    return false;
+  }
   /// Returns the target triple of the primary target.
   const llvm::Triple &getTriple() const {
     return Triple;
diff --git a/clang/lib/Basic/Targets/RISCV.h b/clang/lib/Basic/Targets/RISCV.h
@@ -110,6 +110,8 @@ class RISCVTargetInfo : public TargetInfo {
 
   bool hasBFloat16Type() const override { return true; }
 
+  bool hasStaticRounding() const override { return true; }
+
   CallingConvCheckResult checkCallingConvention(CallingConv CC) const override;
 
   bool useFP16ConversionIntrinsics() const override {
diff --git a/clang/lib/CodeGen/CGCall.cpp b/clang/lib/CodeGen/CGCall.cpp
@@ -30,6 +30,7 @@
 #include "clang/CodeGen/CGFunctionInfo.h"
 #include "clang/CodeGen/SwiftCallingConv.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Assumptions.h"
 #include "llvm/IR/AttributeMask.h"
@@ -5693,6 +5694,9 @@ RValue CodeGenFunction::EmitCall(const CGFunctionInfo &CallInfo,
   AllocAlignAttrEmitter AllocAlignAttrEmitter(*this, TargetDecl, CallArgs);
   Attrs = AllocAlignAttrEmitter.TryEmitAsCallSiteAttribute(Attrs);
 
+  // Prepare execution environment.
+  setRoundingModeForCall(Callee);
+
   // Emit the actual call/invoke instruction.
   llvm::CallBase *CI;
   if (!InvokeDest) {
@@ -5846,6 +5850,9 @@ RValue CodeGenFunction::EmitCall(const CGFunctionInfo &CallInfo,
   // lexical order, so deactivate it and run it manually here.
   CallArgs.freeArgumentMemory(*this);
 
+  // Restore execution environment.
+  restoreRoundingModeAfterCall();
+
   // Extract the return value.
   RValue Ret = [&] {
     switch (RetAI.getKind()) {
@@ -5980,6 +5987,64 @@ RValue CodeGenFunction::EmitCall(const CGFunctionInfo &CallInfo,
   return Ret;
 }
 
+static bool endsWithRoundingModeSuffix(StringRef FuncName) {
+  size_t Underscore = FuncName.find_last_of("_");
+  if (Underscore == StringRef::npos || Underscore < 2)
+    return false;
+  StringRef Suffix = FuncName.substr(Underscore + 1);
+  static const StringRef RMSuffixes[] = {"rtz", "rte", "rtp", "rtn", "rhaz",
+                                         "rz",  "rn",  "ru",  "rd"};
+  for (auto RM : RMSuffixes) {
+    if (Suffix == RM)
+      return true;
+  }
+  return false;
+}
+
+bool CodeGenFunction::requiresDynamicRounding(const CGCallee &Callee) {
+  if (Callee.isOrdinary()) {
+    const Decl *CalleeDecl = Callee.getAbstractInfo().getCalleeDecl().getDecl();
+    if (const FunctionDecl *FD = dyn_cast_or_null<FunctionDecl>(CalleeDecl)) {
+      IdentifierInfo *FuncNameII = FD->getDeclName().getAsIdentifierInfo();
+      if (FuncNameII) {
+        StringRef FuncName = FuncNameII->getName();
+        // If a reserved identifier ends with rounding mode suffix preceded by
+        // underscore, this function does not need the previous dynamic rounding
+        // mode to be set.
+        if (isReservedInAllContexts(
+                FuncNameII->isReserved(getContext().getLangOpts()))) {
+          if (endsWithRoundingModeSuffix(FuncName))
+            return false;
+        }
+      }
+    }
+  }
+  return true;
+}
+
+/// Sets dynamic rounding mode for the function called in the region where
+/// pragma FENV_ROUND is in effect.
+void CodeGenFunction::setRoundingModeForCall(const CGCallee &Callee) {
+  if (Target.hasStaticRounding() || Callee.isBuiltin() ||
+      !requiresDynamicRounding(Callee))
+    return;
+  if (!CurrentRoundingIsStatic || !DynamicRoundingMode)
+    return;
+  Builder.CreateCall(CGM.getIntrinsic(llvm::Intrinsic::set_rounding),
+                     DynamicRoundingMode);
+  CurrentRoundingIsStatic = false;
+}
+
+void CodeGenFunction::restoreRoundingModeAfterCall() {
+  if (Target.hasStaticRounding() || CurFPFeatures.isRoundingModeDynamic())
+    return;
+  if (CurrentRoundingIsStatic || !StaticRoundingMode)
+    return;
+  Builder.CreateCall(CGM.getIntrinsic(llvm::Intrinsic::set_rounding),
+                     StaticRoundingMode);
+  CurrentRoundingIsStatic = true;
+}
+
 CGCallee CGCallee::prepareConcreteCallee(CodeGenFunction &CGF) const {
   if (isVirtual()) {
     const CallExpr *CE = getVirtualCallExpr();
diff --git a/clang/lib/CodeGen/CGStmt.cpp b/clang/lib/CodeGen/CGStmt.cpp
@@ -499,41 +499,47 @@ struct FPControlModesCleanup final : EHScopeStack::Cleanup {
 };
 } // namespace
 
-void CodeGenFunction::emitSetFPControlModes(FPOptions NewFP) {
-  if (NewFP == CurFPFeatures)
+void CodeGenFunction::emitSetFPControlModes(FPOptions NewFP, FPOptions OldFP) {
+  if (NewFP == OldFP)
     return;
 
   // For now only rounding mode is handled.
 
+  if (Target.hasStaticRounding())
+    return;
+
   // If the new rounding mode is unknown in compile-time, it means that the
   // compound statement contains `#pragma STDC FENV_ACCESS ON`. In this case all
   // manipulations on FP environment, including setting and restoring control
   // modes are made by the user.
   if (NewFP.isRoundingModeDynamic())
     return;
 
-  llvm::RoundingMode OldConstRM = CurFPFeatures.getConstRoundingMode();
+  llvm::RoundingMode OldConstRM = OldFP.getConstRoundingMode();
   llvm::RoundingMode NewConstRM = NewFP.getConstRoundingMode();
   if (OldConstRM == NewConstRM)
     return;
 
-  llvm::RoundingMode OldRM = CurFPFeatures.getRoundingMode();
+  llvm::RoundingMode OldRM = OldFP.getRoundingMode();
   if (OldRM == NewConstRM)
     return;
 
-  llvm::Value *PreviousRM = nullptr;
-  if (CurFPFeatures.isRoundingModeDynamic()) {
+  if (OldFP.isRoundingModeDynamic()) {
     llvm::Function *FGetRound = CGM.getIntrinsic(llvm::Intrinsic::get_rounding);
-    PreviousRM = Builder.CreateCall(FGetRound);
+    DynamicRoundingMode = Builder.CreateCall(FGetRound);
   } else {
-    PreviousRM = llvm::ConstantInt::get(Int32Ty, static_cast<uint64_t>(OldRM));
+    DynamicRoundingMode =
+        llvm::ConstantInt::get(Int32Ty, static_cast<uint64_t>(OldRM));
   }
 
   llvm::RoundingMode NewRM = NewFP.getRoundingMode();
-  Builder.CreateIntrinsic(
-      llvm::Intrinsic::set_rounding, {},
-      llvm::ConstantInt::get(Int32Ty, static_cast<uint64_t>(NewRM)));
-  EHStack.pushCleanup<FPControlModesCleanup>(NormalAndEHCleanup, PreviousRM);
+  StaticRoundingMode =
+      llvm::ConstantInt::get(Int32Ty, static_cast<uint64_t>(NewRM));
+  Builder.CreateIntrinsic(llvm::Intrinsic::set_rounding, {},
+                          StaticRoundingMode);
+  EHStack.pushCleanup<FPControlModesCleanup>(NormalAndEHCleanup,
+                                             DynamicRoundingMode);
+  CurrentRoundingIsStatic = true;
 }
 
 /// EmitCompoundStmt - Emit a compound statement {..} node.  If GetLast is true,
@@ -562,8 +568,8 @@ CodeGenFunction::EmitCompoundStmtWithoutScope(const CompoundStmt &S,
   // Optionally set up the new FP environment, if the compound statement
   // contains a pragma that modifies it.
   FPOptions NewFP = S.getNewFPOptions(CurFPFeatures);
-  emitSetFPControlModes(NewFP);
   CGFPOptionsRAII SavedFPFeatues(*this, NewFP);
+  emitSetFPControlModes(NewFP, SavedFPFeatues.getOldFPOptions());
 
   Address RetAlloca = Address::invalid();
 
diff --git a/clang/lib/CodeGen/CodeGenFunction.cpp b/clang/lib/CodeGen/CodeGenFunction.cpp
@@ -147,6 +147,9 @@ CodeGenFunction::CGFPOptionsRAII::CGFPOptionsRAII(CodeGenFunction &CGF,
 void CodeGenFunction::CGFPOptionsRAII::ConstructorHelper(FPOptions FPFeatures) {
   OldFPFeatures = CGF.CurFPFeatures;
   CGF.CurFPFeatures = FPFeatures;
+  OldDynamicRM = CGF.DynamicRoundingMode;
+  OldStaticRM = CGF.StaticRoundingMode;
+  OldCurrentRoundingIsStatic = CGF.CurrentRoundingIsStatic;
 
   OldExcept = CGF.Builder.getDefaultConstrainedExcept();
   OldRounding = CGF.Builder.getDefaultConstrainedRounding();
@@ -191,6 +194,9 @@ void CodeGenFunction::CGFPOptionsRAII::ConstructorHelper(FPOptions FPFeatures) {
 
 CodeGenFunction::CGFPOptionsRAII::~CGFPOptionsRAII() {
   CGF.CurFPFeatures = OldFPFeatures;
+  CGF.DynamicRoundingMode = OldDynamicRM;
+  CGF.StaticRoundingMode = OldStaticRM;
+  CGF.CurrentRoundingIsStatic = OldCurrentRoundingIsStatic;
   CGF.Builder.setDefaultConstrainedExcept(OldExcept);
   CGF.Builder.setDefaultConstrainedRounding(OldRounding);
 }
diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
@@ -817,15 +817,23 @@ class CodeGenFunction : public CodeGenTypeCache {
     CGFPOptionsRAII(CodeGenFunction &CGF, const Expr *E);
     ~CGFPOptionsRAII();
 
+    FPOptions getOldFPOptions() const { return OldFPFeatures; }
+
   private:
     void ConstructorHelper(FPOptions FPFeatures);
     CodeGenFunction &CGF;
     FPOptions OldFPFeatures;
+    llvm::Value *OldDynamicRM;
+    llvm::Value *OldStaticRM;
+    bool OldCurrentRoundingIsStatic;
     llvm::fp::ExceptionBehavior OldExcept;
     llvm::RoundingMode OldRounding;
     std::optional<CGBuilderTy::FastMathFlagGuard> FMFGuard;
   };
   FPOptions CurFPFeatures;
+  llvm::Value *DynamicRoundingMode = nullptr;
+  llvm::Value *StaticRoundingMode = nullptr;
+  bool CurrentRoundingIsStatic = false;
 
 public:
   /// ObjCEHValueStack - Stack of Objective-C exception values, used for
@@ -3326,7 +3334,12 @@ class CodeGenFunction : public CodeGenTypeCache {
 
   /// Optionally emit code that sets required floating-point control modes and
   /// creates corresponding cleanup action.
-  void emitSetFPControlModes(FPOptions NewFP);
+  void emitSetFPControlModes(FPOptions NewFP, FPOptions OldFP);
+
+  void setRoundingModeForCall(const CGCallee &Callee);
+  void restoreRoundingModeAfterCall();
+  bool requiresDynamicRounding(const CGCallee &Callee);
+
 
   //===--------------------------------------------------------------------===//
   //                            Declaration Emission
diff --git a/clang/test/CodeGen/complex-strictfp.c b/clang/test/CodeGen/complex-strictfp.c
@@ -61,7 +61,9 @@ void test3b(void) {
 // CHECK-NEXT:    [[CF_IMAG:%.*]] = load float, ptr getelementptr inbounds ({ float, float }, ptr @cf, i32 0, i32 1), align 4
 // CHECK-NEXT:    [[CONV:%.*]] = call double @llvm.experimental.constrained.fpext.f64.f32(float [[CF_REAL]], metadata !"fpexcept.strict") #[[ATTR3]]
 // CHECK-NEXT:    [[CONV1:%.*]] = call double @llvm.experimental.constrained.fpext.f64.f32(float [[CF_IMAG]], metadata !"fpexcept.strict") #[[ATTR3]]
+// CHECK-NEXT:    call void @llvm.set.rounding(i32 1) #[[ATTR3]]
 // CHECK-NEXT:    [[CALL:%.*]] = call { double, double } @__divdc3(double noundef [[CONV]], double noundef [[CONV1]], double noundef [[G1_REAL]], double noundef [[G1_IMAG]]) #[[ATTR4:[0-9]+]]
+// CHECK-NEXT:    call void @llvm.set.rounding(i32 2) #[[ATTR3]]
 // CHECK-NEXT:    [[TMP0:%.*]] = extractvalue { double, double } [[CALL]], 0
 // CHECK-NEXT:    [[TMP1:%.*]] = extractvalue { double, double } [[CALL]], 1
 // CHECK-NEXT:    [[CONV2:%.*]] = call float @llvm.experimental.constrained.fptrunc.f32.f64(double [[TMP0]], metadata !"round.upward", metadata !"fpexcept.strict") #[[ATTR3]]
diff --git a/clang/test/CodeGen/pragma-fenv_round.c b/clang/test/CodeGen/pragma-fenv_round.c
