[xpupti] Collect peer-to-peer memory copies

libkineto/include/MetadataFieldCatalog.h

@@ -187,6 +187,7 @@ inline constexpr MetadataField<uint64_t> kCorrelation{"correlation"};
 inline constexpr MetadataField<int64_t> kDevice{"device"};
 inline constexpr MetadataField<uint64_t> kEngineIndex{"engine_index"};
 inline constexpr MetadataField<uint64_t> kEngineOrdinal{"engine_ordinal"};
+inline constexpr MetadataField<int64_t> kFromDevice{"fromDevice"};
 inline constexpr MetadataField<uint64_t> kKernelId{"kernel_id"};
 inline constexpr MetadataField<double> kMemoryBandwidthGbps{
     "memory bandwidth (GB/s)"};
@@ -203,4 +204,5 @@ inline constexpr MetadataField<uint64_t> kSyclInvocationId{
     "sycl_invocation_id"};
 inline constexpr MetadataField<uint64_t> kSyclNodeId{"sycl_node_id"};
 inline constexpr MetadataField<uint64_t> kSyclQueue{"sycl queue"};
+inline constexpr MetadataField<int64_t> kToDevice{"toDevice"};
 } // namespace libkineto::XpuMetadataFields

libkineto/src/plugin/xpupti/XpuptiActivityApi.cpp

@@ -15,8 +15,8 @@
 #include <cstdlib>
 #include <filesystem>
 #include <map>
-#include <optional>
 #include <ostream>
+#include <span>
 #include <system_error>
 
 #include <pti/pti.h>
@@ -203,32 +203,39 @@ void warnIfIttNotifyLibInvalid() noexcept {
   }
 }
 
-// The PTI view an activity type is collected through, if it has one of its own.
-// Shared by enable and disable so the two cannot drift apart.
-constexpr std::optional<pti_view_kind> findPtiViewKind(ActivityType activity) {
+template <pti_view_kind... kViews>
+constexpr pti_view_kind kViewKinds[] = {kViews...};
+
+// The PTI views an activity type is collected through, if it has any of its
+// own. Shared by enable and disable so the two cannot drift apart.
+constexpr std::span<const pti_view_kind> findPtiViewKinds(
+    ActivityType activity) {
   switch (activity) {
     case ActivityType::GPU_MEMCPY:
-      return PTI_VIEW_DEVICE_GPU_MEM_COPY;
+      // PTI reports a peer-to-peer copy only through its own view.
+      return kViewKinds<
+          PTI_VIEW_DEVICE_GPU_MEM_COPY,
+          PTI_VIEW_DEVICE_GPU_MEM_COPY_P2P>;
     case ActivityType::GPU_MEMSET:
-      return PTI_VIEW_DEVICE_GPU_MEM_FILL;
+      return kViewKinds<PTI_VIEW_DEVICE_GPU_MEM_FILL>;
     case ActivityType::CONCURRENT_KERNEL:
-      return PTI_VIEW_DEVICE_GPU_KERNEL;
+      return kViewKinds<PTI_VIEW_DEVICE_GPU_KERNEL>;
     case ActivityType::EXTERNAL_CORRELATION:
-      return PTI_VIEW_EXTERNAL_CORRELATION;
+      return kViewKinds<PTI_VIEW_EXTERNAL_CORRELATION>;
     case ActivityType::XPU_RUNTIME:
-      return PTI_VIEW_RUNTIME_API;
+      return kViewKinds<PTI_VIEW_RUNTIME_API>;
     case ActivityType::XPU_DRIVER:
-      return PTI_VIEW_DRIVER_API;
+      return kViewKinds<PTI_VIEW_DRIVER_API>;
     case ActivityType::OVERHEAD:
-      return PTI_VIEW_COLLECTION_OVERHEAD;
+      return kViewKinds<PTI_VIEW_COLLECTION_OVERHEAD>;
     case ActivityType::XPU_SYNC:
-      return PTI_VIEW_DEVICE_SYNCHRONIZATION;
+      return kViewKinds<PTI_VIEW_DEVICE_SYNCHRONIZATION>;
     case ActivityType::COLLECTIVE_COMM:
-      return PTI_VIEW_COMMUNICATION;
+      return kViewKinds<PTI_VIEW_COMMUNICATION>;
     default:
       // XPU_SCOPE_PROFILER among others: enabled by
       // XpuptiScopeProfilerSession's constructor, not by a view.
-      return std::nullopt;
+      return {};
   }
 }
 } // namespace
