address comments, add pushpop corr interface, better test, complete enums

Author: Brian Coutinho

libkineto/include/KinetoDynamicPluginInterface.h

@@ -1,11 +1,11 @@
-#ifndef KINETO_PLUGIN_FORMAT_H
-#define KINETO_PLUGIN_FORMAT_H
+#ifndef KINETO_DYNAMIC_PLUGIN_FORMAT_H
+#define KINETO_DYNAMIC_PLUGIN_FORMAT_H
 
 #include <stdint.h>
 
 #define KINETO_PLUGIN_LIB_DIR_PATH_ENV_VARIABLE "KINETO_PLUGIN_LIB_DIR_PATH"
 
-// This file uses unpadded struct size for version control
+// NOTE: This file uses unpadded struct size for version control
 // To support backward compatibility, do NOT remove, reorder, or resize the
 // existing fields For details, visit
 // https://github.com/annarev/community/commit/09e231324c3b2b1f653c3f385fd2120edd460815
@@ -155,7 +155,7 @@ struct KinetoPlugin_TraceBuilder {
       const struct KinetoPlugin_ProfileEventFlow *pProfileEventFlow);
   // If metadata value is a string itself (e.g., "metadata value"), MUST add
   // additional quote around it (e.g., "\"metadata value\"") If metadata value
