[OpenMP] Provide a flag to disable safety checks for GPU optimizations

[jdoerfert]

Summary

The optimization introduced in D83271 will verify a safety and a
precautionary condition. Until we find a proper solution for the former
we allow users to disable the check. If they do it and violate the
problematic (but probably unlikely) case, they will see UB, most likely
a crash when a non-function pointer is called. The flag has the word
unsafe in the name and the remarks are adjusted to provide some
insight what is happening.

[JonChesterfield]

I think there's an unfortunate interaction with link time optimisation here. If there are external regions, but their code is combined with llvm-link before codegen, then a user could reasonably assume this flag is safe.

Would it would be correct to compile the individual source assuming there may be external uses, then llvm-link the source, then run the pass assuming there are no external uses? That's of interest to the amdgcn case as we can (currently) assume the whole program is available towards the end of the compilation.

[jdoerfert]

In D83832#2152027, @JonChesterfield wrote:

I think there's an unfortunate interaction with link time optimisation here. If there are external regions, but their code is combined with llvm-link before codegen, then a user could reasonably assume this flag is safe.

If you set this flag and reach a parallel region from a target region outside the translation unit, you will experience problems as calling a non-function pointer is likely not to work. If you don't do this, you are all good.

Would it would be correct to compile the individual source assuming there may be external uses, then llvm-link the source, then run the pass assuming there are no external uses? That's of interest to the amdgcn case as we can (currently) assume the whole program is available towards the end of the compilation.

If you run O2 after linking *and* properly set the linkage of all your non-kernel functions to internal at that point, you don't need the flag. All the safe cases will then be handled automatically, all the unsafe ones rejected. At least that is what I suspect will happen.

[JonChesterfield]

I think this is good. It's dangerous, but it's also undocumented and has unsafe in the name.

I should be able to use this to sidestep limitations in the amdgpu function pointer implementation.

clang/test/OpenMP/remarks_parallel_in_target_state_machine.c

