Add support for null fallback

Author: kevinsala

offload/include/Shared/APITypes.h

@@ -102,9 +102,8 @@ struct KernelArgsTy {
   struct {
     uint64_t NoWait : 1; // Was this kernel spawned with a `nowait` clause.
     uint64_t IsCUDA : 1; // Was this kernel spawned via CUDA.
-    uint64_t AllowDynCGroupMemFallback : 1; // Allow fallback for dynamic cgroup
-                                            // mem fallback.
-    uint64_t Unused : 61;
+    uint64_t DynCGroupMemFallback : 2; // The fallback for dynamic cgroup mem.
+    uint64_t Unused : 60;
   } Flags = {0, 0, 0, 0};
   // The number of teams (for x,y,z dimension).
   uint32_t NumTeams[3] = {0, 0, 0};

offload/include/Shared/Environment.h

@@ -92,12 +92,25 @@ struct KernelEnvironmentTy {
   DynamicEnvironmentTy *DynamicEnv = nullptr;
 };
 
+/// The fallback types for the dynamic cgroup memory.
+enum class DynCGroupMemFallbackType : unsigned char {
+  /// None. Used for indicating that no fallback was triggered.
+  None = 0,
+  /// Abort the execution.
+  Abort = None,
+  /// Return null pointer.
+  Null = 1,
+  /// Allocate from a implementation defined memory space.
+  DefaultMem = 2
+};
+
 struct KernelLaunchEnvironmentTy {
   void *ReductionBuffer = nullptr;
-  void *DynCGroupMemFallback = nullptr;
+  void *DynCGroupMemFbPtr = nullptr;
   uint32_t ReductionCnt = 0;
   uint32_t ReductionIterCnt = 0;
   uint32_t DynCGroupMemSize = 0;
+  DynCGroupMemFallbackType DynCGroupMemFb = DynCGroupMemFallbackType::None;
 };
 
 #endif // OMPTARGET_SHARED_ENVIRONMENT_H

offload/plugins-nextgen/common/include/PluginInterface.h

@@ -392,7 +392,8 @@ struct GenericKernelTy {
   /// Return a device pointer to a new kernel launch environment.
   Expected<KernelLaunchEnvironmentTy *> getKernelLaunchEnvironment(
       GenericDeviceTy &GenericDevice, const KernelArgsTy &KernelArgs,
-      void *FallbackBlockMem, AsyncInfoWrapperTy &AsyncInfo) const;
+      uint32_t BlockMemSize, DynCGroupMemFallbackType DynBlockMemFb,
+      void *DynBlockMemFbPtr, AsyncInfoWrapperTy &AsyncInfoWrapper) const;
 
   /// Indicate whether an execution mode is valid.
   static bool isValidExecutionMode(OMPTgtExecModeFlags ExecutionMode) {

offload/plugins-nextgen/common/src/PluginInterface.cpp

@@ -437,7 +437,8 @@ Error GenericKernelTy::init(GenericDeviceTy &GenericDevice,
 Expected<KernelLaunchEnvironmentTy *>
 GenericKernelTy::getKernelLaunchEnvironment(
     GenericDeviceTy &GenericDevice, const KernelArgsTy &KernelArgs,
-    void *FallbackBlockMem, AsyncInfoWrapperTy &AsyncInfoWrapper) const {
+    uint32_t BlockMemSize, DynCGroupMemFallbackType DynBlockMemFb,
+    void *DynBlockMemFbPtr, AsyncInfoWrapperTy &AsyncInfoWrapper) const {
   // Ctor/Dtor have no arguments, replaying uses the original kernel launch
   // environment. Older versions of the compiler do not generate a kernel
   // launch environment.
@@ -479,8 +480,9 @@ GenericKernelTy::getKernelLaunchEnvironment(
     LocalKLE.ReductionBuffer = nullptr;
   }
 
-  LocalKLE.DynCGroupMemSize = KernelArgs.DynCGroupMem;
-  LocalKLE.DynCGroupMemFallback = FallbackBlockMem;
+  LocalKLE.DynCGroupMemSize = BlockMemSize;
+  LocalKLE.DynCGroupMemFbPtr = DynBlockMemFbPtr;
+  LocalKLE.DynCGroupMemFb = DynBlockMemFb;
 
   INFO(OMP_INFOTYPE_DATA_TRANSFER, GenericDevice.getDeviceId(),
        "Copying data from host to device, HstPtr=" DPxMOD ", TgtPtr=" DPxMOD
@@ -539,28 +541,43 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
   if (StaticBlockMemSize > MaxBlockMemSize)
     return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                          "Static block memory size exceeds maximum");
-  else if (!KernelArgs.Flags.AllowDynCGroupMemFallback &&
+  else if (static_cast<DynCGroupMemFallbackType>(
+               KernelArgs.Flags.DynCGroupMemFallback) ==
+               DynCGroupMemFallbackType::Abort &&
            TotalBlockMemSize > MaxBlockMemSize)
     return Plugin::error(
         ErrorCode::INVALID_ARGUMENT,
         "Static and dynamic block memory size exceeds maximum");
 
-  void *FallbackBlockMem = nullptr;
+  void *DynBlockMemFbPtr = nullptr;
+  uint32_t DynBlockMemLaunchSize = DynBlockMemSize;
+
+  DynCGroupMemFallbackType DynBlockMemFb = DynCGroupMemFallbackType::None;
   if (DynBlockMemSize && (!GenericDevice.hasNativeBlockSharedMem() ||
                           TotalBlockMemSize > MaxBlockMemSize)) {
-    auto AllocOrErr = GenericDevice.dataAlloc(
-        NumBlocks[0] * DynBlockMemSize,
-        /*HostPtr=*/nullptr, TargetAllocTy::TARGET_ALLOC_DEVICE);
-    if (!AllocOrErr)
-      return AllocOrErr.takeError();
+    // Launch without native dynamic block memory.
+    DynBlockMemLaunchSize = 0;
+    DynBlockMemFb = static_cast<DynCGroupMemFallbackType>(
+        KernelArgs.Flags.DynCGroupMemFallback);
+    if (DynBlockMemFb == DynCGroupMemFallbackType::DefaultMem) {
+      // Get global memory as fallback.
+      auto AllocOrErr = GenericDevice.dataAlloc(
+          NumBlocks[0] * DynBlockMemSize,
+          /*HostPtr=*/nullptr, TargetAllocTy::TARGET_ALLOC_DEVICE);
+      if (!AllocOrErr)
+        return AllocOrErr.takeError();
 
-    FallbackBlockMem = *AllocOrErr;
-    AsyncInfoWrapper.freeAllocationAfterSynchronization(FallbackBlockMem);
-    DynBlockMemSize = 0;
+      DynBlockMemFbPtr = *AllocOrErr;
+      AsyncInfoWrapper.freeAllocationAfterSynchronization(DynBlockMemFbPtr);
+    } else {
+      // Do not provide any memory as fallback.
+      DynBlockMemSize = 0;
+    }
   }
 
   auto KernelLaunchEnvOrErr = getKernelLaunchEnvironment(
-      GenericDevice, KernelArgs, FallbackBlockMem, AsyncInfoWrapper);
+      GenericDevice, KernelArgs, DynBlockMemSize, DynBlockMemFb,
+      DynBlockMemFbPtr, AsyncInfoWrapper);
   if (!KernelLaunchEnvOrErr)
     return KernelLaunchEnvOrErr.takeError();
 
@@ -591,7 +608,7 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
           printLaunchInfo(GenericDevice, KernelArgs, NumThreads, NumBlocks))
     return Err;
 
-  return launchImpl(GenericDevice, NumThreads, NumBlocks, DynBlockMemSize,
+  return launchImpl(GenericDevice, NumThreads, NumBlocks, DynBlockMemLaunchSize,
                     KernelArgs, LaunchParams, AsyncInfoWrapper);
 }
 

offload/test/offloading/dyn_groupprivate_strict.cpp renamed to offload/test/offloading/dyn_groupprivate.cpp

@@ -1,4 +1,7 @@
-// RUN: %libomptarget-compilexx-run-and-check-generic
+// RUN: %libomptarget-compilexx-generic -fopenmp-version=61
+// RUN: %libomptarget-run-generic | %fcheck-generic
+// RUN: %libomptarget-compileoptxx-generic -fopenmp-version=61
+// RUN: %libomptarget-run-generic | %fcheck-generic
 // REQUIRES: gpu
 
 #include <omp.h>
@@ -9,8 +12,9 @@
 int main() {
   int Result[N], NumThreads;
 
+// Verify the groupprivate buffer works as expected.
 #pragma omp target teams num_teams(1) thread_limit(N)                          \
-    dyn_groupprivate(strict : N * sizeof(Result[0]))                           \
+    dyn_groupprivate(fallback(abort) : N * sizeof(Result[0]))                  \
     map(from : Result, NumThreads)
   {
     int Buffer[N];
@@ -51,8 +55,8 @@ int main() {
   size_t MaxSize = omp_get_groupprivate_limit(0, omp_access_cgroup);
   size_t ExceededSize = MaxSize + 10;
 
-// Verify that the fallback modifier works.
-#pragma omp target dyn_groupprivate(fallback : ExceededSize)                   \
+// Verify that the fallback(default_mem) modifier works.
+#pragma omp target dyn_groupprivate(fallback(default_mem) : ExceededSize)      \
     map(tofrom : Failed)
   {
     int IsFallback;
@@ -66,13 +70,35 @@ int main() {
       ++Failed;
   }
 
-// Verify that the default modifier is fallback.
+// Verify that the fallback(null) modifier works.
+#pragma omp target dyn_groupprivate(fallback(null) : ExceededSize)             \
+    map(tofrom : Failed)
+  {
+    int IsFallback;
+    if ((TmpPtr = omp_get_dyn_groupprivate_ptr(0, &IsFallback)))
+      ++Failed;
+    if ((TmpSize = omp_get_dyn_groupprivate_size()))
+      ++Failed;
+    if (!IsFallback)
+      ++Failed;
+  }
+
+// Verify that the default modifier is fallback(default_mem).
 #pragma omp target dyn_groupprivate(ExceededSize)
   {
+    int IsFallback;
+    if (!omp_get_dyn_groupprivate_ptr(0, &IsFallback))
+      ++Failed;
+    if (!omp_get_dyn_groupprivate_size())
+      ++Failed;
+    if (omp_get_dyn_groupprivate_size() != ExceededSize)
+      ++Failed;
+    if (!IsFallback)
+      ++Failed;
   }
 
-// Verify that the strict modifier works.
-#pragma omp target dyn_groupprivate(strict : N) map(tofrom : Failed)
+// Verify that the fallback(abort) modifier works.
+#pragma omp target dyn_groupprivate(fallback(abort) : N) map(tofrom : Failed)
   {
     int IsFallback;
     if (!omp_get_dyn_groupprivate_ptr(0, &IsFallback))
@@ -85,8 +111,9 @@ int main() {
       ++Failed;
   }
 
-// Verify that the fallback does not trigger when not needed.
-#pragma omp target dyn_groupprivate(fallback : N) map(tofrom : Failed)
+// Verify that the fallback(default_mem) does not trigger when not needed.
+#pragma omp target dyn_groupprivate(fallback(default_mem) : N)                 \
+    map(tofrom : Failed)
   {
     int IsFallback;
     if (!omp_get_dyn_groupprivate_ptr(0, &IsFallback))
@@ -100,7 +127,7 @@ int main() {
   }
 
 // Verify that the clause works when passing a zero size.
-#pragma omp target dyn_groupprivate(strict : 0) map(tofrom : Failed)
+#pragma omp target dyn_groupprivate(fallback(abort) : 0) map(tofrom : Failed)
   {
     int IsFallback;
     if (omp_get_dyn_groupprivate_ptr(0, &IsFallback))
@@ -112,7 +139,8 @@ int main() {
   }
 
 // Verify that the clause works when passing a zero size.
-#pragma omp target dyn_groupprivate(fallback : 0) map(tofrom : Failed)
+#pragma omp target dyn_groupprivate(fallback(default_mem) : 0)                 \
+    map(tofrom : Failed)
   {
     int IsFallback;
     if (omp_get_dyn_groupprivate_ptr(0, &IsFallback))

openmp/device/src/State.cpp

@@ -162,27 +162,27 @@ struct DynCGroupMemTy {
   void init(KernelLaunchEnvironmentTy *KLE, void *NativeDynCGroup) {
     Size = 0;
     Ptr = nullptr;
-    IsFallback = false;
+    Fallback = DynCGroupMemFallbackType::None;
     if (!KLE)
       return;
 
     Size = KLE->DynCGroupMemSize;
-    if (void *Fallback = KLE->DynCGroupMemFallback) {
-      Ptr = static_cast<char *>(Fallback) + Size * omp_get_team_num();
-      IsFallback = true;
-    } else {
+    Fallback = KLE->DynCGroupMemFb;
+    if (Fallback == DynCGroupMemFallbackType::None)
       Ptr = static_cast<char *>(NativeDynCGroup);
-    }
+    else if (Fallback == DynCGroupMemFallbackType::DefaultMem)
+      Ptr = static_cast<char *>(KLE->DynCGroupMemFbPtr) +
+            Size * omp_get_team_num();
   }
 
   char *getPtr(size_t Offset) const { return Ptr + Offset; }
-  bool isFallback() const { return IsFallback; }
+  bool isFallback() const { return Fallback != DynCGroupMemFallbackType::None; }
   size_t getSize() const { return Size; }
 
 private:
   char *Ptr;
   size_t Size;
-  bool IsFallback;
+  DynCGroupMemFallbackType Fallback;
 };
 
 [[clang::loader_uninitialized]] static Local<DynCGroupMemTy> DynCGroupMem;