-  // is a list (e.g., 1, 2, 3), MUST add squre bracket around it (e.g., "[1, 2,
+  // is a list (e.g., 1, 2, 3), MUST add square bracket around it (e.g., "[1, 2,
   // 3]")
   int (*addLastEventMetadata)(
       KinetoPlugin_TraceBuilderHandle *pTraceBuilderHandle, const char *pKey,
@@ -214,7 +214,7 @@ struct KinetoPlugin_ProfilerQuery_Params {
 };
 #define KINETO_PLUGIN_PROFILER_QUERY_PARAMS_UNPADDED_STRUCT_SIZE               \
   KINETO_PLUGIN_UNPADDED_STRUCT_SIZE(struct KinetoPlugin_ProfilerQuery_Params, \
-                                     profilerNameMaxLen)
+                                     supportedActivityTypesMaxLen)
 
 struct KinetoPlugin_ProfilerStart_Params {
   // Always set to KINETO_PLUGIN_PROFILER_START_PARAMS_UNPADDED_STRUCT_SIZE
@@ -253,25 +253,113 @@ struct KinetoPlugin_ProfilerProcessEvents_Params {
   KINETO_PLUGIN_UNPADDED_STRUCT_SIZE(                                          \
       struct KinetoPlugin_ProfilerProcessEvents_Params, pTraceBuilder)
 
+struct KinetoPlugin_ProfilerPushCorrelationId_Params {
+  // Always set to
+  // KINETO_PLUGIN_PROFILER_PUSH_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE
+  size_t unpaddedStructSize;
+
+  // [in] Correlation ID
+  uint64_t correlationId;
+
+  // [in] An instance created via profilerCreate()
+  KinetoPlugin_ProfilerHandle *pProfilerHandle;
+};
+#define KINETO_PLUGIN_PROFILER_PUSH_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE \
+  KINETO_PLUGIN_UNPADDED_STRUCT_SIZE(                                          \
+      struct KinetoPlugin_ProfilerPushCorrelationId_Params, pProfilerHandle)
+
+struct KinetoPlugin_ProfilerPopCorrelationId_Params {
+  // Always set to
+  // KINETO_PLUGIN_PROFILER_POP_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE
+  size_t unpaddedStructSize;
+
+  // [in] An instance created via profilerCreate()
+  KinetoPlugin_ProfilerHandle *pProfilerHandle;
+};
+#define KINETO_PLUGIN_PROFILER_POP_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE  \
+  KINETO_PLUGIN_UNPADDED_STRUCT_SIZE(                                          \
+      struct KinetoPlugin_ProfilerPopCorrelationId_Params, pProfilerHandle)
+
+struct KinetoPlugin_ProfilerPushUserCorrelationId_Params {
+  // Always set to
+  // KINETO_PLUGIN_PROFILER_PUSH_USER_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE
+  size_t unpaddedStructSize;
+
+  // [in] User correlation ID
+  uint64_t userCorrelationId;
+
+  // [in] An instance created via profilerCreate()
+  KinetoPlugin_ProfilerHandle *pProfilerHandle;
+};
+#define KINETO_PLUGIN_PROFILER_PUSH_USER_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE \
+  KINETO_PLUGIN_UNPADDED_STRUCT_SIZE(                                               \
+      struct KinetoPlugin_ProfilerPushUserCorrelationId_Params,                     \
+      pProfilerHandle)
+
+struct KinetoPlugin_ProfilerPopUserCorrelationId_Params {
+  // Always set to
+  // KINETO_PLUGIN_PROFILER_POP_USER_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE
+  size_t unpaddedStructSize;
+
+  // [in] An instance created via profilerCreate()
+  KinetoPlugin_ProfilerHandle *pProfilerHandle;
+};
+#define KINETO_PLUGIN_PROFILER_POP_USER_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE \
+  KINETO_PLUGIN_UNPADDED_STRUCT_SIZE(                                              \
+      struct KinetoPlugin_ProfilerPopUserCorrelationId_Params,                     \
+      pProfilerHandle)
+
 struct KinetoPlugin_ProfilerInterface {
   // Always set to KINETO_PLUGIN_PROFILER_INTERFACE_UNPADDED_STRUCT_SIZE
   size_t unpaddedStructSize;
 
+  // Create an instance of the profiler session
   int (*profilerCreate)(
       struct KinetoPlugin_ProfilerCreate_Params *pProfilerCreateParams);
 
+  // Destroy an instance of the profiler session
   int (*profilerDestroy)(
       struct KinetoPlugin_ProfilerDestroy_Params *pProfilerDestroyParams);
 
+  // Query additional information about the profile such as its name, supported
+  // activity types, etc.
   int (*profilerQuery)(
       struct KinetoPlugin_ProfilerQuery_Params *pProfilerQueryParams);
 
+  // Start the trace collection
   int (*profilerStart)(
       struct KinetoPlugin_ProfilerStart_Params *pProfilerStartParams);
 
+  // Stop the trace collection
   int (*profilerStop)(
       struct KinetoPlugin_ProfilerStop_Params *pProfilerStopParams);
 
+  // The following interfaces are optional and can be used to track correlation
+  // IDs If not implemented, the correlation IDs will not be tracked through the
+  // profiler session.
+
+  // Push a correlation ID to the profiler session
+  int (*profilerPushCorrelationId)(
+      struct KinetoPlugin_ProfilerPushCorrelationId_Params
+          *pProfilerPushCorrelationIdParams);
+
+  // Pop a correlation ID from the profiler session
+  int (*profilerPopCorrelationId)(
+      struct KinetoPlugin_ProfilerPopCorrelationId_Params
+          *pProfilerPopCorrelationIdParams);
+
+  // Push a user correlation ID to the profiler session
+  int (*profilerPushUserCorrelationId)(
+      struct KinetoPlugin_ProfilerPushUserCorrelationId_Params
+          *pProfilerPushUserCorrelationIdParams);
+
+  // Pop a user correlation ID from the profiler session
+  int (*profilerPopUserCorrelationId)(
+      struct KinetoPlugin_ProfilerPopUserCorrelationId_Params
+          *pProfilerPopUserCorrelationIdParams);
+
+  // Process the events collected by the profiler session to the central
+  // profiler. This leverages the trace builder to enqueue events.
   int (*profilerProcessEvents)(struct KinetoPlugin_ProfilerProcessEvents_Params
                                    *pProfilerProcessEventsParams);
 };
@@ -304,4 +392,4 @@ int KinetoPlugin_register(const struct KinetoPlugin_Registry *pRegistry);
 }
 #endif
 
-#endif
+#endif // KINETO_DYNAMIC_PLUGIN_FORMAT_H

libkineto/src/dynamic_plugin/PluginProfiler.h

@@ -95,6 +95,63 @@ class PluginProfilerSession : public IActivityProfilerSession {
     }
   }
 
