flang-new: fix coarse grain fail

Author: ronlieb

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

@@ -1051,6 +1051,9 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
     return OMPX_MinThreadsForLowTripCount;
   }
 
+  virtual uint32_t getUseCoarseGrain() {
+    return OMPX_EnableCoarseAllocs;
+  }
   /// Whether or not to reuse blocks for high trip count loops.
   /// @see OMPX_ReuseBlocksForHighTripCount
   bool getReuseBlocksForHighTripCount() {
@@ -1209,6 +1212,9 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
   BoolEnvar OMPX_ReuseBlocksForHighTripCount =
       BoolEnvar("LIBOMPTARGET_REUSE_BLOCKS_FOR_HIGH_TRIP_COUNT", true);
 
+  /// Envar to enable coasre alocs.
+  BoolEnvar OMPX_EnableCoarseAllocs;
+
 protected:
   /// Environment variables defined by the LLVM OpenMP implementation
   /// regarding the initial number of streams and events.

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

@@ -927,6 +927,7 @@ GenericDeviceTy::GenericDeviceTy(GenericPluginTy &Plugin, int32_t DeviceId,
       OMPX_InitialNumEvents("LIBOMPTARGET_NUM_INITIAL_EVENTS", 1),
       OMPX_NumMultiDevices("LIBOMPTARGET_NUM_MULTI_DEVICES", 0),
       OMPX_EnableRuntimeAutotuning("OMPX_ENABLE_RUNTIME_AUTOTUNING", false),
+      OMPX_EnableCoarseAllocs("OMPX_ENABLE_COARSE_ALLOCS", false),
       DeviceId(DeviceId), GridValues(OMPGridValues),
       PeerAccesses(NumDevices, PeerAccessState::PENDING), PeerAccessesLock(),
       PinnedAllocs(*this), RPCServer(nullptr) {
@@ -2208,7 +2209,7 @@ void *GenericPluginTy::data_alloc(int32_t DeviceId, int64_t Size, void *HostPtr,
 
     // Method has no effect when the CUDA Plugin is used.
     // This method can only be called if HostPtr is not null.
-    if (HostPtr && Kind == TARGET_ALLOC_SHARED)
+    if (HostPtr && Kind == TARGET_ALLOC_SHARED && getDevice(DeviceId).getUseCoarseGrain() )
       set_coarse_grain_mem_region(DeviceId, HostPtr, Size);
 
     return *AllocOrErr;