-// RUN: %clang_cc1                                 -verify=host -Rpass=openmp -fopenmp -x c++ -triple powerpc64le-unknown-unknown -fopenmp-targets=nvptx64-nvidia-cuda -emit-llvm-bc %s -o %t-ppc-host.bc
-// RUN: %clang_cc1                                 -verify      -Rpass=openmp -fopenmp -O2 -x c++ -triple nvptx64-unknown-unknown -fopenmp-targets=nvptx64-nvidia-cuda -emit-llvm %s -fopenmp-is-device -fopenmp-host-ir-file-path %t-ppc-host.bc -o %t.out
-// RUN: %clang_cc1 -fexperimental-new-pass-manager -verify      -Rpass=openmp -fopenmp -O2 -x c++ -triple nvptx64-unknown-unknown -fopenmp-targets=nvptx64-nvidia-cuda -emit-llvm %s -fopenmp-is-device -fopenmp-host-ir-file-path %t-ppc-host.bc -o %t.out
+// RUN: %clang_cc1                                 -verify=host                                                              -Rpass=openmp -fopenmp -x c++ -triple powerpc64le-unknown-unknown -fopenmp-targets=nvptx64-nvidia-cuda -emit-llvm-bc %s -o %t-ppc-host.bc
+// RUN: %clang_cc1                                 -verify=all,safe                                                          -Rpass=openmp -fopenmp -O2 -x c++ -triple nvptx64-unknown-unknown -fopenmp-targets=nvptx64-nvidia-cuda -emit-llvm %s -fopenmp-is-device -fopenmp-host-ir-file-path %t-ppc-host.bc -o %t.out
+// RUN: %clang_cc1 -fexperimental-new-pass-manager -verify=all,safe                                                          -Rpass=openmp -fopenmp -O2 -x c++ -triple nvptx64-unknown-unknown -fopenmp-targets=nvptx64-nvidia-cuda -emit-llvm %s -fopenmp-is-device -fopenmp-host-ir-file-path %t-ppc-host.bc -o %t.out
+// RUN: %clang_cc1                                 -verify=all,force -mllvm -openmp-unsafe-assume-no-external-target-regions -Rpass=openmp -fopenmp -O2 -x c++ -triple nvptx64-unknown-unknown -fopenmp-targets=nvptx64-nvidia-cuda -emit-llvm %s -fopenmp-is-device -fopenmp-host-ir-file-path %t-ppc-host.bc -o %t.out
+// RUN: %clang_cc1 -fexperimental-new-pass-manager -verify=all,force -mllvm -openmp-unsafe-assume-no-external-target-regions -Rpass=openmp -fopenmp -O2 -x c++ -triple nvptx64-unknown-unknown -fopenmp-targets=nvptx64-nvidia-cuda -emit-llvm %s -fopenmp-is-device -fopenmp-host-ir-file-path %t-ppc-host.bc -o %t.out
 
 // host-no-diagnostics
 
 void bar(void) {
 #pragma omp parallel // #1                                                                                                                                                                                                                                                                                                                                           \
-                     // expected-remark@#1 {{Found a parallel region that is called in a target region but not part of a combined target construct nor nesed inside a target construct without intermediate code. This can lead to excessive register usage for unrelated target regions in the same translation unit due to spurious call edges assumed by ptxas.}} \
-                     // expected-remark@#1 {{Parallel region is not known to be called from a unique single target region, maybe the surrounding function has external linkage?; will not attempt to rewrite the state machine use.}}
+                     // all-remark@#1 {{Found a parallel region that is called in a target region but not part of a combined target construct nor nesed inside a target construct without intermediate code. This can lead to excessive register usage for unrelated target regions in the same translation unit due to spurious call edges assumed by ptxas.}} \
+                     // safe-remark@#1 {{Parallel region is not known to be called from a unique single target region, maybe the surrounding function has external linkage?; will not attempt to rewrite the state machine use.}}         \
+                     // force-remark@#1 {{[UNSAFE] Parallel region is not known to be called from a unique single target region, maybe the surrounding function has external linkage?; will rewrite the state machine use due to command line flag, this can lead to undefined behavior if the parallel region is called from a target region outside this translation unit.}} \
+                     // force-remark@#1 {{Specialize parallel region that is only reached from a single target region to avoid spurious call edges and excessive register usage in other target regions. (parallel region ID: __omp_outlined__2_wrapper, kernel ID: <NONE>}}
   {
   }
 }
 
 void foo(void) {
 #pragma omp target teams // #2                                                                                                                                                                      \
-                         // expected-remark@#2 {{Target region containing the parallel region that is specialized. (parallel region ID: __omp_outlined__1_wrapper, kernel ID: __omp_offloading_22}} \
-                         // expected-remark@#2 {{Target region containing the parallel region that is specialized. (parallel region ID: __omp_outlined__3_wrapper, kernel ID: __omp_offloading_22}}
+                         // all-remark@#2 {{Target region containing the parallel region that is specialized. (parallel region ID: __omp_outlined__1_wrapper, kernel ID: __omp_offloading_22}} \
+                         // all-remark@#2 {{Target region containing the parallel region that is specialized. (parallel region ID: __omp_outlined__3_wrapper, kernel ID: __omp_offloading_22}}
   {
 #pragma omp parallel // #3                                                                                                                                                                                                                                                                                                                                           \
-                     // expected-remark@#3 {{Found a parallel region that is called in a target region but not part of a combined target construct nor nesed inside a target construct without intermediate code. This can lead to excessive register usage for unrelated target regions in the same translation unit due to spurious call edges assumed by ptxas.}} \
-                     // expected-remark@#3 {{Specialize parallel region that is only reached from a single target region to avoid spurious call edges and excessive register usage in other target regions. (parallel region ID: __omp_outlined__1_wrapper, kernel ID: __omp_offloading_22}}
+                     // all-remark@#3 {{Found a parallel region that is called in a target region but not part of a combined target construct nor nesed inside a target construct without intermediate code. This can lead to excessive register usage for unrelated target regions in the same translation unit due to spurious call edges assumed by ptxas.}} \
+                     // all-remark@#3 {{Specialize parallel region that is only reached from a single target region to avoid spurious call edges and excessive register usage in other target regions. (parallel region ID: __omp_outlined__1_wrapper, kernel ID: __omp_offloading_22}}
     {
     }
     bar();
 #pragma omp parallel // #4                                                                                                                                                                                                                                                                                                                                           \
-                     // expected-remark@#4 {{Found a parallel region that is called in a target region but not part of a combined target construct nor nesed inside a target construct without intermediate code. This can lead to excessive register usage for unrelated target regions in the same translation unit due to spurious call edges assumed by ptxas.}} \
-                     // expected-remark@#4 {{Specialize parallel region that is only reached from a single target region to avoid spurious call edges and excessive register usage in other target regions. (parallel region ID: __omp_outlined__3_wrapper, kernel ID: __omp_offloading_22}}
+                     // all-remark@#4 {{Found a parallel region that is called in a target region but not part of a combined target construct nor nesed inside a target construct without intermediate code. This can lead to excessive register usage for unrelated target regions in the same translation unit due to spurious call edges assumed by ptxas.}} \
+                     // all-remark@#4 {{Specialize parallel region that is only reached from a single target region to avoid spurious call edges and excessive register usage in other target regions. (parallel region ID: __omp_outlined__3_wrapper, kernel ID: __omp_offloading_22}}
     {
     }
   }
 }
 
 void spmd(void) {
   // Verify we do not emit the remarks above for "SPMD" regions.
 #pragma omp target teams
 #pragma omp parallel
   {
   }
 
 #pragma omp target teams distribute parallel for
   for (int i = 0; i < 100; ++i) {
   }
 }
 
