[OPENMP] fix D35162 to remove changes to plugins/exports

[gregrodgers]

Ignore this, I was just trying to fix https://reviews.llvm.org/D135162 and accidentally created a new revision

openmp/libomptarget/include/device.h

 struct DeviceTy {
   int32_t DeviceID;
   RTLInfoTy *RTL;
   int32_t RTLDeviceID;
 
   bool IsInit;
   std::once_flag InitFlag;
   bool HasPendingGlobals;
+  /// The physical number of team processors. For cuda, this is number of SMs,
+  /// for AMD, this is number of CUs. Field used by ompx_get_team_procs(devid).
+  int32_t TeamProcs;
 
   /// Host data to device map type with a wrapper key indirection that allows
   /// concurrent modification of the entries without invalidating the underlying
   /// entries.
   using HostDataToTargetListTy =
       std::set<HostDataToTargetMapKeyTy, std::less<>>;
 
   /// The HDTTMap is a protected object that can only be accessed by one thread
struct DeviceTy {
   /// is fulfilled.
   int32_t waitEvent(void *Event, AsyncInfoTy &AsyncInfo);
 
   /// Synchronize the event. It is expected to block the thread.
   int32_t syncEvent(void *Event);
 
   /// Destroy the event.
   int32_t destroyEvent(void *Event);
+
+  void setTeamProcs(int32_t num_team_procs) { TeamProcs = num_team_procs; }
+  int32_t getTeamProcs() { return TeamProcs; }
+
   /// }
 
 private:
   // Call to RTL
   void init(); // To be called only via DeviceTy::initOnce()
 
   /// Deinitialize the device (and plugin).
   void deinit();

openmp/libomptarget/include/omptarget.h

struct __tgt_device_info {
   void *Context = nullptr;
   void *Device = nullptr;
 };
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+int ompx_get_team_procs(int device_num);
 int omp_get_num_devices(void);
 int omp_get_device_num(void);
 int omp_get_initial_device(void);
 void *omp_target_alloc(size_t Size, int DeviceNum);
 void omp_target_free(void *DevicePtr, int DeviceNum);
 int omp_target_is_present(const void *Ptr, int DeviceNum);
 int omp_target_memcpy(void *Dst, const void *Src, size_t Length,
                       size_t DstOffset, size_t SrcOffset, int DstDevice,

openmp/libomptarget/include/omptargetplugin.h

 
 int32_t __tgt_rtl_destroy_event(int32_t ID, void *Event);
 // }
 
 int32_t __tgt_rtl_init_async_info(int32_t ID, __tgt_async_info **AsyncInfoPtr);
 int32_t __tgt_rtl_init_device_info(int32_t ID, __tgt_device_info *DeviceInfoPtr,
                                    const char **ErrStr);
 
+// Number of available physical processors to execute teams. AMD calls these
+// CUs. Nvidia calls them SMs. For CPUs modeling teams, they could be sockets.
+int32_t __tgt_rtl_number_of_team_procs(int32_t device_num);
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // _OMPTARGETPLUGIN_H_

openmp/libomptarget/include/rtl.h

struct RTLInfoTy {
   typedef int32_t(record_event_ty)(int32_t, void *, __tgt_async_info *);
   typedef int32_t(wait_event_ty)(int32_t, void *, __tgt_async_info *);
   typedef int32_t(sync_event_ty)(int32_t, void *);
   typedef int32_t(destroy_event_ty)(int32_t, void *);
   typedef int32_t(release_async_info_ty)(int32_t, __tgt_async_info *);
   typedef int32_t(init_async_info_ty)(int32_t, __tgt_async_info **);
   typedef int64_t(init_device_into_ty)(int64_t, __tgt_device_info *,
                                        const char **);
+  typedef int32_t(number_of_team_procs_ty)(int32_t);
 
   int32_t Idx = -1;             // RTL index, index is the number of devices
                                 // of other RTLs that were registered before,
                                 // i.e. the OpenMP index of the first device
                                 // to be registered with this RTL.
   int32_t NumberOfDevices = -1; // Number of devices this RTL deals with.
 
   std::unique_ptr<llvm::sys::DynamicLibrary> LibraryHandler;
#endif
   print_device_info_ty *print_device_info = nullptr;
   create_event_ty *create_event = nullptr;
   record_event_ty *record_event = nullptr;
   wait_event_ty *wait_event = nullptr;
   sync_event_ty *sync_event = nullptr;
   destroy_event_ty *destroy_event = nullptr;
   init_async_info_ty *init_async_info = nullptr;
   init_device_into_ty *init_device_info = nullptr;
+  number_of_team_procs_ty *number_of_team_procs = nullptr;
   release_async_info_ty *release_async_info = nullptr;
 
   // Are there images associated with this RTL.
   bool IsUsed = false;
 
   llvm::DenseSet<const __tgt_device_image *> UsedImages;
 
   // Mutex for thread-safety when calling RTL interface functions.

openmp/libomptarget/plugins/amdgpu/src/rtl.cpp

bool IsImageCompatibleWithEnv(const char *ImgInfo, std::string EnvInfo) {
   return true;
 }
 
 extern "C" {
 int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *Image) {
   return elfMachineIdIsAmdgcn(Image);
 }
 
+int __tgt_rtl_number_of_team_procs(int DeviceId) {
+  return DeviceInfo().ComputeUnits[DeviceId];
+}
+
 int32_t __tgt_rtl_is_valid_binary_info(__tgt_device_image *image,
                                        __tgt_image_info *info) {
   if (!__tgt_rtl_is_valid_binary(image))
     return false;
 
   // A subarchitecture was not specified. Assume it is compatible.
   if (!info->Arch)
     return true;

openmp/libomptarget/plugins/cuda/src/rtl.cpp

class DeviceRTLTy {
   using StreamPoolTy = ResourcePoolTy<StreamAllocatorTy>;
   std::vector<std::unique_ptr<StreamPoolTy>> StreamPool;
 
   using EventPoolTy = ResourcePoolTy<EventAllocatorTy>;
   std::vector<std::unique_ptr<EventPoolTy>> EventPool;
 
   std::vector<DeviceDataTy> DeviceData;
   std::vector<std::vector<CUmodule>> Modules;
+  std::vector<int> NumberOfTeamProcs;
 
   /// Vector of flags indicating the initalization status of all associated
   /// devices.
   std::vector<bool> InitializedFlags;
 
   enum class PeerAccessState : uint8_t { Unkown, Yes, No };
   std::vector<std::vector<PeerAccessState>> PeerAccessMatrix;
   std::mutex PeerAccessMatrixLock;
DeviceRTLTy()
 
     if (NumberOfDevices == 0) {
       DP("There are no devices supporting CUDA.\n");
       return;
     }
 
     DeviceData.resize(NumberOfDevices);
     Modules.resize(NumberOfDevices);
+    NumberOfTeamProcs.resize(NumberOfDevices);
     StreamPool.resize(NumberOfDevices);
     EventPool.resize(NumberOfDevices);
     PeerAccessMatrix.resize(NumberOfDevices);
     for (auto &V : PeerAccessMatrix)
       V.resize(NumberOfDevices, PeerAccessState::Unkown);
 
     // Get environment variables regarding teams
     if (const char *EnvStr = getenv("OMP_TEAM_LIMIT")) {
public:
 
   // Check whether a given DeviceId is valid
   bool isValidDeviceId(const int DeviceId) const {
     return DeviceId >= 0 && DeviceId < NumberOfDevices;
   }
 
   int getNumOfDevices() const { return NumberOfDevices; }
 
+  int getNumOfTeamProcs(int devid) const { return NumberOfTeamProcs[devid]; }
+
   void setRequiresFlag(const int64_t Flags) { this->RequiresFlags = Flags; }
 
   int initDevice(const int DeviceId) {
     CUdevice Device;
 
     DP("Getting device %d\n", DeviceId);
     CUresult Err = cuDeviceGet(&Device, DeviceId);
     if (!checkResult(Err, "Error returned from cuDeviceGet\n"))
if (Err != CUDA_SUCCESS) {
       DP("Error getting max grid dimension, use default value %d\n",
          DeviceRTLTy::DefaultNumTeams);
       DeviceData[DeviceId].BlocksPerGrid = DeviceRTLTy::DefaultNumTeams;
     } else {
       DP("Using %d CUDA blocks per grid\n", MaxGridDimX);
       DeviceData[DeviceId].BlocksPerGrid = MaxGridDimX;
     }
 
+    // Query attributes to for number of SMs for ompx_get_team_procs(devid)
+    int TmpTeamProcs;
+    Err = cuDeviceGetAttribute(
+        &TmpTeamProcs, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, Device);
+    if (Err != CUDA_SUCCESS) {
+      DP("Error: on CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, using %d\n", 16);
+      NumberOfTeamProcs[DeviceId] = 16;
+    } else {
+      DP("Device %d has %d procs for team execution\n", DeviceId, TmpTeamProcs);
+      NumberOfTeamProcs[DeviceId] = TmpTeamProcs;
+    }
+
     // We are only exploiting threads along the x axis.
     int MaxBlockDimX;
     Err = cuDeviceGetAttribute(&MaxBlockDimX,
                                CU_DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_X, Device);
     if (Err != CUDA_SUCCESS) {
       DP("Error getting max block dimension, use default value %d\n",
          DeviceRTLTy::DefaultNumThreads);
       DeviceData[DeviceId].ThreadsPerBlock = DeviceRTLTy::DefaultNumThreads;
int32_t __tgt_rtl_is_valid_binary_info(__tgt_device_image *Image,
   }
 
   DP("Image has compatible compute capability: %s\n", Info->Arch);
   return true;
 }
 
 int32_t __tgt_rtl_number_of_devices() { return DeviceRTL.getNumOfDevices(); }
 
+int32_t __tgt_rtl_number_of_team_procs(int devid) {
+  return DeviceRTL.getNumOfTeamProcs(devid);
+}
+
 int64_t __tgt_rtl_init_requires(int64_t RequiresFlags) {
   DP("Init requires flags to %" PRId64 "\n", RequiresFlags);
   DeviceRTL.setRequiresFlag(RequiresFlags);
   return RequiresFlags;
 }
 
 int32_t __tgt_rtl_is_data_exchangable(int32_t SrcDevId, int DstDevId) {
   if (DeviceRTL.isValidDeviceId(SrcDevId) &&

openmp/libomptarget/plugins/generic-elf-64bit/src/rtl.cpp

 
 #ifndef TARGET_ELF_ID
 #define TARGET_ELF_ID 0
 #endif
 
 #include "elf_common.h"
 
 #define NUMBER_OF_DEVICES 4
+#define NUMBER_OF_TEAM_PROCS 1
 #define OFFLOAD_SECTION_NAME "omp_offloading_entries"
 
 /// Array of Dynamic libraries loaded for this target.
 struct DynLibTy {
   std::string FileName;
   std::unique_ptr<DynamicLibrary> DynLib;
 };
 
#if TARGET_ELF_ID < 1
   return 0;
 #else
   return elf_check_machine(Image, TARGET_ELF_ID);
 #endif
 }
 
 int32_t __tgt_rtl_number_of_devices() { return NUMBER_OF_DEVICES; }
 
+// __tgt_rtl_number_of_team_procs supports ompx_get_team_procs(devid).
+// Imples that support multiple host teams may want to change this.
+int32_t __tgt_rtl_number_of_team_procs(int32_t device_id) {
+  return NUMBER_OF_TEAM_PROCS;
+}
+
 int32_t __tgt_rtl_init_device(int32_t DeviceId) { return OFFLOAD_SUCCESS; }
 
 __tgt_target_table *__tgt_rtl_load_binary(int32_t DeviceId,
                                           __tgt_device_image *Image) {
 
   DP("Dev %d: load binary from " DPxMOD " image\n", DeviceId,
      DPxPTR(Image->ImageStart));
 

openmp/libomptarget/plugins/ve/src/rtl.cpp

static int target_run_function_wait(uint32_t DeviceID, uint64_t FuncAddr,
   }
   return OFFLOAD_SUCCESS;
 }
 
 // Return the number of available devices of the type supported by the
 // target RTL.
 int32_t __tgt_rtl_number_of_devices(void) { return DeviceInfo.NodeIds.size(); }
 
+int32_t __tgt_rtl_number_of_team_procs(int device_id) {
+  return DeviceInfo.ProcHandles[device_id].size();
+}
+
 // Return an integer different from zero if the provided device image can be
 // supported by the runtime. The functionality is similar to comparing the
 // result of __tgt__rtl__load__binary to NULL. However, this is meant to be a
 // lightweight query to determine if the RTL is suitable for an image without
 // having to load the library, which can be expensive.
 int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *Image) {
 #if TARGET_ELF_ID < 1
   return 0;

openmp/libomptarget/src/api.cpp

EXTERN int omp_get_num_devices(void) {
   size_t DevicesSize = PM->Devices.size();
   PM->RTLsMtx.unlock();
 
   DP("Call to omp_get_num_devices returning %zd\n", DevicesSize);
 
   return DevicesSize;
 }
 
+EXTERN int ompx_get_team_procs(int device_num) {
+  if (!deviceIsReady(device_num)) {
+    DP("Device %d did not initialize\n", device_num);
+    // return 1 team proc for initial/host device
+    return 1;
+  }
+  TIMESCOPE();
+  PM->RTLsMtx.lock();
+  int TeamProcs = PM->Devices[device_num]->getTeamProcs();
+  PM->RTLsMtx.unlock();
+  DP("Call to ompx_get_team_procs returning %d\n", TeamProcs);
+  return TeamProcs;
+}
+
 EXTERN int omp_get_device_num(void) {
   TIMESCOPE();
   int HostDevice = omp_get_initial_device();
 
   DP("Call to omp_get_device_num returning %d\n", HostDevice);
 
   return HostDevice;
 }

openmp/libomptarget/src/device.cpp

 /// Init device, should not be called directly.
 void DeviceTy::init() {
   // Make call to init_requires if it exists for this plugin.
   if (RTL->init_requires)
     RTL->init_requires(PM->RTLs.RequiresFlags);
   int32_t Ret = RTL->init_device(RTLDeviceID);
   if (Ret != OFFLOAD_SUCCESS)
     return;
+  setTeamProcs(RTL->number_of_team_procs(RTLDeviceID));
 
   IsInit = true;
 }
 
 /// Thread-safe method to initialize the device only once.
 int32_t DeviceTy::initOnce() {
   std::call_once(InitFlag, &DeviceTy::init, this);
 

openmp/libomptarget/src/exports

global:
     __tgt_target_nowait_mapper;
     __tgt_target_teams_nowait_mapper;
     __tgt_target_kernel;
     __tgt_target_kernel_nowait;
     __tgt_mapper_num_components;
     __tgt_push_mapper_component;
     __kmpc_push_target_tripcount;
     __kmpc_push_target_tripcount_mapper;
+    ompx_get_team_procs;
     omp_get_num_devices;
     omp_get_device_num;
     omp_get_initial_device;
     omp_target_alloc;
     omp_target_free;
     omp_target_is_present;
     omp_target_memcpy;
     omp_target_memcpy_rect;

openmp/libomptarget/src/rtl.cpp

for (auto *Name : RTLNames) {
     bool ValidPlugin = true;
 
     if (!(*((void **)&R.is_valid_binary) =
               DynLibrary->getAddressOfSymbol("__tgt_rtl_is_valid_binary")))
       ValidPlugin = false;
     if (!(*((void **)&R.number_of_devices) =
               DynLibrary->getAddressOfSymbol("__tgt_rtl_number_of_devices")))
       ValidPlugin = false;
+    if (!(*((void **)&R.number_of_team_procs) =
+              DynLibrary->getAddressOfSymbol("__tgt_rtl_number_of_team_procs")))
+      ValidPlugin = false;
     if (!(*((void **)&R.init_device) =
               DynLibrary->getAddressOfSymbol("__tgt_rtl_init_device")))
       ValidPlugin = false;
     if (!(*((void **)&R.load_binary) =
               DynLibrary->getAddressOfSymbol("__tgt_rtl_load_binary")))
       ValidPlugin = false;
     if (!(*((void **)&R.data_alloc) =
               DynLibrary->getAddressOfSymbol("__tgt_rtl_data_alloc")))

openmp/libomptarget/test/api/ompx_get_team_procs.c

+// RUN: %libomptarget-compile-run-and-check-generic
+
+#include <omp.h>
+#include <stdio.h>
+
+int test_omp_get_team_procs(int devid) {
+  /* checks that ompx_get_team_procs() > 0 */
+  int team_procs = ompx_get_team_procs(devid);
+  printf("team_procs(%d) = %d\n", devid, team_procs);
+
+#pragma omp target
+  {}
+
+  return (team_procs > 0);
+}
+
+int main() {
+  int i;
+  int failed = 0;
+
+  if (!test_omp_get_team_procs(omp_get_initial_device())) {
+    failed++;
+  }
+  if (!test_omp_get_team_procs(omp_get_default_device())) {
+    failed++;
+  }
+  if (failed)
+    printf("FAIL\n");
+  else
+    printf("PASS\n");
+  return failed;
+}
+
+// CHECK: PASS

openmp/runtime/src/include/omp.h.var

#   if defined(_OPENMP) && _OPENMP >= 201811
     #pragma omp end declare variant
 #   endif
 
     /* OpenMP 5.2 */
     extern int __KAI_KMPC_CONVENTION omp_in_explicit_task(void);
 
     /* LLVM Extensions */
     extern void *llvm_omp_target_dynamic_shared_alloc();
+    extern int __KAI_KMPC_CONVENTION ompx_get_team_procs(int);
 
 #   undef __KAI_KMPC_CONVENTION
 #   undef __KMP_IMP
 
     /* Warning:
        The following typedefs are not standard, deprecated and will be removed in a future release.
     */
     typedef int     omp_int_t;
     typedef double  omp_wtime_t;
 
 #   ifdef __cplusplus
     }
 #   endif
 
 #endif /* __OMP_H */