[Offload] Add `MAX_WORK_GROUP_SIZE` device info query (#143718)

This adds a new device info query for the maximum workgroup/block size
for each dimension.

1 files changed, 40 insertions, 0 deletions

diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index e7da4eddce54..f9da63843670 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -302,6 +302,41 @@ Error olGetDeviceInfoImplDetail(ol_device_handle_t Device,
                      "plugin did not provide a response for this information");
   };
 
+  auto getInfoXyz =
+      [&](std::vector<std::string> Names) -> llvm::Expected<ol_dimensions_t> {
+    for (auto &Name : Names) {
+      if (auto Entry = Device->Info.get(Name)) {
+        auto Node = *Entry;
+        ol_dimensions_t Out{0, 0, 0};
+
+        auto getField = [&](StringRef Name, uint32_t &Dest) {
+          if (auto F = Node->get(Name)) {
+            if (!std::holds_alternative<size_t>((*F)->Value))
+              return makeError(
+                  ErrorCode::BACKEND_FAILURE,
+                  "plugin returned incorrect type for dimensions element");
+            Dest = std::get<size_t>((*F)->Value);
+          } else
+            return makeError(ErrorCode::BACKEND_FAILURE,
+                             "plugin didn't provide all values for dimensions");
+          return Plugin::success();
+        };
+
+        if (auto Res = getField("x", Out.x))
+          return Res;
+        if (auto Res = getField("y", Out.y))
+          return Res;
+        if (auto Res = getField("z", Out.z))
+          return Res;
+
+        return Out;
+      }
+    }
+
+    return makeError(ErrorCode::UNIMPLEMENTED,
+                     "plugin did not provide a response for this information");
+  };
+
   switch (PropName) {
   case OL_DEVICE_INFO_PLATFORM:
     return Info.write<void *>(Device->Platform);
@@ -314,6 +349,9 @@ Error olGetDeviceInfoImplDetail(ol_device_handle_t Device,
   case OL_DEVICE_INFO_DRIVER_VERSION:
     return Info.writeString(
         getInfoString({"CUDA Driver Version", "HSA Runtime Version"}));
+  case OL_DEVICE_INFO_MAX_WORK_GROUP_SIZE:
+    return Info.write(getInfoXyz({"Workgroup Max Size per Dimension" /*AMD*/,
+                                  "Maximum Block Dimensions" /*CUDA*/}));
   default:
     return createOffloadError(ErrorCode::INVALID_ENUMERATION,
                               "getDeviceInfo enum '%i' is invalid", PropName);
@@ -339,6 +377,8 @@ Error olGetDeviceInfoImplDetailHost(ol_device_handle_t Device,
     return Info.writeString("Liboffload");
   case OL_DEVICE_INFO_DRIVER_VERSION:
     return Info.writeString(LLVM_VERSION_STRING);
+  case OL_DEVICE_INFO_MAX_WORK_GROUP_SIZE:
+    return Info.write<ol_dimensions_t>(ol_dimensions_t{1, 1, 1});
   default:
     return createOffloadError(ErrorCode::INVALID_ENUMERATION,
                               "getDeviceInfo enum '%i' is invalid", PropName);