-// expected-remark@* {{OpenMP runtime call __kmpc_global_thread_num moved to}}
-// expected-remark@* {{OpenMP runtime call __kmpc_global_thread_num deduplicated}}
+// all-remark@* {{OpenMP runtime call __kmpc_global_thread_num moved to}}
+// all-remark@* {{OpenMP runtime call __kmpc_global_thread_num deduplicated}}

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 using namespace omp;
 
 #define DEBUG_TYPE "openmp-opt"
 
 static cl::opt<bool> DisableOpenMPOptimizations(
     "openmp-opt-disable", cl::ZeroOrMore,
     cl::desc("Disable OpenMP specific optimizations."), cl::Hidden,
     cl::init(false));
+static cl::opt<bool> UnsafeAssumeNoExternalTargetRegions(
+    "openmp-unsafe-assume-no-external-target-regions", cl::ZeroOrMore,
+    cl::desc("[UNSAFE!] Assume no external target regions when specializing "
+             "OpenMP parallel regions to avoid spurious register usage issues, "
+             "combine with `-Rpass=openmp`"),
+    cl::Hidden, cl::init(false));
 
 static cl::opt<bool> PrintICVValues("openmp-print-icv-values", cl::init(false),
                                     cl::Hidden);
 static cl::opt<bool> PrintOpenMPKernels("openmp-print-gpu-kernels",
                                         cl::init(false), cl::Hidden);
 
 STATISTIC(NumOpenMPRuntimeCallsDeduplicated,
           "Number of OpenMP runtime calls deduplicated");
for (Function *F : SCC) {
     }
 
     // Even if we have __kmpc_kernel_prepare_parallel calls, we (for now) give
     // up if the function is not called from a unique kernel.
     Kernel K = getUniqueKernelFor(*F);
     if (!K) {
       {
         auto Remark = [&](OptimizationRemark OR) {
-          return OR << "Parallel region is not known to be called from a "
+          return OR << (UnsafeAssumeNoExternalTargetRegions ? "[UNSAFE] " : "")
+                    << "Parallel region is not known to be called from a "
                        "unique single target region, maybe the surrounding "
-                       "function has external linkage?; will not attempt to "
-                       "rewrite the state machine use.";
+                       "function has external linkage?"
+                    << (UnsafeAssumeNoExternalTargetRegions
+                            ? "; will rewrite the state machine use due to "
+                              "command line flag, this can lead to undefined "
+                              "behavior if the parallel region is called from "
+                              "a target region outside this translation unit."
+                            : "; will not attempt to rewrite the state machine "
+                              "use.");
         };
         emitRemarkOnFunction(F, "OpenMPParallelRegionInMultipleKernesl",
                              Remark);
       }
+      if (!UnsafeAssumeNoExternalTargetRegions)
         continue;
     }
 
     // We now know F is a parallel body function called only from the kernel K.
     // We also identified the state machine uses in which we replace the
     // function pointer by a new global symbol for identification purposes. This
     // ensures only direct calls to the function are left.
 
     {
       auto RemarkParalleRegion = [&](OptimizationRemark OR) {
         return OR << "Specialize parallel region that is only reached from a "
                      "single target region to avoid spurious call edges and "
                      "excessive register usage in other target regions. "
                      "(parallel region ID: "
                   << ore::NV("OpenMPParallelRegion", F->getName())
                   << ", kernel ID: "
-                  << ore::NV("OpenMPTargetRegion", K->getName()) << ")";
+                  << ore::NV("OpenMPTargetRegion", K ? K->getName() : "<NONE>")
+                  << ")";
       };
       emitRemarkOnFunction(F, "OpenMPParallelRegionInNonSPMD",
                            RemarkParalleRegion);
+
       auto RemarkKernel = [&](OptimizationRemark OR) {
         return OR << "Target region containing the parallel region that is "
                      "specialized. (parallel region ID: "
                   << ore::NV("OpenMPParallelRegion", F->getName())
                   << ", kernel ID: "
                   << ore::NV("OpenMPTargetRegion", K->getName()) << ")";
       };
+      if (K)
         emitRemarkOnFunction(K, "OpenMPParallelRegionInNonSPMD", RemarkKernel);
     }
 
     Module &M = *F->getParent();
     Type *Int8Ty = Type::getInt8Ty(M.getContext());
 
     auto *ID = new GlobalVariable(
         M, Int8Ty, /* isConstant */ true, GlobalValue::PrivateLinkage,
         UndefValue::get(Int8Ty), F->getName() + ".ID");