Introduce XPU scope profiler extending existing XPU profiler plugin

libkineto/include/ActivityType.h

@@ -50,7 +50,9 @@ enum class ActivityType {
   PRIVATEUSE1_RUNTIME = 24, // host side privateUse1 runtime events
   PRIVATEUSE1_DRIVER = 25, // host side privateUse1 driver events
 
-  ENUM_COUNT = 26, // This is to add buffer and not used for any profiling logic. Add
+  XPU_SCOPE_PROFILER = 26, // XPUPTI Profiler scope for performance metrics
+
+  ENUM_COUNT = 27, // This is to add buffer and not used for any profiling logic. Add
   // your new type before it.
   OPTIONAL_ACTIVITY_TYPE_START = GLOW_RUNTIME,
 };
@@ -94,6 +96,7 @@ inline constexpr std::array<_ActivityTypeName, activityTypeCount + 1> _activityT
     {"collective_comm", ActivityType::COLLECTIVE_COMM},
     {"privateuse1_runtime", ActivityType::PRIVATEUSE1_RUNTIME},
     {"privateuse1_driver", ActivityType::PRIVATEUSE1_DRIVER},
+    {"xpu_scope_profiler", ActivityType::XPU_SCOPE_PROFILER},
     {"ENUM_COUNT", ActivityType::ENUM_COUNT},
 }};
 

libkineto/libkineto_defs.bzl

@@ -52,7 +52,10 @@ def get_libkineto_xpupti_srcs(with_api = True):
         "src/plugin/xpupti/XpuptiActivityProfiler.cpp",
         "src/plugin/xpupti/XpuptiActivityProfilerSession.cpp",
         "src/plugin/xpupti/XpuptiProfilerMacros.cpp",
+        "src/plugin/xpupti/XpuptiScopeProfilerApi.cpp",
         "src/plugin/xpupti/XpuptiScopeProfilerConfig.cpp",
+        "src/plugin/xpupti/XpuptiScopeProfilerHandlers.cpp",
+        "src/plugin/xpupti/XpuptiScopeProfilerSession.cpp",
     ] + (get_libkineto_cpu_only_srcs(with_api))
 
 def get_libkineto_aiupti_srcs(with_api = True):

libkineto/src/output_json.cpp

@@ -760,9 +760,13 @@ void ChromeTraceLogger::handleActivity(const libkineto::ITraceActivity& op) {
     return;
   }
 
