[Offload] Add type information to device info nodes

[RossBrunton]

Rather than being "stringly typed", store values as a std::variant that
can hold various types. This means that liboffload doesn't have to do
any string parsing for integer/bool device info keys.

[llvmbot]

@llvm/pr-subscribers-offload

Author: Ross Brunton (RossBrunton)

Changes

Rather than being "stringly typed", store values as a std::variant that
can hold various types. This means that liboffload doesn't have to do
any string parsing for integer/bool device info keys.

---

Full diff: https://github.com/llvm/llvm-project/pull/144535.diff

2 Files Affected:

• (modified) offload/liboffload/src/OffloadImpl.cpp (+9-9)
• (modified) offload/plugins-nextgen/common/include/PluginInterface.h (+30-12)

diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index 770c212d804d2..9faa07bf80227 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -228,23 +228,23 @@ Error olGetDeviceInfoImplDetail(ol_device_handle_t Device,
   ReturnHelper ReturnValue(PropSize, PropValue, PropSizeRet);
 
   // Find the info if it exists under any of the given names
-  auto GetInfo = [&](std::vector<std::string> Names) {
+  auto GetInfoString = [&](std::vector<std::string> Names) {
     if (Device == HostDevice())
-      return std::string("Host");
+      return "Host";
 
     if (!Device->Device)
-      return std::string("");
+      return "";
 
     auto Info = Device->Device->obtainInfoImpl();
     if (auto Err = Info.takeError())
-      return std::string("");
+      return "";
 
     for (auto Name : Names) {
       if (auto Entry = Info->get(Name))
-        return (*Entry)->Value;
+        return std::get<std::string>((*Entry)->Value).c_str();
     }
 
-    return std::string("");
+    return "";
   };
 
   switch (PropName) {
@@ -254,12 +254,12 @@ Error olGetDeviceInfoImplDetail(ol_device_handle_t Device,
     return Device == HostDevice() ? ReturnValue(OL_DEVICE_TYPE_HOST)
                                   : ReturnValue(OL_DEVICE_TYPE_GPU);
   case OL_DEVICE_INFO_NAME:
-    return ReturnValue(GetInfo({"Device Name"}).c_str());
+    return ReturnValue(GetInfoString({"Device Name"}));
   case OL_DEVICE_INFO_VENDOR:
-    return ReturnValue(GetInfo({"Vendor Name"}).c_str());
+    return ReturnValue(GetInfoString({"Vendor Name"}));
   case OL_DEVICE_INFO_DRIVER_VERSION:
     return ReturnValue(
-        GetInfo({"CUDA Driver Version", "HSA Runtime Version"}).c_str());
+        GetInfoString({"CUDA Driver Version", "HSA Runtime Version"}));
   default:
     return createOffloadError(ErrorCode::INVALID_ENUMERATION,
                               "getDeviceInfo enum '%i' is invalid", PropName);
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index f5d995532b7a5..d76c415ce4cdd 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -17,6 +17,7 @@
 #include <list>
 #include <map>
 #include <shared_mutex>
+#include <variant>
 #include <vector>
 
 #include "ExclusiveAccess.h"
@@ -122,7 +123,9 @@ struct InfoTreeNode {
   static constexpr uint64_t IndentSize = 4;
 
   std::string Key;
-  std::string Value;
+  struct None {};
+  using VariantType = std::variant<uint64_t, std::string, bool, None>;
+  VariantType Value;
   std::string Units;
   // Need to specify a default value number of elements here as `InfoTreeNode`'s
   // size is unknown. This is a vector (rather than a Key->Value map) since:
@@ -131,15 +134,15 @@ struct InfoTreeNode {
   // * The same key can appear multiple times
   std::unique_ptr<llvm::SmallVector<InfoTreeNode, 8>> Children;
 
-  InfoTreeNode() : InfoTreeNode("", "", "") {}
-  InfoTreeNode(std::string Key, std::string Value, std::string Units)
+  InfoTreeNode() : InfoTreeNode("", None{}, "") {}
+  InfoTreeNode(std::string Key, VariantType Value, std::string Units)
       : Key(Key), Value(Value), Units(Units) {}
 
   /// Add a new info entry as a child of this node. The entry requires at least
   /// a key string in \p Key. The value in \p Value is optional and can be any
   /// type that is representable as a string. The units in \p Units is optional
   /// and must be a string.
-  template <typename T = std::string>
+  template <typename T = None>
   InfoTreeNode *add(std::string Key, T Value = T(),
                     const std::string &Units = std::string()) {
     assert(!Key.empty() && "Invalid info key");
@@ -147,15 +150,15 @@ struct InfoTreeNode {
     if (!Children)
       Children = std::make_unique<llvm::SmallVector<InfoTreeNode, 8>>();
 
-    std::string ValueStr;
-    if constexpr (std::is_same_v<T, bool>)
-      ValueStr = Value ? "Yes" : "No";
+    VariantType ValueVariant;
+    if constexpr (std::is_same_v<T, bool> || std::is_same_v<T, None>)
+      ValueVariant = Value;
     else if constexpr (std::is_arithmetic_v<T>)
-      ValueStr = std::to_string(Value);
+      ValueVariant = static_cast<uint64_t>(Value);
     else
-      ValueStr = Value;
+      ValueVariant = std::string{Value};
 
-    return &Children->emplace_back(Key, ValueStr, Units);
+    return &Children->emplace_back(Key, ValueVariant, Units);
   }
 
   std::optional<InfoTreeNode *> get(StringRef Key) {
@@ -184,8 +187,23 @@ struct InfoTreeNode {
           MaxKeySize - (Key.size() + KeyIndentSize) + IndentSize;
 
       llvm::outs() << std::string(KeyIndentSize, ' ') << Key
-                   << std::string(ValIndentSize, ' ') << Value
-                   << (Units.empty() ? "" : " ") << Units << "\n";
+                   << std::string(ValIndentSize, ' ');
+      std::visit(
+          [](auto &&V) {
+            using T = std::decay_t<decltype(V)>;
+            if constexpr (std::is_same_v<T, std::string>)
+              llvm::outs() << V;
+            else if constexpr (std::is_same_v<T, bool>)
+              llvm::outs() << (V ? "Yes" : "No");
+            else if constexpr (std::is_same_v<T, uint64_t>)
+              llvm::outs() << V;
+            else if constexpr (std::is_same_v<T, None>) {
+              // Do nothing
+            } else
+              static_assert(false, "doPrint visit not exhaustive");
+          },
+          Value);
+      llvm::outs() << (Units.empty() ? "" : " ") << Units << "\n";
     }
 
     // Print children

[jhuber6]

I'm not exactly sure what this nets us. Usually for these lookups you just do type erasure on a void pointer and expect the user to provide the proper size.

[RossBrunton]

@jhuber6

I'm not exactly sure what this nets us. Usually for these lookups you just do type erasure on a void pointer and expect the user to provide the proper size.

This is primarily a response to your concerns here: #143718 (comment)

I figure rather than fiddling around using conversion functions between int and strings, if we know the data type is an int we may as well just store it as an int.

We could store it as a void * and cast it to whatever type we need when we use it, but type information is required for llvm-offload-device-info to print it correctly and provides us with some free error checking.

[jhuber6]

I still feel like there's a more elegant way to store this device information struct but I don't have an immediate solution in the back of my head so this is fine for now.