add the device info API to trace builder

Author: Brian Coutinho

libkineto/include/KinetoDynamicPluginInterface.h

@@ -111,6 +111,21 @@ struct KinetoPlugin_ProfileEventFlow {
   KINETO_PLUGIN_UNPADDED_STRUCT_SIZE(struct KinetoPlugin_ProfileEventFlow,     \
                                      isStartPoint)
 
+struct KinetoPlugin_ProfileDeviceInfo {
+  // Always set to KINETO_PLUGIN_PROFILE_DEVICE_INFO_UNPADDED_STRUCT_SIZE
+  size_t unpaddedStructSize;
+  // Device ID
+  int64_t deviceId;
+  // Device sort index
+  int64_t sortIndex;
+  // Device name
+  const char *pName;
+  // Device label
+  const char *pLabel;
+};
+#define KINETO_PLUGIN_PROFILE_DEVICE_INFO_UNPADDED_STRUCT_SIZE                 \
+  KINETO_PLUGIN_UNPADDED_STRUCT_SIZE(struct KinetoPlugin_ProfileDeviceInfo,    \
+                                     pLabel)
 struct KinetoPlugin_ProfileResourceInfo {
   // Always set to KINETO_PLUGIN_PROFILE_RESOURCE_INFO_UNPADDED_STRUCT_SIZE
   size_t unpaddedStructSize;
@@ -161,6 +176,10 @@ struct KinetoPlugin_TraceBuilder {
       KinetoPlugin_TraceBuilderHandle *pTraceBuilderHandle, const char *pKey,
       const char *pValue);
 
+  int (*addDeviceInfo)(
+      KinetoPlugin_TraceBuilderHandle *pTraceBuilderHandle,
+      const struct KinetoPlugin_ProfileDeviceInfo *pProfileDeviceInfo);
+
   int (*addResourceInfo)(
       KinetoPlugin_TraceBuilderHandle *pTraceBuilderHandle,
       const struct KinetoPlugin_ProfileResourceInfo *pProfileResourceInfo);

libkineto/src/dynamic_plugin/PluginProfiler.h

@@ -191,6 +191,7 @@ class PluginProfilerSession : public IActivityProfilerSession {
 
     // Take ownership of trace buffer from builder
     traceBuffer_ = pluginTraceBuilder.getTraceBuffer();
+    deviceInfos_ = pluginTraceBuilder.getDeviceInfos();
     resourceInfos_ = pluginTraceBuilder.getResourceInfos();
 
     // Log events
@@ -207,7 +208,12 @@ class PluginProfilerSession : public IActivityProfilerSession {
   }
 
   // returns device info used in this trace, could be nullptr
-  std::unique_ptr<DeviceInfo> getDeviceInfo() override { return {}; }
+  std::unique_ptr<DeviceInfo> getDeviceInfo() override {
+    if (deviceInfos_.empty()) {
+      return {};
+    }
+    return std::make_unique<DeviceInfo>(deviceInfos_[0]);
+  }
 
   // returns resource info used in this trace, could be empty
   std::vector<ResourceInfo> getResourceInfos() override {
@@ -221,6 +227,7 @@ class PluginProfilerSession : public IActivityProfilerSession {
 
 private:
   std::unique_ptr<CpuTraceBuffer> traceBuffer_;
+  std::vector<DeviceInfo> deviceInfos_;
   std::vector<ResourceInfo> resourceInfos_;
   const KinetoPlugin_ProfilerInterface profiler_;
   KinetoPlugin_ProfilerHandle *pProfilerHandle_ = nullptr;

libkineto/src/dynamic_plugin/PluginTraceBuilder.h

@@ -29,7 +29,6 @@ class PluginTraceBuilder {
 
     if (pProfileEvent == nullptr) {
       LOG(ERROR) << "Failed to add event of nullptr";
-
       return -1;
     }
 
@@ -38,7 +37,6 @@ class PluginTraceBuilder {
     if (pProfileEvent->unpaddedStructSize <
         KINETO_PLUGIN_PROFILER_PROCESS_EVENTS_PARAMS_UNPADDED_STRUCT_SIZE) {
       LOG(ERROR) << "Profile event has an incompatible version";
-
       return -1;
     }
 
@@ -66,18 +64,15 @@ class PluginTraceBuilder {
 
     if (pName == nullptr) {
       LOG(ERROR) << "Failed to set last event name of nullptr";
-
       return -1;
     }
 
     if (buffer_->activities.empty()) {
       LOG(ERROR) << "Failed to set last event flow as there is no last event";
-
       return -1;
     }
 
     buffer_->activities.back()->activityName.assign(pName);
-
     return 0;
   }
 
@@ -97,13 +92,11 @@ class PluginTraceBuilder {
     if (pProfileEventFlow->unpaddedStructSize <
         KINETO_PLUGIN_PROFILE_EVENT_FLOW_UNPADDED_STRUCT_SIZE) {
       LOG(ERROR) << "Profile event flow has an incompatible version";
-
       return -1;
     }
 
     if (buffer_->activities.empty()) {
       LOG(ERROR) << "Failed to set last event flow as there is no last event";
-
       return -1;
     }
 
@@ -123,19 +116,43 @@ class PluginTraceBuilder {
 
     if (pKey == nullptr || pValue == nullptr) {
       LOG(ERROR) << "Failed to set last event metadata of nullptr";
-
       return -1;
     }
 
     if (buffer_->activities.empty()) {
       LOG(ERROR)
           << "Failed to set last event metadata as there is no last event";
-
       return -1;
     }
 
     buffer_->activities.back()->addMetadata(std::string{pKey},
                                             std::string{pValue});
+    return 0;
+  }
+
+  int addDeviceInfo(const KinetoPlugin_ProfileDeviceInfo *pProfileDeviceInfo) {
+    if (pProfileDeviceInfo == nullptr) {
+      LOG(ERROR) << "Failed to add device info of nullptr";
+      return -1;
+    }
+
+    // Handle versioning
+    // Currently expect the exact same version
+    if (pProfileDeviceInfo->unpaddedStructSize <
+        KINETO_PLUGIN_PROFILE_DEVICE_INFO_UNPADDED_STRUCT_SIZE) {
+      LOG(ERROR) << "Profile device info has an incompatible version";
+      return -1;
+    }
+
+    DeviceInfo deviceInfo(
+        pProfileDeviceInfo->deviceId, pProfileDeviceInfo->sortIndex,
+        pProfileDeviceInfo->pName
+            ? std::string(pProfileDeviceInfo->pName)
+            : std::to_string(pProfileDeviceInfo->deviceId),
+        pProfileDeviceInfo->pLabel ? std::string(pProfileDeviceInfo->pLabel)
+                                   : "");
+
+    deviceInfos_.push_back(deviceInfo);
 
     return 0;
   }
@@ -144,7 +161,6 @@ class PluginTraceBuilder {
       const KinetoPlugin_ProfileResourceInfo *pProfileResourceInfo) {
     if (pProfileResourceInfo == nullptr) {
       LOG(ERROR) << "Failed to add resource info of nullptr";
-
       return -1;
     }
 
@@ -153,7 +169,6 @@ class PluginTraceBuilder {
     if (pProfileResourceInfo->unpaddedStructSize <
         KINETO_PLUGIN_PROFILE_RESOURCE_INFO_UNPADDED_STRUCT_SIZE) {
       LOG(ERROR) << "Profile resource info has an incompatible version";
-
       return -1;
     }
 
@@ -178,6 +193,7 @@ class PluginTraceBuilder {
         .setLastEventName = cSetLastEventName,
         .setLastEventFlow = cSetLastEventFlow,
         .addLastEventMetadata = cAddLastEventMetadata,
+        .addDeviceInfo = cAddDeviceInfo,
         .addResourceInfo = cAddResourceInfo};
   }
 
@@ -186,6 +202,8 @@ class PluginTraceBuilder {
     return std::move(buffer_);
   }
 
+  std::vector<DeviceInfo> getDeviceInfos() { return deviceInfos_; }
+
   std::vector<ResourceInfo> getResourceInfos() { return resourceInfos_; }
 
 private:
@@ -220,6 +238,14 @@ class PluginTraceBuilder {
     return pPluginTraceBuilder->addLastEventMetadata(pKey, pValue);
   }
 
+  static int
+  cAddDeviceInfo(KinetoPlugin_TraceBuilderHandle *pTraceBuilderHandle,
+                 const KinetoPlugin_ProfileDeviceInfo *pProfileDeviceInfo) {
+    auto pPluginTraceBuilder =
+        reinterpret_cast<PluginTraceBuilder *>(pTraceBuilderHandle);
+    return pPluginTraceBuilder->addDeviceInfo(pProfileDeviceInfo);
+  }
+
   static int cAddResourceInfo(
       KinetoPlugin_TraceBuilderHandle *pTraceBuilderHandle,
       const KinetoPlugin_ProfileResourceInfo *pProfileResourceInfo) {
@@ -229,6 +255,7 @@ class PluginTraceBuilder {
   }
 
   std::unique_ptr<CpuTraceBuffer> buffer_;
+  std::vector<DeviceInfo> deviceInfos_;
   std::vector<ResourceInfo> resourceInfos_;
 };