-  if (op.type() == ActivityType::MTIA_COUNTERS) {
-    handleCounterEvent(op);
-    return;
+  switch (op.type()) {
+    case ActivityType::MTIA_COUNTERS:
+    case ActivityType::XPU_SCOPE_PROFILER:
+      handleCounterEvent(op);
+      return;
+    default:
+      break;
   }
 
   int64_t ts = op.timestamp();

libkineto/src/plugin/xpupti/XpuptiActivityApi.cpp

@@ -245,6 +245,17 @@ void XpuptiActivityApi::enableXpuptiActivities(
         XPUPTI_CALL(ptiViewEnable(PTI_VIEW_DRIVER_API));
         break;
 
+      case ActivityType::XPU_SCOPE_PROFILER:
+#if PTI_VERSION_AT_LEAST(0, 15)
+        // This case is handled in constructor of
+        // XpuptiScopeProfilerSession
+#else
+        throw std::runtime_error(
+            "Intel® PTI version required to use scope profiler is at least 0.15 "
+            "(available with Intel® oneAPI in version at least 2025.3.1).");
+#endif
+        break;
+
       case ActivityType::OVERHEAD:
         XPUPTI_CALL(ptiViewEnable(PTI_VIEW_COLLECTION_OVERHEAD));
         break;

libkineto/src/plugin/xpupti/XpuptiActivityHandlers.cpp

@@ -225,6 +225,14 @@ void XpuptiActivityProfilerSession::handleRuntimeKernelMemcpyMemsetActivities(
     trace_activity->resource = activity->_sycl_queue_id;
     trace_activity->flow.start = 0;
 
+    if constexpr (handleKernelActivities) {
+      kernelActivities_[activity->_kernel_id].emplace(
+          trace_activity->startTime,
+          trace_activity->endTime,
+          trace_activity->device,
+          trace_activity->resource);
+    }
+
     addResouceInfo(trace_activity->device, trace_activity->resource);
   }
 

libkineto/src/plugin/xpupti/XpuptiActivityProfiler.cpp

@@ -7,8 +7,8 @@
  */
 
 #include "XpuptiActivityProfiler.h"
-#include "XpuptiActivityApi.h"
-#include "XpuptiActivityProfilerSession.h"
+#include "XpuptiScopeProfilerApi.h"
+#include "XpuptiScopeProfilerSession.h"
 
 #include <fmt/ranges.h>
 #include <sycl/sycl.hpp>
@@ -66,8 +66,13 @@ std::unique_ptr<libkineto::IActivityProfilerSession> XPUActivityProfiler::
     configure(
         const std::set<ActivityType>& activity_types,
         const libkineto::Config& config) {
+#if PTI_VERSION_AT_LEAST(0, 15)
+  return std::make_unique<XpuptiScopeProfilerSession>(
+      XpuptiActivityApi::singleton(), name(), config, activity_types);
+#else
   return std::make_unique<XpuptiActivityProfilerSession>(
       XpuptiActivityApi::singleton(), name(), config, activity_types);
+#endif
 }
 
 std::unique_ptr<libkineto::IActivityProfilerSession> XPUActivityProfiler::

libkineto/src/plugin/xpupti/XpuptiActivityProfilerSession.h

@@ -95,7 +95,7 @@ class XpuptiActivityProfilerSession : public libkineto::IActivityProfilerSession
 
   void addResouceInfo(int32_t device_id, int32_t sycl_queue_id);
 
- private:
+ protected:
   static uint32_t iterationCount_;
   static std::vector<DeviceUUIDsT> deviceUUIDs_;
   static std::unordered_set<std::string_view> correlateRuntimeOps_;
@@ -117,6 +117,27 @@ class XpuptiActivityProfilerSession : public libkineto::IActivityProfilerSession
   std::unique_ptr<const libkineto::Config> config_{nullptr};
   const std::set<ActivityType>& activity_types_;
   std::string name_;
+
+  struct KernelActivity {
+    void emplace(
+        int64_t startTime,
+        int64_t endTime,
+        int32_t device,
+        int32_t resource) {
+      startTime_ = startTime;
+      endTime_ = endTime;
+      device_ = device;
+      resource_ = resource;
+    }
+
+    int64_t startTime_{0};
+    int64_t endTime_{0};
+    int32_t device_{0};
+    int32_t resource_{0};
+  };
+
+  std::unordered_map<uint64_t, KernelActivity> kernelActivities_;
+  uint64_t lastKernelActivityEndTime_{0};
 };
 
 } // namespace KINETO_NAMESPACE

libkineto/src/plugin/xpupti/XpuptiScopeProfilerApi.cpp

