[OpenMP] Introduce an assumption to ignore possible external callers

[jdoerfert]

Summary

In this patch we utilize the new assumption function attributes to
improve the OpenMPOpt behavior in the presence of externally visible
functions containing parallel regions when we compile for the GPU.

The existing OpenMP target state machine optimization will replace the
use of parallel region function pointers with dedicated handles if all
uses of the parallel region function pointer can be replaced. External
callers that are inside of OpenMP target regions will require the
function pointer to be used as there will be an indirect call. We now
allow the user to assume that external callers do not exist, or that
they might exist but not call from within an OpenMP (target) region.
The assumption in the IR is "omp_no_external_caller_in_target_region"
which can be created with OpenMP assumes

omp begin assumes ext_omp_no_external_caller_in_target_region
void foo() {
  #pragma omp parallel
  { ... }
}
omp end assumes

or the generic clang assume attribute

void foo() __attribute__((assume("omp_no_external_caller_in_target_region"))) {
  #pragma omp parallel
  { ... }
}

[jdoerfert]

Add remark and documentation

[jdoerfert]

Add remark summary

[JonChesterfield]

Is this for cases where we are compiling a subset of the target code, i.e. without link time optimisation?

It's interesting that we might want a static function on the gpu and an external one on the cpu. The user could presumably make it static and provide a different, host-only function which calls it as a workaround.

I'm nervous about this one because it looks like a property the compiler should already be able to derive (given LTO, anyway), and I don't see how we can diagnose when the user annotation is inaccurate, e.g. because of changes elsewhere in the codebase after it was added.

[jdoerfert]

In D93079#2451782, @JonChesterfield wrote:

Is this for cases where we are compiling a subset of the target code, i.e. without link time optimisation?

Upstream cannot to LTO on any target code as of now :(

It's interesting that we might want a static function on the gpu and an external one on the cpu. The user could presumably make it static and provide a different, host-only function which calls it as a workaround.

If they have control. Code in a header for example might not be changeable. Though I will add it to the remark description for sure.

I'm nervous about this one because it looks like a property the compiler should already be able to derive (given LTO, anyway), and I don't see how we can diagnose when the user annotation is inaccurate, e.g. because of changes elsewhere in the codebase after it was added.

This is always an issue with assumes/assertions/annotations/... Unsure what to say. I'd be fine with retiring this one as soon as we can do LTO upstream, and I am all in favor of making that the only/default mode then too.

[jdoerfert]

Rebase on top of D93439

[JonChesterfield]

Ah, this may be one of the things aomp changed for nvptx. We did something around linking but I've never looked into exactly what. In that case fine by me - it's a sharp edge, but I'd expect it to make a difference in benchmarks until whole program optimisation is online.

I think the strategy we want to converge on is equivalent to:

N openmp source files -> bitcode -> llvm-link -> mlink-builtin-bitcode=nvptx.devicertl -> opt -> ptxas

which interacts poorly with the compile/link separation. So presumably we should blend in the devicertl in an lld plugin, as that's when users expect the whole program to be available.

clang/include/clang/Basic/AttrDocs.td

 While LLVM plugins might provide more assumption strings, the default LLVM
 optimization passes are aware of the following assumptions:
 
   .. code-block:: none
 
     "omp_no_openmp"
     "omp_no_openmp_routines"
     "omp_no_parallelism"
+    "omp_no_external_caller_in_target_region"
 
-The OpenMP standard defines the meaning of OpenMP assumptions ("omp_XYZ" is
-spelled "XYZ" in the `OpenMP 5.1 Standard`_).
+The OpenMP standard defines the meaning of the first three OpenMP assumptions
+("omp_XYZ" is spelled "XYZ" in the `OpenMP 5.1 Standard`_) while other are
+clang extensions spelled described `here
+<http://openmp.llvm.org/docs/remarks/OptimizationRemarks.html>`_ and spelled
+without the `omp` in the OpenMP assume directive syntax.
 
 .. _`OpenMP 5.1 Standard`: https://www.openmp.org/spec-html/5.1/openmpsu37.html#x56-560002.5.2
 
 }];
 }
 
 def NoStackProtectorDocs : Documentation {
   let Category = DocCatFunction;

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 #include "llvm/ADT/EnumeratedArray.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
+#include "llvm/IR/Assumptions.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/Attributor.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
 
 using namespace llvm;
for (int Idx = 0; Idx < OMPInfoCache.ICVs.size() - 1; ++Idx) {
         return false;
       };
 
       GetterRFI.foreachUse(SCC, CreateAA);
     }
   }
 };
 
