[xpupti] Support more than 1 device in ptiMetricsScopeConfigure

libkineto/src/plugin/xpupti/XpuptiScopeProfilerApi.cpp

@@ -6,15 +6,40 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <fmt/core.h>
+
 #include <algorithm>
 #include <iterator>
+#include <span>
 #include <stdexcept>
+#include <utility>
 
 #include "XpuptiScopeProfilerApi.h"
 #include "XpuptiScopeProfilerConfig.h"
 
 namespace KINETO_NAMESPACE {
 
+std::vector<pti_device_handle_t> selectDeviceHandles(
+    std::span<const pti_device_handle_t> handles,
+    std::span<const int> indices) {
+  const auto outOfRange = [handles](int idx) {
+    return idx < 0 || std::cmp_greater_equal(idx, handles.size());
+  };
+  if (const auto bad = std::ranges::find_if(indices, outOfRange);
+      bad != indices.end()) {
+    throw std::runtime_error(fmt::format(
+        "XPUPTI_PROFILER_DEVICES index {} is out of range; {} XPU device(s) available",
+        *bad,
+        handles.size()));
+  }
+  // Gather: map each requested index to its device handle, preserving order.
+  std::vector<pti_device_handle_t> selected(indices.size());
+  std::ranges::transform(indices, selected.begin(), [handles](int idx) {
+    return handles[static_cast<std::size_t>(idx)];
+  });
+  return selected;
+}
+
 XpuptiScopeProfilerApi::safe_pti_scope_collection_handle_t::
     safe_pti_scope_collection_handle_t(std::exception_ptr& exceptFromDestructor)
     : exceptFromDestructor_(exceptFromDestructor) {
@@ -67,13 +92,53 @@ void XpuptiScopeProfilerApi::enableScopeProfiler(const Config& cfg) {
   }
 
   scopeHandleOpt_.emplace(exceptFromScopeHandleDestructor_);
+
+  const auto& requestedDevices = spcfg.xpuptiProfilerDevices();
+
+#if PTI_VERSION_AT_LEAST(0, 18)
+  if (requestedDevices.empty()) {
+    // Default: profile every available device (PTI auto-detect mode).
+    XPUPTI_CALL(ptiMetricsScopeConfigure(
+        *scopeHandleOpt_,
+        collectionMode,
+        /*devices_to_profile=*/nullptr,
+        /*device_count=*/0,
+        metricNames.data(),
+        metricNames.size()));
+  } else {
+    // Explicit subset: map requested indices to device handles.
+    auto selectedHandles = selectDeviceHandles(
+        {devicesHandles.get(), deviceCount}, requestedDevices);
+    XPUPTI_CALL(ptiMetricsScopeConfigure(
+        *scopeHandleOpt_,
+        collectionMode,
+        selectedHandles.data(),
+        static_cast<uint32_t>(selectedHandles.size()),
+        metricNames.data(),
+        metricNames.size()));
+  }
+#else
+  // PTI < 0.18 (pre PTI-363): multi-device metrics scope is not available;
+  // ptiMetricsScopeConfigure accepts only a single device.
+  if (requestedDevices.size() > 1) {
+    throw std::runtime_error(
+        "XPUPTI_PROFILER_DEVICES lists more than one device, but this build "
+        "links PTI < 0.18 which supports only single-device metrics scope. "
+        "Rebuild against PTI >= 0.18 for multi-device support.");
+  }
+  // Point at the single requested device (default: first device).
+  pti_device_handle_t singleHandle = requestedDevices.empty()
+      ? devicesHandles[0]
+      : selectDeviceHandles({devicesHandles.get(), deviceCount}, requestedDevices)
+            .front();
   XPUPTI_CALL(ptiMetricsScopeConfigure(
       *scopeHandleOpt_,
       collectionMode,
-      devicesHandles.get(),
-      ((void)deviceCount, 1), // Only 1 device is currently supported
+      &singleHandle,
+      1,
       metricNames.data(),
       metricNames.size()));
+#endif
 
   uint64_t expectedKernels = spcfg.xpuptiProfilerMaxScopes();
   size_t estimatedCollectionBufferSize = 0;

libkineto/src/plugin/xpupti/XpuptiScopeProfilerApi.h

@@ -9,6 +9,8 @@
 #pragma once
 
 #include <optional>
+#include <span>
+#include <vector>
 
 #include <pti/pti_metrics_scope.h>
 
@@ -51,4 +53,10 @@ class XpuptiScopeProfilerApi {
   std::exception_ptr exceptFromScopeHandleDestructor_;
 };
 
+// Map requested device indices to PTI device handles, preserving order.
+// Throws std::runtime_error if any index is out of [0, deviceCount).
+std::vector<pti_device_handle_t> selectDeviceHandles(
+    std::span<const pti_device_handle_t> handles,
+    std::span<const int> indices);
+
 } // namespace KINETO_NAMESPACE

libkineto/src/plugin/xpupti/XpuptiScopeProfilerConfig.cpp

@@ -14,6 +14,10 @@
 #include <fmt/ostream.h>
 #include <fmt/ranges.h>
 
+#include <algorithm>
+#include <iterator>
+#include <ranges>
+
 namespace KINETO_NAMESPACE {
 
 // number of scopes affect the size of counter data binary used by
@@ -33,6 +37,14 @@ bool XpuptiScopeProfilerConfig::handleOption(
     xpuptiProfilerPerKernel_ = toBool(val);
   } else if (name == "XPUPTI_PROFILER_MAX_SCOPES"sv) {
     xpuptiProfilerMaxScopes_ = toInt64(val);
+  } else if (name == "XPUPTI_PROFILER_DEVICES"sv) {
+    const auto tokens = splitAndTrim(val, ',');
+    const auto nonEmpty = [](const std::string& tok) { return !tok.empty(); };
+    const auto toIndex = [this](const std::string& tok) { return toInt32(tok); };
+    xpuptiProfilerDevices_.clear();
+    std::ranges::copy(
+        tokens | std::views::filter(nonEmpty) | std::views::transform(toIndex),
+        std::back_inserter(xpuptiProfilerDevices_));
   } else {
     return false;
   }
@@ -53,10 +65,14 @@ void XpuptiScopeProfilerConfig::printActivityProfilerConfig(
         s,
         "Xpupti Profiler metrics : {}\n"
         "Xpupti Profiler measure per kernel : {}\n"
-        "Xpupti Profiler max scopes : {}\n",
+        "Xpupti Profiler max scopes : {}\n"
+        "Xpupti Profiler devices : {}\n",
         fmt::join(activitiesXpuptiMetrics_, ", "),
         xpuptiProfilerPerKernel_,
-        xpuptiProfilerMaxScopes_);
+        xpuptiProfilerMaxScopes_,
+        xpuptiProfilerDevices_.empty()
+            ? std::string("all")
+            : fmt::format("{}", fmt::join(xpuptiProfilerDevices_, ", ")));
   }
 }
 

