Fixes

Author: kevinsala

offload/libomptarget/device.cpp

@@ -373,7 +373,14 @@ bool DeviceTy::isAccessiblePtr(const void *Ptr, size_t Size) {
 }
 
 uint64_t DeviceTy::getMaxSharedTeamMemory() {
-  using DeviceQueryKind = llvm::omp::target::plugin::DeviceQueryKind;
-  return RTL->query_device_info(
-      RTLDeviceID, DeviceQueryKind::DEVICE_QUERY_MAX_SHARED_TEAM_MEM);
+  InfoTreeNode Info = RTL->query_device_info(RTLDeviceID);
+
+  auto EntryOpt = Info.get(DeviceInfo::WORK_GROUP_SHARED_MEM_SIZE);
+  if (!EntryOpt)
+    return 0;
+
+  auto Entry = *EntryOpt;
+  if (!std::holds_alternative<uint64_t>(Entry->Value))
+    return 0;
+  return std::get<uint64_t>(Entry->Value);
 }

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

@@ -299,10 +299,6 @@ struct InfoTreeNode {
   }
 };
 
-enum class DeviceQueryKind {
-  DEVICE_QUERY_MAX_SHARED_TEAM_MEM = 0,
-};
-
 /// Class wrapping a __tgt_device_image and its offload entry table on a
 /// specific device. This class is responsible for storing and managing
 /// the offload entries for an image on a device.
@@ -1515,12 +1511,12 @@ struct GenericPluginTy {
   /// Query the current state of an asynchronous queue.
   int32_t query_async(int32_t DeviceId, __tgt_async_info *AsyncInfoPtr);
 
+  /// Obtain information about the given device.
+  InfoTreeNode obtain_device_info(int32_t DeviceId);
+
   /// Prints information about the given devices supported by the plugin.
   void print_device_info(int32_t DeviceId);
 
-  /// Retrieve information about the given device.
-  int64_t query_device_info(int32_t DeviceId, DeviceQueryKind Query);
-
   /// Creates an event in the given plugin if supported.
   int32_t create_event(int32_t DeviceId, void **EventPtr);
 

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

@@ -2037,23 +2037,22 @@ int32_t GenericPluginTy::query_async(int32_t DeviceId,
   return OFFLOAD_SUCCESS;
 }
 
+InfoTreeNode GenericPluginTy::obtain_device_info(int32_t DeviceId) {
+  auto InfoOrErr = getDevice(DeviceId).obtainInfo();
+  if (auto Err = InfoOrErr.takeError()) {
+    REPORT("Failure to obtain device %d info: %s\n", DeviceId,
+           toString(std::move(Err)).data());
+    return InfoTreeNode{};
+  }
+  return *InfoOrErr;
+}
+
 void GenericPluginTy::print_device_info(int32_t DeviceId) {
   if (auto Err = getDevice(DeviceId).printInfo())
     REPORT("Failure to print device %d info: %s\n", DeviceId,
            toString(std::move(Err)).data());
 }
 
-int64_t GenericPluginTy::query_device_info(int32_t DeviceId,
-                                           DeviceQueryKind Query) {
-  const GenericDeviceTy &Device = getDevice(DeviceId);
-
-  switch (Query) {
-  case DeviceQueryKind::DEVICE_QUERY_MAX_SHARED_TEAM_MEM:
-    return Device.getMaxBlockSharedMemSize();
-  }
-  return 0;
-}
-
 int32_t GenericPluginTy::create_event(int32_t DeviceId, void **EventPtr) {
   auto Err = getDevice(DeviceId).createEvent(EventPtr);
   if (Err) {