+  void pushCorrelationId(uint64_t id) override {
+    KinetoPlugin_ProfilerPushCorrelationId_Params pushCorrelationIdParams{
+        KINETO_PLUGIN_PROFILER_PUSH_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE};
+    pushCorrelationIdParams.pProfilerHandle = pProfilerHandle_;
+    pushCorrelationIdParams.correlationId = id;
+
+    int errorCode =
+        profiler_.profilerPushCorrelationId(&pushCorrelationIdParams);
+    if (errorCode != 0) {
+      LOG(ERROR) << "Plugin profiler " << name_
+                 << " failed at profilerPushCorrelationId() with error "
+                 << errorCode;
+    }
+  }
+
+  void popCorrelationId() override {
+    KinetoPlugin_ProfilerPopCorrelationId_Params popCorrelationIdParams{
+        KINETO_PLUGIN_PROFILER_POP_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE};
+    popCorrelationIdParams.pProfilerHandle = pProfilerHandle_;
+
+    int errorCode = profiler_.profilerPopCorrelationId(&popCorrelationIdParams);
+    if (errorCode != 0) {
+      LOG(ERROR) << "Plugin profiler " << name_
+                 << " failed at profilerPopCorrelationId() with error "
+                 << errorCode;
+    }
+  }
+
+  void pushUserCorrelationId(uint64_t id) override {
+    KinetoPlugin_ProfilerPushUserCorrelationId_Params pushUserCorrelationIdParams{
+        KINETO_PLUGIN_PROFILER_PUSH_USER_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE};
+    pushUserCorrelationIdParams.pProfilerHandle = pProfilerHandle_;
+    pushUserCorrelationIdParams.userCorrelationId = id;
+
+    int errorCode =
+        profiler_.profilerPushUserCorrelationId(&pushUserCorrelationIdParams);
+    if (errorCode != 0) {
+      LOG(ERROR) << "Plugin profiler " << name_
+                 << " failed at profilerPushUserCorrelationId() with error "
+                 << errorCode;
+    }
+  }
+
+  void popUserCorrelationId() override {
+    KinetoPlugin_ProfilerPopUserCorrelationId_Params popUserCorrelationIdParams{
+        KINETO_PLUGIN_PROFILER_POP_USER_CORRELATION_ID_PARAMS_UNPADDED_STRUCT_SIZE};
+    popUserCorrelationIdParams.pProfilerHandle = pProfilerHandle_;
+
+    int errorCode =
+        profiler_.profilerPopUserCorrelationId(&popUserCorrelationIdParams);
+    if (errorCode != 0) {
+      LOG(ERROR) << "Plugin profiler " << name_
+                 << " failed at profilerPopUserCorrelationId() with error "
+                 << errorCode;
+    }
+  }
+
   TraceStatus status() { return status_; }
 
   // returns errors with this trace