libkineto/src/plugin/xpupti/XpuptiScopeProfilerConfig.h

@@ -44,6 +44,10 @@ class XpuptiScopeProfilerConfig : public AbstractConfig {
     return xpuptiProfilerMaxScopes_;
   }
 
+  const std::vector<int>& xpuptiProfilerDevices() const {
+    return xpuptiProfilerDevices_;
+  }
+
   void setClientDefaults() override {
     setDefaults();
   }
@@ -79,6 +83,10 @@ class XpuptiScopeProfilerConfig : public AbstractConfig {
   // max number of scopes to configure the profiler for.
   // this has to be set before hand to reserve space for the output
   int64_t xpuptiProfilerMaxScopes_ = 0;
+
+  // Explicit XPU device indices to profile with the scope profiler.
+  // Empty means "all available devices" (PTI auto-detect).
+  std::vector<int> xpuptiProfilerDevices_;
 };
 
 } // namespace KINETO_NAMESPACE

libkineto/test/xpupti/XpuptiScopeProfilerConfigTest.cpp

@@ -6,13 +6,16 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include "src/plugin/xpupti/XpuptiScopeProfilerApi.h"
 #include "src/plugin/xpupti/XpuptiScopeProfilerConfig.h"
 
 #include <fmt/core.h>
 #include <fmt/ranges.h>
 
 #include <gtest/gtest.h>
 
