[OpenMP] Apply OpenMP assumptions to applicable call sites

This patch adds OpenMP assumption attributes to call sites in applicable
regions. Currently this applies the caller's assumption attributes to
any calls contained within it. So, if a call occurs inside an OpenMP
assumes region to a function outside that region, we will assume that
call respects the assumptions. This is primarily useful for inline
assembly calls used heavily in the OpenMP GPU device runtime, which
allows us to then make judgements about what the ASM will do.

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D110655

clang/lib/CodeGen/CGCall.cpp

@@ -1747,6 +1747,26 @@ static void AddAttributesFromFunctionProtoType(ASTContext &Ctx,
     FuncAttrs.addAttribute(llvm::Attribute::NoUnwind);
 }
 
+static void AddAttributesFromAssumes(llvm::AttrBuilder &FuncAttrs,
+                                     const Decl *Callee, const Decl *Caller,
+                                     bool AssumptionOnCallSite) {
+  if (!Callee)
+    return;
+
+  SmallVector<StringRef, 4> Attrs;
+
+  for (const AssumptionAttr *AA : Callee->specific_attrs<AssumptionAttr>())
+    AA->getAssumption().split(Attrs, ",");
+
+  if (Caller && Caller->hasAttrs() && AssumptionOnCallSite)
+    for (const AssumptionAttr *AA : Caller->specific_attrs<AssumptionAttr>())
+      AA->getAssumption().split(Attrs, ",");
+
+  if (!Attrs.empty())
+    FuncAttrs.addAttribute(llvm::AssumptionAttrKey,
+                           llvm::join(Attrs.begin(), Attrs.end(), ","));
+}
+
 bool CodeGenModule::MayDropFunctionReturn(const ASTContext &Context,
                                           QualType ReturnType) {
   // We can't just discard the return value for a record type with a
@@ -1997,12 +2017,10 @@ static bool DetermineNoUndef(QualType QTy, CodeGenTypes &Types,
 ///     attributes that restrict how the frontend generates code must be
 ///     added here rather than getDefaultFunctionAttributes.
 ///
-void CodeGenModule::ConstructAttributeList(StringRef Name,
+void CodeGenModule::ConstructAttributeList(
-                                           const CGFunctionInfo &FI,
+    StringRef Name, const CGFunctionInfo &FI, CGCalleeInfo CalleeInfo,
-                                           CGCalleeInfo CalleeInfo,
+    llvm::AttributeList &AttrList, unsigned &CallingConv, bool AttrOnCallSite,
-                                           llvm::AttributeList &AttrList,
+    bool IsThunk, const Decl *Caller) {
-                                           unsigned &CallingConv,
-                                           bool AttrOnCallSite, bool IsThunk) {
   llvm::AttrBuilder FuncAttrs;
   llvm::AttrBuilder RetAttrs;
 
@@ -2020,6 +2038,13 @@ void CodeGenModule::ConstructAttributeList(StringRef Name,
 
   const Decl *TargetDecl = CalleeInfo.getCalleeDecl().getDecl();
 
+  // Only attach assumptions to call sites in OpenMP mode.
+  bool AssumptionOnCallSite = getLangOpts().OpenMP && AttrOnCallSite;
+
+  // Attach assumption attributes to the declaration. If this is a call
+  // site, attach assumptions from the caller to the call as well.
+  AddAttributesFromAssumes(FuncAttrs, TargetDecl, Caller, AssumptionOnCallSite);
+
   bool HasOptnone = false;
   // The NoBuiltinAttr attached to the target FunctionDecl.
   const NoBuiltinAttr *NBA = nullptr;
@@ -2119,18 +2144,6 @@ void CodeGenModule::ConstructAttributeList(StringRef Name,
                                llvm::toStringRef(CodeGenOpts.UniformWGSize));
       }
     }
-
-    std::string AssumptionValueStr;
-    for (AssumptionAttr *AssumptionA :
-         TargetDecl->specific_attrs<AssumptionAttr>()) {
-      std::string AS = AssumptionA->getAssumption().str();
-      if (!AS.empty() && !AssumptionValueStr.empty())
-        AssumptionValueStr += ",";
-      AssumptionValueStr += AS;
-    }
-
-    if (!AssumptionValueStr.empty())
-      FuncAttrs.addAttribute(llvm::AssumptionAttrKey, AssumptionValueStr);
   }
 
   // Attach "no-builtins" attributes to:
@@ -5165,7 +5178,7 @@ RValue CodeGenFunction::EmitCall(const CGFunctionInfo &CallInfo,
   CGM.ConstructAttributeList(CalleePtr->getName(), CallInfo,
                              Callee.getAbstractInfo(), Attrs, CallingConv,
                              /*AttrOnCallSite=*/true,
-                             /*IsThunk=*/false);
+                             /*IsThunk=*/false, CurFuncDecl);
 
   if (const FunctionDecl *FD = dyn_cast_or_null<FunctionDecl>(CurFuncDecl))
     if (FD->hasAttr<StrictFPAttr>())

clang/lib/CodeGen/CGStmt.cpp

@@ -26,6 +26,7 @@
 #include "clang/Basic/TargetInfo.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/IR/Assumptions.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/InlineAsm.h"
 #include "llvm/IR/Intrinsics.h"
@@ -2209,6 +2210,20 @@ static void UpdateAsmCallInst(llvm::CallBase &Result, bool HasSideEffect,
       Result.addFnAttr(llvm::Attribute::ReadOnly);
   }
 
+  // Attach OpenMP assumption attributes from the caller, if they exist.
+  if (CGF.CGM.getLangOpts().OpenMP) {
+    SmallVector<StringRef, 4> Attrs;
+
+    for (const AssumptionAttr *AA :
+         CGF.CurFuncDecl->specific_attrs<AssumptionAttr>())
+      AA->getAssumption().split(Attrs, ",");
+
+    if (!Attrs.empty())
+      Result.addFnAttr(
+          llvm::Attribute::get(CGF.getLLVMContext(), llvm::AssumptionAttrKey,
+                               llvm::join(Attrs.begin(), Attrs.end(), ",")));
+  }
+
   // Slap the source location of the inline asm into a !srcloc metadata on the
   // call.
   if (const auto *gccAsmStmt = dyn_cast<GCCAsmStmt>(&S))

clang/lib/CodeGen/CGVTables.cpp

@@ -428,7 +428,7 @@ void CodeGenFunction::EmitMustTailThunk(GlobalDecl GD,
   llvm::AttributeList Attrs;
   CGM.ConstructAttributeList(Callee.getCallee()->getName(), *CurFnInfo, GD,
                              Attrs, CallingConv, /*AttrOnCallSite=*/true,
-                             /*IsThunk=*/false);
+                             /*IsThunk=*/false, CurFuncDecl);
   Call->setAttributes(Attrs);
   Call->setCallingConv(static_cast<llvm::CallingConv::ID>(CallingConv));
 

clang/lib/CodeGen/CodeGenModule.h

@@ -1200,7 +1200,8 @@ public:
   void ConstructAttributeList(StringRef Name, const CGFunctionInfo &Info,
                               CGCalleeInfo CalleeInfo,
                               llvm::AttributeList &Attrs, unsigned &CallingConv,
-                              bool AttrOnCallSite, bool IsThunk);
+                              bool AttrOnCallSite, bool IsThunk,
+                              const Decl *Caller = nullptr);
 
   /// Adds attributes to F according to our CodeGenOptions and LangOptions, as
   /// though we had emitted it ourselves.  We remove any attributes on F that

clang/test/OpenMP/assumes_codegen.cpp

@@ -67,6 +67,20 @@ int lambda_outer() {
 }
 #pragma omp end assumes
 
+void no_assume() {
+  foo();
+}
+
+#pragma omp begin assumes ext_call_site
+void assume() {
+  foo();
+}
+
+void assembly() {
+  asm ("nop");
+}
+#pragma omp end assumes ext_call_site
+
 // AST:      void foo() __attribute__((assume("omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses"))) __attribute__((assume("omp_no_openmp"))) {
 // AST-NEXT: }
 // AST-NEXT: class BAR {
@@ -115,29 +129,41 @@ int lambda_outer() {
 // CHECK: define{{.*}} void @_Z3barv()
 // CHECK-SAME: [[attr1:#[0-9]]]
 // CHECK:   call{{.*}} @_ZN3BARC1Ev(%class.BAR*{{.*}} %b)
-// CHECK-SAME: [[attr9:#[0-9]]]
+// CHECK-SAME: [[attr10:#[0-9]]]
 // CHECK: define{{.*}} void @_ZN3BARC1Ev(%class.BAR*{{.*}} %this)
 // CHECK-SAME: [[attr2:#[0-9]]]
 // CHECK:   call{{.*}} @_ZN3BARC2Ev(%class.BAR*{{.*}} %this1)
-// CHECK-SAME: [[attr9]]
+// CHECK-SAME: [[attr10]]
 // CHECK: define{{.*}} void @_ZN3BARC2Ev(%class.BAR*{{.*}} %this)
 // CHECK-SAME: [[attr3:#[0-9]]]
 // CHECK: define{{.*}} void @_Z3bazv()
 // CHECK-SAME: [[attr4:#[0-9]]]
 // CHECK:   call{{.*}} @_ZN3BAZIfEC1Ev(%class.BAZ*{{.*}} %b)
-// CHECK-SAME: [[attr10:#[0-9]]]
+// CHECK-SAME: [[attr11:#[0-9]]]
 // CHECK: define{{.*}} void @_ZN3BAZIfEC1Ev(%class.BAZ*{{.*}} %this)
 // CHECK-SAME: [[attr5:#[0-9]]]
 // CHECK:   call{{.*}} @_ZN3BAZIfEC2Ev(%class.BAZ*{{.*}} %this1)
-// CHECK-SAME: [[attr10]]
+// CHECK-SAME: [[attr12:#[0-9]]]
 // CHECK: define{{.*}} void @_ZN3BAZIfEC2Ev(%class.BAZ*{{.*}} %this)
 // CHECK-SAME: [[attr6:#[0-9]]]
 // CHECK: define{{.*}} i32 @_Z12lambda_outerv()
 // CHECK-SAME: [[attr7:#[0-9]]]
 // CHECK: call{{.*}} @"_ZZ12lambda_outervENK3$_0clEv"
-// CHECK-SAME: [[attr11:#[0-9]]]
+// CHECK-SAME: [[attr13:#[0-9]]]
 // CHECK: define{{.*}} i32 @"_ZZ12lambda_outervENK3$_0clEv"(%class.anon*{{.*}} %this)
 // CHECK-SAME: [[attr8:#[0-9]]]
+// CHECK: define{{.*}} void @_Z9no_assumev()
+// CHECK-SAME: [[attr0:#[0-9]]]
+// CHECK: call{{.*}} @_Z3foov()
+// CHECK-SAME: [[attr14:#[0-9]]]
+// CHECK: define{{.*}} void @_Z6assumev()
+// CHECK-SAME: [[attr9:#[0-9]]]
+// CHECK: call{{.*}} @_Z3foov()
+// CHECK-SAME: [[attr15:#[0-9]]]
+// CHECK: define{{.*}} void @_Z8assemblyv()
+// CHECK-SAME: [[attr9:#[0-9]]]
+// CHECK: call{{.*}} void asm sideeffect "nop", "~{dirflag},~{fpsr},~{flags}"()
+// CHECK-SAME: [[attr16:#[0-9]]]
 
 // CHECK:     attributes [[attr0]]
 // CHECK-SAME:  "llvm.assume"="omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
@@ -158,8 +184,18 @@ int lambda_outer() {
 // CHECK:     attributes [[attr8]]
 // CHECK-SAME:  "llvm.assume"="ompx_lambda_assumption,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
 // CHECK:     attributes [[attr9]]
-// CHECK-SAME:  "llvm.assume"="ompx_range_bar_only,ompx_range_bar_only_2,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
+// CHECK-SAME:  "llvm.assume"="ompx_call_site,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
 // CHECK:     attributes [[attr10]]
-// CHECK-SAME:  "llvm.assume"="ompx_1234,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
+// CHECK-SAME:  "llvm.assume"="ompx_range_bar_only,ompx_range_bar_only_2,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp,ompx_range_bar_only,ompx_range_bar_only_2,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
 // CHECK:     attributes [[attr11]]
-// CHECK-SAME:  "llvm.assume"="ompx_lambda_assumption,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
+// CHECK-SAME:  "llvm.assume"="ompx_1234,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp,ompx_1234,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp,ompx_1234,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
+// CHECK:     attributes [[attr12]]
+// CHECK-SAME:  "llvm.assume"="ompx_1234,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp,ompx_1234,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
+// CHECK:     attributes [[attr13]]
+// CHECK-SAME:  "llvm.assume"="ompx_lambda_assumption,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp,ompx_lambda_assumption,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
+// CHECK:     attributes [[attr14]]
+// CHECK-SAME:  "llvm.assume"="omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
+// CHECK:     attributes [[attr15]]
+// CHECK-SAME:  "llvm.assume"="omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp,ompx_call_site,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"
+// CHECK:     attributes [[attr16]]
+// CHECK-SAME:  "llvm.assume"="ompx_call_site,omp_no_openmp_routines,ompx_another_warning,ompx_after_invalid_clauses,omp_no_openmp"

clang/test/OpenMP/assumes_include_nvptx.cpp

@@ -24,7 +24,7 @@
 // CHECK:       attributes [[attr1]]
 // CHECK-SAME:  "llvm.assume"="ompx_check_that_this_is_attached_to_included_functions_and_template_instantiations"
 // CHECK:       attributes [[attr2]]
-// CHECK-SAME:  "llvm.assume"="ompx_check_that_this_is_attached_to_included_functions_and_template_instantiations"
+// CHECK-SAME:  "llvm.assume"="ompx_check_that_this_is_attached_to_included_functions_and_template_instantiations,ompx_check_that_this_is_attached_to_included_functions_and_template_instantiations"
 
 
 template <typename T>