@@ -169,8 +226,13 @@ class PluginProfiler : public IActivityProfiler {
 
 public:
   PluginProfiler(const KinetoPlugin_ProfilerInterface &profiler)
-      : profiler_(profiler) {
-    validateProfiler();
+      : profiler_(profiler), isValid_(true) {
+    isValid_ = validateProfiler();
+    if (!isValid_) {
+      LOG(ERROR) << "Plugin profiler " << name_
+                 << " is not valid; skipping registration";
+      return;
+    }
 
     char profilerName[32];
 
@@ -216,6 +278,13 @@ class PluginProfiler : public IActivityProfiler {
   std::unique_ptr<IActivityProfilerSession>
   configure(const std::set<ActivityType> &activity_types,
             const Config & /*config*/) override {
+    // Check if profiler is valid
+    if (!isValid_) {
+      LOG(ERROR) << "Plugin profiler " << name_
+                 << " is not valid, cannot configure";
+      return nullptr;
+    }
+
     // Check if the plugin supports ANY of the requested activity types
     bool hasSupported = false;
     for (const auto &activity_type : activity_types) {
@@ -245,7 +314,9 @@ class PluginProfiler : public IActivityProfiler {
   }
 
 private:
-  void validateProfiler() {
+  bool validateProfiler() {
+    bool isValid = true;
+
     // Handle versioning
     // Currently expect the exact same version
     if (profiler_.unpaddedStructSize <
@@ -265,14 +336,128 @@ class PluginProfiler : public IActivityProfiler {
       profiler_.profilerStop = [](struct KinetoPlugin_ProfilerStop_Params *) {
         return -1;
       };
+      profiler_.profilerPushCorrelationId =
+          [](struct KinetoPlugin_ProfilerPushCorrelationId_Params *) {
+            return -1;
+          };
+      profiler_.profilerPopCorrelationId =
+          [](struct KinetoPlugin_ProfilerPopCorrelationId_Params *) {
+            return -1;
+          };
+      profiler_.profilerPushUserCorrelationId =
+          [](struct KinetoPlugin_ProfilerPushUserCorrelationId_Params *) {
+            return -1;
+          };
+      profiler_.profilerPopUserCorrelationId =
+          [](struct KinetoPlugin_ProfilerPopUserCorrelationId_Params *) {
+            return -1;
+          };
       profiler_.profilerProcessEvents =
           [](struct KinetoPlugin_ProfilerProcessEvents_Params *) { return -1; };
+      return false;
+    }
+
+    // Check if individual function pointers are implemented
+    // For critical functions, set them to error-returning stubs if nullptr and
+    // mark as invalid
+    if (profiler_.profilerCreate == nullptr) {
+      LOG(ERROR) << "Plugin profiler profilerCreate is not implemented";
+      profiler_.profilerCreate =
+          [](struct KinetoPlugin_ProfilerCreate_Params *) { return -1; };
+      isValid = false;
+    }
+
+    if (profiler_.profilerDestroy == nullptr) {
+      LOG(ERROR) << "Plugin profiler profilerDestroy is not implemented";
+      profiler_.profilerDestroy =
+          [](struct KinetoPlugin_ProfilerDestroy_Params *) { return -1; };
+      isValid = false;
+    }
+
+    if (profiler_.profilerQuery == nullptr) {
+      LOG(ERROR) << "Plugin profiler profilerQuery is not implemented";
+      profiler_.profilerQuery = [](struct KinetoPlugin_ProfilerQuery_Params *) {
+        return -1;
+      };
+      isValid = false;
+    }
+
+    if (profiler_.profilerStart == nullptr) {
+      LOG(ERROR) << "Plugin profiler profilerStart is not implemented";
+      profiler_.profilerStart = [](struct KinetoPlugin_ProfilerStart_Params *) {
+        return -1;
+      };
+      isValid = false;
     }
+
+    if (profiler_.profilerStop == nullptr) {
+      LOG(ERROR) << "Plugin profiler profilerStop is not implemented";
+      profiler_.profilerStop = [](struct KinetoPlugin_ProfilerStop_Params *) {
+        return -1;
+      };
+      isValid = false;
+    }
+
+    if (profiler_.profilerProcessEvents == nullptr) {
+      LOG(ERROR) << "Plugin profiler profilerProcessEvents is not implemented";
+      profiler_.profilerProcessEvents =
+          [](struct KinetoPlugin_ProfilerProcessEvents_Params *) { return -1; };
+      isValid = false;
+    }
+
+    // For correlation functions, provide default warning implementations if not
+    // provided These don't affect validity
+    if (profiler_.profilerPushCorrelationId == nullptr) {
+      LOG(INFO) << "Plugin profiler profilerPushCorrelationId is not "
+                   "implemented, using default stub";
+      profiler_.profilerPushCorrelationId =
+          [](struct KinetoPlugin_ProfilerPushCorrelationId_Params *) {
+            LOG(WARNING) << "profilerPushCorrelationId called but not "
+                            "implemented by plugin";
+            return 0;
+          };
+    }
+
+    if (profiler_.profilerPopCorrelationId == nullptr) {
+      LOG(INFO) << "Plugin profiler profilerPopCorrelationId is not "
+                   "implemented, using default stub";
+      profiler_.profilerPopCorrelationId =
+          [](struct KinetoPlugin_ProfilerPopCorrelationId_Params *) {
+            LOG(WARNING) << "profilerPopCorrelationId called but not "
+                            "implemented by plugin";
+            return 0;
+          };
+    }
+
+    if (profiler_.profilerPushUserCorrelationId == nullptr) {
+      LOG(INFO) << "Plugin profiler profilerPushUserCorrelationId is not "
+                   "implemented, using default stub";
+      profiler_.profilerPushUserCorrelationId =
+          [](struct KinetoPlugin_ProfilerPushUserCorrelationId_Params *) {
+            LOG(WARNING) << "profilerPushUserCorrelationId called but not "
+                            "implemented by plugin";
+            return 0;
+          };
+    }
+
+    if (profiler_.profilerPopUserCorrelationId == nullptr) {
+      LOG(INFO) << "Plugin profiler profilerPopUserCorrelationId is not "
+                   "implemented, using default stub";
+      profiler_.profilerPopUserCorrelationId =
+          [](struct KinetoPlugin_ProfilerPopUserCorrelationId_Params *) {
+            LOG(WARNING) << "profilerPopUserCorrelationId called but not "
+                            "implemented by plugin";
+            return 0;
+          };
+    }
+
+    return isValid;
   }
 
   KinetoPlugin_ProfilerInterface profiler_;
   std::string name_;
   std::set<ActivityType> supportedActivities_;
+  bool isValid_;
 };
 
 } // namespace libkineto