@@ -0,0 +1,178 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "XpuptiScopeProfilerApi.h"
+
+#if PTI_VERSION_AT_LEAST(0, 15)
+
+#include "XpuptiScopeProfilerConfig.h"
+
+#include <algorithm>
+#include <iterator>
+#include <stdexcept>
+
+#endif
+
+namespace KINETO_NAMESPACE {
+
+#if PTI_VERSION_AT_LEAST(0, 15)
+
+XpuptiScopeProfilerApi::safe_pti_scope_collection_handle_t::
+    safe_pti_scope_collection_handle_t(std::exception_ptr& exceptFromDestructor)
+    : exceptFromDestructor_(exceptFromDestructor) {
+  XPUPTI_CALL(ptiMetricsScopeEnable(&handle_));
+}
+
+XpuptiScopeProfilerApi::safe_pti_scope_collection_handle_t::
+    ~safe_pti_scope_collection_handle_t() noexcept {
+  try {
+    XPUPTI_CALL(ptiMetricsScopeDisable(handle_));
+  } catch (...) {
+    exceptFromDestructor_ = std::current_exception();
+  }
+}
+
+void XpuptiScopeProfilerApi::enableScopeProfiler(const Config& cfg) {
+  uint32_t deviceCount = 0;
+  XPUPTI_CALL(ptiMetricsGetDevices(nullptr, &deviceCount));
+
+  if (deviceCount == 0) {
+    throw std::runtime_error("No XPU devices available");
+  }
+
+  auto devices = std::make_unique<pti_device_properties_t[]>(deviceCount);
+  XPUPTI_CALL(ptiMetricsGetDevices(devices.get(), &deviceCount));
+
+  auto devicesHandles = std::make_unique<pti_device_handle_t[]>(deviceCount);
+  for (uint32_t i = 0; i < deviceCount; ++i) {
+    devicesHandles[i] = devices[i]._handle;
+  }
+
+  const auto& spcfg = XpuptiScopeProfilerConfig::get(cfg);
+  const auto& activitiesXpuptiMetrics = spcfg.activitiesXpuptiMetrics();
+
+  std::vector<const char*> metricNames;
+  metricNames.reserve(activitiesXpuptiMetrics.size());
+  std::transform(
+      activitiesXpuptiMetrics.begin(),
+      activitiesXpuptiMetrics.end(),
+      std::back_inserter(metricNames),
+      [](const std::string& s) { return s.c_str(); });
+
+  pti_metrics_scope_mode_t collectionMode = spcfg.xpuptiProfilerPerKernel()
+      ? PTI_METRICS_SCOPE_AUTO_KERNEL
+      : PTI_METRICS_SCOPE_USER;
+
+  if (collectionMode == PTI_METRICS_SCOPE_USER) {
+    throw std::runtime_error(
+        "XPUPTI_PROFILER_ENABLE_PER_KERNEL has to be set to 1. Other variants are currently not supported.");
+  }
+
+  scopeHandleOpt_.emplace(exceptFromScopeHandleDestructor_);
+  XPUPTI_CALL(ptiMetricsScopeConfigure(
+      *scopeHandleOpt_,
+      collectionMode,
+      devicesHandles.get(),
+      ((void)deviceCount, 1), // Only 1 device is currently supported
+      metricNames.data(),
+      metricNames.size()));
+
+  uint64_t expectedKernels = spcfg.xpuptiProfilerMaxScopes();
+  size_t estimatedCollectionBufferSize = 0;
+  XPUPTI_CALL(ptiMetricsScopeQueryCollectionBufferSize(
+      *scopeHandleOpt_, expectedKernels, &estimatedCollectionBufferSize));
+
+  XPUPTI_CALL(ptiMetricsScopeSetCollectionBufferSize(
+      *scopeHandleOpt_, estimatedCollectionBufferSize));
+}
+
+void XpuptiScopeProfilerApi::disableScopeProfiler() {
+  scopeHandleOpt_.reset();
+  if (exceptFromScopeHandleDestructor_) {
+    std::rethrow_exception(exceptFromScopeHandleDestructor_);
+  }
+}
+
+void XpuptiScopeProfilerApi::startScopeActivity() {
+  if (scopeHandleOpt_) {
+    XPUPTI_CALL(ptiMetricsScopeStartCollection(*scopeHandleOpt_));
+  }
+}
+
+void XpuptiScopeProfilerApi::stopScopeActivity() {
+  if (scopeHandleOpt_) {
+    XPUPTI_CALL(ptiMetricsScopeStopCollection(*scopeHandleOpt_));
+  }
+}
+
+static size_t IntDivRoundUp(size_t a, size_t b) {
+  return (a + b - 1) / b;
+}
+
+void XpuptiScopeProfilerApi::processScopeTrace(
+    std::function<void(
+        const pti_metrics_scope_record_t*,
+        const pti_metrics_scope_record_metadata_t& metadata)> handler) {
+  if (scopeHandleOpt_) {
+    pti_metrics_scope_record_metadata_t metadata;
+    metadata._struct_size = sizeof(pti_metrics_scope_record_metadata_t);
+
+    XPUPTI_CALL(ptiMetricsScopeGetMetricsMetadata(*scopeHandleOpt_, &metadata));
+
+    uint64_t collectionBuffersCount = 0;
+    XPUPTI_CALL(ptiMetricsScopeGetCollectionBuffersCount(
+        *scopeHandleOpt_, &collectionBuffersCount));
+
+    for (uint64_t bufferId = 0; bufferId < collectionBuffersCount; ++bufferId) {
+      void* collectionBuffer = nullptr;
+      size_t actualCollectionBufferSize = 0;
+      XPUPTI_CALL(ptiMetricsScopeGetCollectionBuffer(
+          *scopeHandleOpt_,
+          bufferId,
+          &collectionBuffer,
+          &actualCollectionBufferSize));
+
+      pti_metrics_scope_collection_buffer_properties_t metricsBufferProps;
+      metricsBufferProps._struct_size =
+          sizeof(pti_metrics_scope_collection_buffer_properties_t);
+      XPUPTI_CALL(ptiMetricsScopeGetCollectionBufferProperties(
+          *scopeHandleOpt_, collectionBuffer, &metricsBufferProps));
+
+      size_t requiredMetricsBufferSize = 0;
+      size_t recordsCount = 0;
+      XPUPTI_CALL(ptiMetricsScopeQueryMetricsBufferSize(
+          *scopeHandleOpt_,
+          collectionBuffer,
+          &requiredMetricsBufferSize,
+          &recordsCount));
+
+      if (recordsCount > 0) {
+        auto metricsBuffer =
+            std::make_unique<pti_metrics_scope_record_t[]>(IntDivRoundUp(
+                requiredMetricsBufferSize, sizeof(pti_metrics_scope_record_t)));
+
+        size_t actualRecordsCount = 0;
+        XPUPTI_CALL(ptiMetricsScopeCalculateMetrics(
+            *scopeHandleOpt_,
+            collectionBuffer,
+            metricsBuffer.get(),
+            requiredMetricsBufferSize,
+            &actualRecordsCount));
+
+        for (size_t recordId = 0; recordId < actualRecordsCount; ++recordId) {
+          auto record = metricsBuffer.get() + recordId;
+          handler(record, metadata);
+        }
+      }
+    }
+  }
+}
+
+#endif
+
+} // namespace KINETO_NAMESPACE