+#include <array>
+
 namespace KN = KINETO_NAMESPACE;
 
 class XpuptiScopeProfilerConfigTest : public ::testing::Test {
@@ -68,3 +71,66 @@ TEST_F(XpuptiScopeProfilerConfigTest, ScopesDefaults) {
   EXPECT_EQ(user_scopes, 10);
   EXPECT_EQ(auto_scopes, 1500);
 }
+
+TEST_F(XpuptiScopeProfilerConfigTest, DevicesDefaultEmpty) {
+  KN::Config cfg;
+  EXPECT_TRUE(cfg.parse("XPUPTI_PROFILER_METRICS = metric1"));
+  // No XPUPTI_PROFILER_DEVICES set -> empty means "all devices".
+  const KN::XpuptiScopeProfilerConfig& xpupti_cfg =
+      KN::XpuptiScopeProfilerConfig::get(cfg);
+  EXPECT_TRUE(xpupti_cfg.xpuptiProfilerDevices().empty());
+}
+
+TEST_F(XpuptiScopeProfilerConfigTest, DevicesParsedList) {
+  KN::Config cfg;
+  EXPECT_TRUE(cfg.parse("XPUPTI_PROFILER_METRICS = metric1"));
+  EXPECT_TRUE(cfg.parse("XPUPTI_PROFILER_DEVICES = 0, 2, 3"));
+  const KN::XpuptiScopeProfilerConfig& xpupti_cfg =
+      KN::XpuptiScopeProfilerConfig::get(cfg);
+  const std::vector<int> expected{0, 2, 3};
+  EXPECT_EQ(xpupti_cfg.xpuptiProfilerDevices(), expected);
+}
+
+TEST_F(XpuptiScopeProfilerConfigTest, DevicesSingle) {
+  KN::Config cfg;
+  EXPECT_TRUE(cfg.parse("XPUPTI_PROFILER_DEVICES = 1"));
+  const KN::XpuptiScopeProfilerConfig& xpupti_cfg =
+      KN::XpuptiScopeProfilerConfig::get(cfg);
+  const std::vector<int> expected{1};
+  EXPECT_EQ(xpupti_cfg.xpuptiProfilerDevices(), expected);
+}
+
+TEST_F(XpuptiScopeProfilerConfigTest, SelectDeviceHandlesSubset) {
+  // Fabricate 4 distinct opaque handles; the helper never dereferences them.
+  std::array<pti_device_handle_t, 4> handles{
+      reinterpret_cast<pti_device_handle_t>(0x10),
+      reinterpret_cast<pti_device_handle_t>(0x20),
+      reinterpret_cast<pti_device_handle_t>(0x30),
+      reinterpret_cast<pti_device_handle_t>(0x40)};
+  const std::vector<int> indices{0, 2, 3};
+  auto out = KN::selectDeviceHandles(handles, indices);
+  ASSERT_EQ(out.size(), 3u);
+  EXPECT_EQ(out[0], handles[0]);
+  EXPECT_EQ(out[1], handles[2]);
+  EXPECT_EQ(out[2], handles[3]);
+}
+
+TEST_F(XpuptiScopeProfilerConfigTest, SelectDeviceHandlesOutOfRangeThrows) {
+  std::array<pti_device_handle_t, 2> handles{
+      reinterpret_cast<pti_device_handle_t>(0x10),
+      reinterpret_cast<pti_device_handle_t>(0x20)};
+  const std::vector<int> indices{0, 5};
+  EXPECT_THROW(
+      KN::selectDeviceHandles(handles, indices),
+      std::runtime_error);
+}
+
+TEST_F(XpuptiScopeProfilerConfigTest, SelectDeviceHandlesNegativeThrows) {
+  std::array<pti_device_handle_t, 2> handles{
+      reinterpret_cast<pti_device_handle_t>(0x10),
+      reinterpret_cast<pti_device_handle_t>(0x20)};
+  const std::vector<int> indices{-1};
+  EXPECT_THROW(
+      KN::selectDeviceHandles(handles, indices),
+      std::runtime_error);
+}

libkineto/test/xpupti/XpuptiScopeProfilerTest.cpp

@@ -28,7 +28,10 @@ class XpuptiScopeProfilerTest : public ::testing::Test {
   }
 };
 
-void RunTest(std::string_view perKernel, unsigned maxScopes) {
+void RunTest(
+    std::string_view perKernel,
+    unsigned maxScopes,
+    std::string_view devices = "") {
   KN::Config cfg;
 
   std::vector<std::string_view> metrics = {
@@ -45,6 +48,10 @@ void RunTest(std::string_view perKernel, unsigned maxScopes) {
       fmt::format("XPUPTI_PROFILER_ENABLE_PER_KERNEL = {}", perKernel)));
   EXPECT_TRUE(
       cfg.parse(fmt::format("XPUPTI_PROFILER_MAX_SCOPES = {}", maxScopes)));
+  if (!devices.empty()) {
+    EXPECT_TRUE(cfg.parse(
+        fmt::format("XPUPTI_PROFILER_DEVICES = {}", devices)));
+  }
 
   std::set<KN::ActivityType> activities{
       KN::ActivityType::GPU_MEMCPY,
@@ -131,3 +138,13 @@ TEST_F(XpuptiScopeProfilerTest, PerKernelScope) {
 TEST_F(XpuptiScopeProfilerTest, UserScope) {
   RunTest("false", 159);
 }
+
+#if PTI_VERSION_AT_LEAST(0, 18)
+// Exercises the explicit device-subset path (selectDeviceHandles ->
+// ptiMetricsScopeConfigure with a device handle array). The gtest workload
+// uses a single queue/device, so requesting device 0 reproduces the same
+// activities as PerKernelScope while covering the multi-device code path.
+TEST_F(XpuptiScopeProfilerTest, PerKernelScopeExplicitDevice0) {
+  RunTest("true", 314, "0");
+}
+#endif