+/// The "omp_no_external_caller_in_target_region" assumption guarantees that
+/// there are no external caller of a function which are inside an OpenMP
+/// target region.
+static KnownAssumptionString
+    NoExternalCallerInTargetRegion("omp_no_external_caller_in_target_region");
+
 Kernel OpenMPOpt::getUniqueKernelFor(Function &F) {
   if (!OMPInfoCache.ModuleSlice.count(&F))
     return nullptr;
 
   // Use a scope to keep the lifetime of the CachedKernel short.
   {
     Optional<Kernel> &CachedKernel = UniqueKernelMap[&F];
     if (CachedKernel)
       return *CachedKernel;
 
     // TODO: We should use an AA to create an (optimistic and callback
     //       call-aware) call graph. For now we stick to simple patterns that
     //       are less powerful, basically the worst fixpoint.
     if (isKernel(F)) {
       CachedKernel = Kernel(&F);
       return *CachedKernel;
     }
 
     CachedKernel = nullptr;
-    if (!F.hasLocalLinkage()) {
+    if (!F.hasLocalLinkage() &&
+        !hasAssumption(F, NoExternalCallerInTargetRegion)) {
 
       // See https://openmp.llvm.org/remarks/OptimizationRemarks.html
       auto Remark = [&](OptimizationRemark OR) {
         return OR << "[OMP100] Potentially unknown OpenMP target region caller";
       };
       emitRemarkOnFunction(&F, "OMP100", Remark);
 
       return nullptr;

llvm/test/Transforms/OpenMP/gpu_state_machine_function_ptr_replacement.ll

 ; RUN: opt -S -passes=openmpopt -pass-remarks=openmp-opt -openmp-print-gpu-kernels < %s | FileCheck %s
 ; RUN: opt -S        -openmpopt -pass-remarks=openmp-opt -openmp-print-gpu-kernels < %s | FileCheck %s
 
 ; C input used for this test:
 
 ; void bar(void) {
 ;     #pragma omp parallel
 ;     { }
 ; }
+; __attribute__((assume("no_external_callers")))
+; void baz(void) {
+;     #pragma omp parallel
+;     { }
+; }
 ; void foo(void) {
 ;   #pragma omp target teams
 ;   {
 ;     #pragma omp parallel
 ;     {}
 ;     bar();
+;     baz();
 ;     #pragma omp parallel
 ;     {}
 ;   }
 ; }
 
 ; Verify we replace the function pointer uses for the first and last outlined
 ; region (1 and 3) but not for the middle one (2) because it could be called from
 ; another kernel.
 
 ; CHECK-DAG: @__omp_outlined__1_wrapper.ID = private constant i8 undef
+; CHECK-DAG: @__omp_outlined__2b_wrapper.ID = private constant i8 undef
 ; CHECK-DAG: @__omp_outlined__3_wrapper.ID = private constant i8 undef
 
 ; CHECK-DAG:   icmp eq i8* %5, @__omp_outlined__1_wrapper.ID
+; CHECK-DAG:   icmp eq i8* %b6, @__omp_outlined__2b_wrapper.ID
 ; CHECK-DAG:   icmp eq i8* %7, @__omp_outlined__3_wrapper.ID
 
 ; CHECK-DAG:   call void @__kmpc_kernel_prepare_parallel(i8* @__omp_outlined__1_wrapper.ID)
-; CHECK-DAG:   call void @__kmpc_kernel_prepare_parallel(i8* bitcast (void ()* @__omp_outlined__2_wrapper to i8*))
+; CHECK-DAG:   call void @__kmpc_kernel_prepare_parallel(i8* bitcast (void ()* @__omp_outlined__2a_wrapper to i8*))
+; CHECK-DAG:   call void @__kmpc_kernel_prepare_parallel(i8* @__omp_outlined__2b_wrapper.ID)
 ; CHECK-DAG:   call void @__kmpc_kernel_prepare_parallel(i8* @__omp_outlined__3_wrapper.ID)
 
 
 %struct.ident_t = type { i32, i32, i32, i32, i8* }
 
 define internal void @__omp_offloading_35_a1e179_foo_l7_worker() {
 entry:
   %work_fn = alloca i8*, align 8
.execute.parallel:                                ; preds = %.select.workers
   br i1 %work_match, label %.execute.fn, label %.check.next
 
 .execute.fn:                                      ; preds = %.execute.parallel
   call void @__omp_outlined__1_wrapper()
   br label %.terminate.parallel
 
 .check.next:                                      ; preds = %.execute.parallel
   %6 = load i8*, i8** %work_fn, align 8
-  %work_match1 = icmp eq i8* %6, bitcast (void ()* @__omp_outlined__2_wrapper to i8*)
-  br i1 %work_match1, label %.execute.fn2, label %.check.next3
+  %work_match1 = icmp eq i8* %6, bitcast (void ()* @__omp_outlined__2a_wrapper to i8*)
+  br i1 %work_match1, label %.execute.fn2a, label %.check.next2
+
+.execute.fn2a:                                     ; preds = %.check.next
+  call void @__omp_outlined__2a_wrapper()
+  br label %.terminate.parallel
+
+.check.next2:                                      ; preds = %.execute.parallel
+  %b6 = load i8*, i8** %work_fn, align 8
+  %work_match1b = icmp eq i8* %b6, bitcast (void ()* @__omp_outlined__2b_wrapper to i8*)
+  br i1 %work_match1b, label %.execute.fn2b, label %.check.next3
 
-.execute.fn2:                                     ; preds = %.check.next
-  call void @__omp_outlined__2_wrapper()
+.execute.fn2b:                                     ; preds = %.check.next
+  call void @__omp_outlined__2b_wrapper()
   br label %.terminate.parallel
 
 .check.next3:                                     ; preds = %.check.next
   %7 = load i8*, i8** %work_fn, align 8
   %work_match4 = icmp eq i8* %7, bitcast (void ()* @__omp_outlined__3_wrapper to i8*)
   br i1 %work_match4, label %.execute.fn5, label %.check.next6
 
 .execute.fn5:                                     ; preds = %.check.next3
define weak void @__omp_offloading_35_a1e179_foo_l7() {
   call void @__omp_offloading_35_a1e179_foo_l7_worker()
   call void @__omp_outlined__()
   ret void
 }
 
 define internal void @__omp_outlined__() {
   call void @__kmpc_kernel_prepare_parallel(i8* bitcast (void ()* @__omp_outlined__1_wrapper to i8*))
   call void @bar()
+  call void @baz()
   call void @__kmpc_kernel_prepare_parallel(i8* bitcast (void ()* @__omp_outlined__3_wrapper to i8*))
   ret void
 }
 
 define internal void @__omp_outlined__1() {
   ret void
 }
 
 define internal void @__omp_outlined__1_wrapper() {
   call void @__omp_outlined__1()
   ret void
 }
 
 define hidden void @bar() {
-  call void @__kmpc_kernel_prepare_parallel(i8* bitcast (void ()* @__omp_outlined__2_wrapper to i8*))
+  call void @__kmpc_kernel_prepare_parallel(i8* bitcast (void ()* @__omp_outlined__2a_wrapper to i8*))
+  ret void
+}
+
+define hidden void @baz() #0 {
+  call void @__kmpc_kernel_prepare_parallel(i8* bitcast (void ()* @__omp_outlined__2b_wrapper to i8*))
+  ret void
+}
+
+define internal void @__omp_outlined__2a_wrapper() {
   ret void
 }
 
-define internal void @__omp_outlined__2_wrapper() {
+define internal void @__omp_outlined__2b_wrapper() {
   ret void
 }
 
 define internal void @__omp_outlined__3_wrapper() {
   ret void
 }
 
 declare void @__kmpc_kernel_prepare_parallel(i8* %WorkFn)
 
 declare zeroext i1 @__kmpc_kernel_parallel(i8** nocapture %WorkFn)
 
 declare void @__kmpc_kernel_end_parallel()
 
 declare void @__kmpc_barrier_simple_spmd(%struct.ident_t* nocapture readnone %loc_ref, i32 %tid)
 
 declare i32 @__kmpc_global_thread_num(%struct.ident_t* nocapture readnone)
 
+attributes #0 = { "llvm.assume"="abc,omp_no_external_caller_in_target_region,123" }
 
 !nvvm.annotations = !{!0}
 
 !0 = !{void ()* @__omp_offloading_35_a1e179_foo_l7, !"kernel", i32 1}

openmp/docs/remarks/OptimizationRemarks.rst

 regions that reach the parallel region we can avoid this function pointer
 passing scheme and often improve the register usage on the GPU. However, If a
 parallel region on the GPU is in a function with external linkage we may not
 know all callers statically. If there are outside callers within target
 regions, this remark is to be ignored. If there are no such callers, users can
 modify the linkage and thereby help optimization with a `static` or
 `__attribute__((internal))` function annotation. If changing the linkage is
 impossible, e.g., because there are outside callers on the host, one can split
-the function into an external visible interface which is not compiled for
-the target and an internal implementation which is compiled for the target
-and should be called from within the target region.
+the function into an external visible interface which is not compiled for the
+target and an internal implementation which is compiled for the target and
+should be called from within the target region. Finally,  we provide an
+assumption for the target side which can be spelled either as
+`__attribute__((assume("omp_no_external_caller_in_target_region")))` or as
+OpenMP assumption, i.a., `#pragma omp begin assume
+ext_no_external_caller_in_target_region`. The assumption effectively allows the
+compiler to assume no caller outside the current translation unit will call the
+function from a target region, hence the function is not called from outside
+translation units on the device.