libkineto/src/plugin/xpupti/XpuptiScopeProfilerApi.h

@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include "XpuptiActivityApi.h"
+
+#if PTI_VERSION_AT_LEAST(0, 15)
+
+#include <pti/pti_metrics_scope.h>
+
+#include <optional>
+
+namespace KINETO_NAMESPACE {
+
+class Config;
+
+class XpuptiScopeProfilerApi {
+ public:
+  XpuptiScopeProfilerApi() = default;
+  XpuptiScopeProfilerApi(const XpuptiScopeProfilerApi&) = delete;
+  XpuptiScopeProfilerApi& operator=(const XpuptiScopeProfilerApi&) = delete;
+
+  ~XpuptiScopeProfilerApi() = default;
+
+  void enableScopeProfiler(const Config&);
+  void disableScopeProfiler();
+  void startScopeActivity();
+  void stopScopeActivity();
+
+  void processScopeTrace(
+      std::function<void(
+          const pti_metrics_scope_record_t*,
+          const pti_metrics_scope_record_metadata_t& metadata)> handler);
+
+ private:
+  struct safe_pti_scope_collection_handle_t {
+    safe_pti_scope_collection_handle_t(
+        std::exception_ptr& exceptFromDestructor);
+    ~safe_pti_scope_collection_handle_t() noexcept;
+
+    operator pti_scope_collection_handle_t() {
+      return handle_;
+    }
+
+    pti_scope_collection_handle_t handle_{};
+    std::exception_ptr& exceptFromDestructor_;
+  };
+
+  std::optional<safe_pti_scope_collection_handle_t> scopeHandleOpt_;
+  std::exception_ptr exceptFromScopeHandleDestructor_;
+};
+
+} // namespace KINETO_NAMESPACE
+
+#endif