[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