@@ -240,11 +247,13 @@ void XpuptiActivityApi::enableXpuptiActivities(
 
   externalCorrelationEnabled_ = false;
   selected_activities.forEach([this](ActivityType activity) {
-    const auto view = findPtiViewKind(activity);
-    if (not view) {
+    const auto views = findPtiViewKinds(activity);
+    if (views.empty()) {
       return;
     }
-    XPUPTI_CALL(ptiViewEnable(*view));
+    for (const auto view : views) {
+      XPUPTI_CALL(ptiViewEnable(view));
+    }
 
     // What enabling a view takes on top of the view itself.
     switch (activity) {
@@ -270,8 +279,8 @@ void XpuptiActivityApi::enableXpuptiActivities(
 void XpuptiActivityApi::disablePtiActivities(
     const ActivityTypeMask& selected_activities) {
   selected_activities.forEach([](ActivityType activity) {
-    if (const auto view = findPtiViewKind(activity)) {
-      XPUPTI_CALL(ptiViewDisable(*view));
+    for (const auto view : findPtiViewKinds(activity)) {
+      XPUPTI_CALL(ptiViewDisable(view));
     }
   });
   externalCorrelationEnabled_ = false;

libkineto/src/plugin/xpupti/XpuptiActivityHandlers.cpp

@@ -207,7 +207,9 @@ void XpuptiActivityProfilerSession::handleRuntimeKernelMemcpyMemsetActivities(
       std::is_same_v<pti_view_memory_record_type, pti_view_record_api_t>;
   constexpr bool handleKernelActivities =
       std::is_same_v<pti_view_memory_record_type, pti_view_record_kernel_t>;
-  constexpr bool handleMemcpyActivities =
+  constexpr bool handleMemcpyP2PActivities =
+      std::is_same_v<pti_view_memory_record_type, pti_view_record_memcpy_p2p_t>;
+  constexpr bool handleMemcpyActivities = handleMemcpyP2PActivities ||
       std::is_same_v<pti_view_memory_record_type, pti_view_record_memcpy_t>;
   constexpr bool handleMemsetActivities =
       std::is_same_v<pti_view_memory_record_type, pti_view_record_memfill_t>;
@@ -257,7 +259,17 @@ void XpuptiActivityProfilerSession::handleRuntimeKernelMemcpyMemsetActivities(
     trace_activity->device = activity->_process_id;
     trace_activity->resource = activity->_thread_id;
   } else {
-    trace_activity->device = getDeviceIdxFromUUID(activity->_device_uuid);
+    if constexpr (handleMemcpyP2PActivities) {
+      // The copy is submitted to, and runs on an engine of, the source device.
+      trace_activity->device = getDeviceIdxFromUUID(activity->_src_uuid);
+      trace_activity->addMetadata(
+          XpuFields::kFromDevice, static_cast<int64_t>(trace_activity->device));
+      trace_activity->addMetadata(
+          XpuFields::kToDevice,
+          static_cast<int64_t>(getDeviceIdxFromUUID(activity->_dst_uuid)));
+    } else {
+      trace_activity->device = getDeviceIdxFromUUID(activity->_device_uuid);
+    }
     trace_activity->resource = activity->_sycl_queue_id;
 
     if constexpr (handleKernelActivities) {
@@ -612,6 +624,12 @@ void XpuptiActivityProfilerSession::handlePtiActivity(
           reinterpret_cast<const pti_view_record_memcpy_t*>(record),
           logger);
       break;
+    case PTI_VIEW_DEVICE_GPU_MEM_COPY_P2P:
+      handleRuntimeKernelMemcpyMemsetActivities(
+          ActivityType::GPU_MEMCPY,
+          reinterpret_cast<const pti_view_record_memcpy_p2p_t*>(record),
+          logger);
+      break;
     case PTI_VIEW_DEVICE_GPU_MEM_FILL:
       handleRuntimeKernelMemcpyMemsetActivities(
           ActivityType::GPU_MEMSET,

libkineto/src/plugin/xpupti/XpuptiActivityProfilerSession.h

@@ -44,11 +44,13 @@ using DeviceUUIDsT = std::array<unsigned char, 16>;
 inline constexpr bool kReportsEngineIds = true;
 using pti_view_record_kernel_t = pti_view_record_kernel_v2;
 using pti_view_record_memcpy_t = pti_view_record_memory_copy_v2;
+using pti_view_record_memcpy_p2p_t = pti_view_record_memory_copy_p2p_v2;
 using pti_view_record_memfill_t = pti_view_record_memory_fill_v2;
 #else
 inline constexpr bool kReportsEngineIds = false;
 using pti_view_record_kernel_t = pti_view_record_kernel;
 using pti_view_record_memcpy_t = pti_view_record_memory_copy;
+using pti_view_record_memcpy_p2p_t = pti_view_record_memory_copy_p2p;
 using pti_view_record_memfill_t = pti_view_record_memory_fill;
 #endif
 

libkineto/test/xpupti/XpuptiActivityHandlersTest.cpp

@@ -37,6 +37,9 @@ static_assert(
 static_assert(
     offsetof(pti_view_record_memory_fill, _view_kind) == 0,
     "base record must be first member");
+static_assert(
+    offsetof(pti_view_record_memory_copy_p2p, _view_kind) == 0,
+    "base record must be first member");
 
 // Mock XpuptiActivityApi that delivers hand-crafted PTI records
 // through the virtual processActivities without needing PTI runtime.
@@ -652,6 +655,44 @@ TEST_F(XpuptiActivityHandlersTest, MemcpyAndMemsetRecordsAreFlowDestinations) {
   EXPECT_FALSE(memset_activity.flowStart());
 }
 
+TEST_F(XpuptiActivityHandlersTest, PeerToPeerMemcpyIsGpuMemcpy) {
+  constexpr uint32_t kCorrelationId = 53;
+
+  KN::pti_view_record_memcpy_p2p_t p2p_record{};
+  p2p_record._view_kind._view_kind = PTI_VIEW_DEVICE_GPU_MEM_COPY_P2P;
+  p2p_record._name = "zeCommandListAppendMemoryCopy(D2D - P2P)";
+  p2p_record._start_timestamp = 200;
+  p2p_record._end_timestamp = 240;
+  p2p_record._thread_id = 7;
+  p2p_record._correlation_id = kCorrelationId;
+  p2p_record._sycl_queue_id = 3;
+  p2p_record._mem_op_id = 5;
+  p2p_record._bytes = 4096;
+  p2p_record._memcpy_type = PTI_VIEW_MEMCPY_TYPE_D2D;
+  p2p_record._mem_src = PTI_VIEW_MEMORY_TYPE_DEVICE;
+  p2p_record._mem_dst = PTI_VIEW_MEMORY_TYPE_DEVICE;
+
+  mockApi_.records.push_back(&p2p_record._view_kind);
+
+  auto traceBuffer = processAndGetTrace();
+  ASSERT_EQ(traceBuffer->activities.size(), 1);
+
+  auto& activity = *traceBuffer->activities[0];
+  EXPECT_EQ(activity.type(), ActivityType::GPU_MEMCPY);
+  EXPECT_EQ(activity.name(), "Memcpy D2D (Device -> Device)");
+  EXPECT_EQ(activity.resourceId(), 3);
+  EXPECT_EQ(activity.flowId(), kCorrelationId);
+  EXPECT_FALSE(activity.flowStart());
+  EXPECT_EQ(
+      activity.getMetadataValue(XpuMetadataFields::kBytes),
+      std::optional<uint64_t>{4096});
+  EXPECT_EQ(
+      activity.getMetadataValue(XpuMetadataFields::kFromDevice),
+      std::optional<int64_t>{activity.deviceId()});
+  EXPECT_TRUE(
+      activity.getMetadataValue(XpuMetadataFields::kToDevice).has_value());
+}
+
 // --- Mixed dispatch test ---
 
 TEST_F(XpuptiActivityHandlersTest, MixedCommunicationAndSynchronization) {

libkineto/test/xpupti/XpuptiProfilerTest.cpp

@@ -9,12 +9,17 @@
 #include "XpuptiTestUtilities.h"
 
 #include "include/GenericTraceActivity.h"
+#include "include/MetadataFieldCatalog.h"
 #include "include/libkineto.h"
 #include "src/plugin/xpupti/XpuptiActivityProfiler.h"
+#include "test/xpupti/compute/XpuptiPeerCopyCompute.h"
 
+#include <fmt/ranges.h>
 #include <gtest/gtest.h>
 #include <nlohmann/json.hpp>
 
+#include <chrono>
+#include <optional>
 #include <set>
 #include <unordered_set>
 
@@ -313,3 +318,47 @@ TEST(XpuptiProfilerTest, GpuUserAnnotationLinkedActivityRetry) {
       runGpuUserAnnotationCase(LinkedActivityMode::MissFirstThenLinked);
   expectTwoUserAnnotations(*pBuffer);
 }
+
+TEST(XpuptiProfilerTest, PeerToPeerMemcpy) {
+  KN::XPUActivityProfiler profiler;
+  auto pSession =
+      profiler.configure({KN::ActivityType::GPU_MEMCPY}, KN::Config{});
+
+  const auto now = [] {
+    return std::chrono::duration_cast<std::chrono::nanoseconds>(
+               std::chrono::system_clock::now().time_since_epoch())
+        .count();
+  };
+  const int64_t startTime = now();
+  pSession->start();
+  const bool copied = CopyPeerToPeerOnXpu(1024);
+  pSession->stop();
+  const int64_t endTime = now();
+  if (!copied) {
+    GTEST_SKIP() << "needs two GPUs with peer access";
+  }
+
+  TestActivityLogger logger;
+  pSession->processTrace(
+      logger,
+      [](int32_t) -> const KN::ITraceActivity* { return nullptr; },
+      startTime,
+      endTime);
+  EXPECT_TRUE(pSession->errors().empty())
+      << fmt::format("{}", fmt::join(pSession->errors(), ","));
+
+  auto pBuffer = pSession->getTraceBuffer();
+  ASSERT_EQ(pBuffer->activities.size(), 1);
+  const auto& activity = *pBuffer->activities[0];
+  EXPECT_EQ(activity.type(), KN::ActivityType::GPU_MEMCPY);
+  EXPECT_EQ(
+      activity.getMetadataValue(libkineto::XpuMetadataFields::kBytes),
+      std::optional<uint64_t>{1024});
+  const auto fromDevice =
+      activity.getMetadataValue(libkineto::XpuMetadataFields::kFromDevice);
+  const auto toDevice =
+      activity.getMetadataValue(libkineto::XpuMetadataFields::kToDevice);
+  EXPECT_EQ(fromDevice, std::optional<int64_t>{activity.deviceId()});
+  ASSERT_TRUE(toDevice.has_value());
+  EXPECT_NE(*fromDevice, *toDevice);
+}

libkineto/test/xpupti/XpuptiTestUtilities.cpp

@@ -202,34 +202,6 @@ static unsigned acceptOverheadActivities(
   return count;
 }
 
-class TestActivityLogger : public KN::ActivityLogger {
-  void handleDeviceInfo(
-      [[maybe_unused]] const KN::DeviceInfo& info,
-      [[maybe_unused]] int64_t time) override {}
-  void handleResourceInfo(
-      [[maybe_unused]] const KN::ResourceInfo& info,
-      [[maybe_unused]] int64_t time) override {}
-  void handleOverheadInfo(
-      [[maybe_unused]] const KN::ActivityLogger::OverheadInfo& info,
-      [[maybe_unused]] int64_t time) override {}
-  void handleTraceSpan([[maybe_unused]] const KN::TraceSpan& span) override {}
-  void handleActivity(
-      [[maybe_unused]] const KN::ITraceActivity& activity) override {}
-  void handleGenericActivity(
-      [[maybe_unused]] const KN::GenericTraceActivity& activity) override {}
-  void handleTraceStart(
-      [[maybe_unused]] const std::unordered_map<std::string, std::string>&
-          metadata,
-      [[maybe_unused]] const std::string& device_properties) override {}
-  void finalizeMemoryTrace(
-      [[maybe_unused]] const std::string&,
-      [[maybe_unused]] const KN::Config&) override {}
-  void finalizeTrace(
-      [[maybe_unused]] const KN::Config& config,
-      [[maybe_unused]] std::unique_ptr<KN::ActivityBuffers> buffers,
-      [[maybe_unused]] int64_t endTime) override {}
-};
-
 std::pair<
     std::unique_ptr<KN::IActivityProfilerSession>,
     std::unique_ptr<KN::CpuTraceBuffer>>

libkineto/test/xpupti/XpuptiTestUtilities.h

@@ -7,13 +7,44 @@
  */
 
 #include "include/output_base.h"
+#include "src/ActivityBuffers.h"
 
 #include <functional>
 
 namespace KN = KINETO_NAMESPACE;
 
 bool IsEnvVerbose();
 
+// Discards everything: tests inspect the session's trace buffer instead.
+class TestActivityLogger : public KN::ActivityLogger {
+ public:
+  void handleDeviceInfo(
+      [[maybe_unused]] const KN::DeviceInfo& info,
+      [[maybe_unused]] int64_t time) override {}
+  void handleResourceInfo(
+      [[maybe_unused]] const KN::ResourceInfo& info,
+      [[maybe_unused]] int64_t time) override {}
+  void handleOverheadInfo(
+      [[maybe_unused]] const KN::ActivityLogger::OverheadInfo& info,
+      [[maybe_unused]] int64_t time) override {}
+  void handleTraceSpan([[maybe_unused]] const KN::TraceSpan& span) override {}
+  void handleActivity(
+      [[maybe_unused]] const KN::ITraceActivity& activity) override {}
+  void handleGenericActivity(
+      [[maybe_unused]] const KN::GenericTraceActivity& activity) override {}
+  void handleTraceStart(
+      [[maybe_unused]] const std::unordered_map<std::string, std::string>&
+          metadata,
+      [[maybe_unused]] const std::string& device_properties) override {}
+  void finalizeMemoryTrace(
+      [[maybe_unused]] const std::string&,
+      [[maybe_unused]] const KN::Config&) override {}
+  void finalizeTrace(
+      [[maybe_unused]] const KN::Config& config,
+      [[maybe_unused]] std::unique_ptr<KN::ActivityBuffers> buffers,
+      [[maybe_unused]] int64_t endTime) override {}
+};
+
 std::pair<
     std::unique_ptr<KN::IActivityProfilerSession>,
     std::unique_ptr<KN::CpuTraceBuffer>>

libkineto/test/xpupti/compute/CMakeLists.txt

@@ -20,7 +20,9 @@ endif()
 # it without seeing any SYCL. ComputeOnXpu is exported via the API macro in the
 # header (dllexport on Windows / default visibility on ELF) - no need for
 # WINDOWS_EXPORT_ALL_SYMBOLS or an exe<->dll symbol cycle.
-add_library(xpupti_compute SHARED XpuptiScopeProfilerCompute.cpp)
+add_library(xpupti_compute SHARED
+  XpuptiScopeProfilerCompute.cpp
+  XpuptiPeerCopyCompute.cpp)
 target_compile_options(xpupti_compute PRIVATE -fsycl)
 target_link_options(xpupti_compute PRIVATE -fsycl)