[SYCL][XPTI] Enable PI calls notifications with arguments (#3973)

This PR adds capability to capture PI calls' arguments and pass this data to XPTI subscribers. The sample pi-trace library demonstrates how one can parse this data on the subscriber side. An example usage of this utility would be:
```bash
XPTI_TRACE_ENABLE=1 XPTI_FRAMEWORK_DISPATCHER=lib/libxptifw.so XPTI_SUBSCRIBERS=lib/libpi_trace.so ./my_app
```

Author: Alexander Batashev

sycl/include/CL/sycl/detail/common.hpp

@@ -31,6 +31,9 @@ namespace detail {
 constexpr const char *SYCL_STREAM_NAME = "sycl";
 // Stream name being used for traces generated from the SYCL plugin layer
 constexpr const char *SYCL_PICALL_STREAM_NAME = "sycl.pi";
+// Stream name being used for traces generated from PI calls. This stream
+// contains information about function arguments.
+constexpr const char *SYCL_PIDEBUGCALL_STREAM_NAME = "sycl.pi.debug";
 // Data structure that captures the user code location information using the
 // builtin capabilities of the compiler
 struct code_location {

sycl/include/CL/sycl/detail/pi.def

@@ -14,6 +14,10 @@
 // This is for convinience of doing same thing for all interfaces, e.g.
 // declare, define, initialize.
 //
+// This list is used to define PiApiKind enum, which is part of external
+// interface. To avoid ABI breakage, please, add new entries to the end of the
+// list.
+//
 // Platform
 _PI_API(piPlatformsGet)
 _PI_API(piPlatformGetInfo)

sycl/include/CL/sycl/detail/pi.hpp

@@ -184,6 +184,26 @@ uint64_t emitFunctionBeginTrace(const char *FName);
 /// \param FName The name of the PI API call
 void emitFunctionEndTrace(uint64_t CorrelationID, const char *FName);
 
+/// Notifies XPTI subscribers about PI function calls and packs call arguments.
+///
+/// \param FuncID is the API hash ID from PiApiID type trait.
+/// \param FName The name of the PI API call.
+/// \param ArgsData is a pointer to packed function call arguments.
+uint64_t emitFunctionWithArgsBeginTrace(uint32_t FuncID, const char *FName,
+                                        unsigned char *ArgsData);
+
+/// Notifies XPTI subscribers about PI function call result.
+///
+/// \param CorrelationID The correlation ID for the API call generated by the
+/// emitFunctionWithArgsBeginTrace() call.
+/// \param FuncID is the API hash ID from PiApiID type trait.
+/// \param FName The name of the PI API call.
+/// \param ArgsData is a pointer to packed function call arguments.
+/// \param Result is function call result value.
+void emitFunctionWithArgsEndTrace(uint64_t CorrelationID, uint32_t FuncID,
+                                  const char *FName, unsigned char *ArgsData,
+                                  pi_result Result);
+
 // A wrapper for passing around byte array properties
 class ByteArray {
 public:
@@ -393,3 +413,5 @@ namespace RT = cl::sycl::detail::pi;
 
 } // namespace sycl
 } // __SYCL_INLINE_NAMESPACE(cl)
+
+#undef _PI_API

sycl/include/CL/sycl/detail/type_traits.hpp

@@ -14,6 +14,7 @@
 #include <CL/sycl/detail/type_list.hpp>
 
 #include <array>
+#include <tuple>
 #include <type_traits>
 
 __SYCL_INLINE_NAMESPACE(cl) {
@@ -344,6 +345,12 @@ template <access::address_space AS, class DataT>
 using const_if_const_AS = DataT;
 #endif
 
+template <typename T> struct function_traits {};
+
+template <typename Ret, typename... Args> struct function_traits<Ret(Args...)> {
+  using ret_type = Ret;
+  using args_type = std::tuple<Args...>;
+};
 
 } // namespace detail
 } // namespace sycl

sycl/source/detail/pi.cpp

@@ -46,6 +46,8 @@ namespace detail {
 xpti_td *GSYCLGraphEvent = nullptr;
 /// Event to be used by PI layer related activities
 xpti_td *GPICallEvent = nullptr;
+/// Event to be used by PI layer calls with arguments
+xpti_td *GPIArgCallEvent = nullptr;
 /// Constants being used as placeholder until one is able to reliably get the
 /// version of the SYCL runtime
 constexpr uint32_t GMajVer = 1;
@@ -135,6 +137,42 @@ void emitFunctionEndTrace(uint64_t CorrelationID, const char *FName) {
 #endif // XPTI_ENABLE_INSTRUMENTATION
 }
 
+uint64_t emitFunctionWithArgsBeginTrace(uint32_t FuncID, const char *FuncName,
+                                        unsigned char *ArgsData) {
+  uint64_t CorrelationID = 0;
+#ifdef XPTI_ENABLE_INSTRUMENTATION
+  if (xptiTraceEnabled()) {
+    uint8_t StreamID = xptiRegisterStream(SYCL_PIDEBUGCALL_STREAM_NAME);
+    CorrelationID = xptiGetUniqueId();
+
+    xpti::function_with_args_t Payload{FuncID, FuncName, ArgsData, nullptr,
+                                       nullptr};
+
+    xptiNotifySubscribers(
+        StreamID, (uint16_t)xpti::trace_point_type_t::function_with_args_begin,
+        GPIArgCallEvent, nullptr, CorrelationID, &Payload);
+  }
+#endif
+  return CorrelationID;
+}
+
+void emitFunctionWithArgsEndTrace(uint64_t CorrelationID, uint32_t FuncID,
+                                  const char *FuncName, unsigned char *ArgsData,
+                                  pi_result Result) {
+#ifdef XPTI_ENABLE_INSTRUMENTATION
+  if (xptiTraceEnabled()) {
+    uint8_t StreamID = xptiRegisterStream(SYCL_PIDEBUGCALL_STREAM_NAME);
+
+    xpti::function_with_args_t Payload{FuncID, FuncName, ArgsData, &Result,
+                                       nullptr};
+
+    xptiNotifySubscribers(
+        StreamID, (uint16_t)xpti::trace_point_type_t::function_with_args_end,
+        GPIArgCallEvent, nullptr, CorrelationID, &Payload);
+  }
+#endif
+}
+
 void contextSetExtendedDeleter(const cl::sycl::context &context,
                                pi_context_extended_deleter func,
                                void *user_data) {
@@ -430,6 +468,14 @@ static void initializePlugins(std::vector<plugin> *Plugins) {
   GPICallEvent =
       xptiMakeEvent("PI Layer", &PIPayload, xpti::trace_algorithm_event,
                     xpti_at::active, &PiInstanceNo);
+
+  xptiInitialize(SYCL_PIDEBUGCALL_STREAM_NAME, GMajVer, GMinVer, GVerStr);
+  xpti::payload_t PIArgPayload(
+      "Plugin Interface Layer (with function arguments)");
+  uint64_t PiArgInstanceNo;
+  GPIArgCallEvent = xptiMakeEvent("PI Layer with arguments", &PIArgPayload,
+                                  xpti::trace_algorithm_event, xpti_at::active,
+                                  &PiArgInstanceNo);
 #endif
 }
 

sycl/source/detail/plugin.hpp

@@ -10,6 +10,7 @@
 #include <CL/sycl/backend_types.hpp>
 #include <CL/sycl/detail/common.hpp>
 #include <CL/sycl/detail/pi.hpp>
+#include <CL/sycl/detail/type_traits.hpp>
 #include <CL/sycl/stl.hpp>
 #include <detail/plugin_printers.hpp>
 #include <memory>
@@ -25,7 +26,62 @@ namespace sycl {
 namespace detail {
 #ifdef XPTI_ENABLE_INSTRUMENTATION
 extern xpti::trace_event_data_t *GPICallEvent;
+extern xpti::trace_event_data_t *GPIArgCallEvent;
 #endif
+
+template <PiApiKind Kind, size_t Idx, typename... Args>
+struct array_fill_helper;
+
+template <PiApiKind Kind> struct PiApiArgTuple;
+
+#define _PI_API(api)                                                           \
+  template <> struct PiApiArgTuple<PiApiKind::api> {                           \
+    using type = typename function_traits<decltype(api)>::args_type;           \
+  };
+
+#include <CL/sycl/detail/pi.def>
+#undef _PI_API
+
+template <PiApiKind Kind, size_t Idx, typename T>
+struct array_fill_helper<Kind, Idx, T> {
+  static void fill(unsigned char *Dst, T &&Arg) {
+    using ArgsTuple = typename PiApiArgTuple<Kind>::type;
+    // C-style cast is required here.
+    auto RealArg = (std::tuple_element_t<Idx, ArgsTuple>)(Arg);
+    *(std::remove_cv_t<std::tuple_element_t<Idx, ArgsTuple>> *)Dst = RealArg;
+  }
+};
+
+template <PiApiKind Kind, size_t Idx, typename T, typename... Args>
+struct array_fill_helper<Kind, Idx, T, Args...> {
+  static void fill(unsigned char *Dst, const T &&Arg, Args &&... Rest) {
+    using ArgsTuple = typename PiApiArgTuple<Kind>::type;
+    // C-style cast is required here.
+    auto RealArg = (std::tuple_element_t<Idx, ArgsTuple>)(Arg);
+    *(std::remove_cv_t<std::tuple_element_t<Idx, ArgsTuple>> *)Dst = RealArg;
+    array_fill_helper<Kind, Idx + 1, Args...>::fill(
+        Dst + sizeof(decltype(RealArg)), std::forward<Args>(Rest)...);
+  }
+};
+
+template <typename... Ts>
+constexpr size_t totalSize(const std::tuple<Ts...> &) {
+  return (sizeof(Ts) + ...);
+}
+
+template <PiApiKind Kind, typename... ArgsT>
+auto packCallArguments(ArgsT &&... Args) {
+  using ArgsTuple = typename PiApiArgTuple<Kind>::type;
+
+  constexpr size_t TotalSize = totalSize(ArgsTuple{});
+
+  std::array<unsigned char, TotalSize> ArgsData;
+  array_fill_helper<Kind, 0, ArgsT...>::fill(ArgsData.data(),
+                                             std::forward<ArgsT>(Args)...);
+
+  return ArgsData;
+}
+
 /// The plugin class provides a unified interface to the underlying low-level
 /// runtimes for the device-agnostic SYCL runtime.
 ///
@@ -85,6 +141,10 @@ class plugin {
     // the per_instance_user_data field.
     const char *PIFnName = PiCallInfo.getFuncName();
     uint64_t CorrelationID = pi::emitFunctionBeginTrace(PIFnName);
+    auto ArgsData =
+        packCallArguments<PiApiOffset>(std::forward<ArgsT>(Args)...);
+    uint64_t CorrelationIDWithArgs = pi::emitFunctionWithArgsBeginTrace(
+        static_cast<uint32_t>(PiApiOffset), PIFnName, ArgsData.data());
 #endif
     RT::PiResult R;
     if (pi::trace(pi::TraceLevel::PI_TRACE_CALLS)) {
@@ -103,6 +163,9 @@ class plugin {
 #ifdef XPTI_ENABLE_INSTRUMENTATION
     // Close the function begin with a call to function end
     pi::emitFunctionEndTrace(CorrelationID, PIFnName);
+    pi::emitFunctionWithArgsEndTrace(CorrelationIDWithArgs,
+                                     static_cast<uint32_t>(PiApiOffset),
+                                     PIFnName, ArgsData.data(), R);
 #endif
     return R;
   }

sycl/tools/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_subdirectory(sycl-ls)
+add_subdirectory(pi-trace)
 
 # TODO: move each tool in its own sub-directory
 add_executable(get_device_count_by_type get_device_count_by_type.cpp)

sycl/tools/pi-trace/CMakeLists.txt

@@ -0,0 +1,14 @@
+add_library(pi_trace SHARED pi_trace.cpp)
+target_link_libraries(pi_trace PRIVATE xptifw)
+target_include_directories(pi_trace PRIVATE "${XPTI_SOURCE_DIR}/include")
+target_include_directories(pi_trace PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/../xpti_helpers/")
+target_include_directories(pi_trace PRIVATE "${sycl_inc_dir}")
+target_include_directories(pi_trace PRIVATE "${sycl_src_dir}")
+
+if(UNIX)
+  target_link_libraries(pi_trace PRIVATE dl)
+endif()
+
+if (XPTI_ENABLE_TBB)
+  target_link_libraries(pi_trace PRIVATE tbb)
+endif()

sycl/tools/pi-trace/pi_trace.cpp

@@ -0,0 +1,84 @@
+//==----------- pi_trace.cpp.cpp -------------------------------------------==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+/// \file pi_trace.cpp
+/// A sample XPTI subscriber to demonstrate how to collect PI function call
+/// arguments.
+
+#include "xpti_trace_framework.h"
+
+#include "pi_arguments_handler.hpp"
+
+#include <detail/plugin_printers.hpp>
+
+#include <iostream>
+#include <mutex>
+#include <string>
+#include <string_view>
+#include <thread>
+
+static uint8_t GStreamID = 0;
+std::mutex GIOMutex;
+
+sycl::xpti_helpers::PiArgumentsHandler ArgHandler;
+
+// The lone callback function we are going to use to demonstrate how to attach
+// the collector to the running executable
+XPTI_CALLBACK_API void tpCallback(uint16_t trace_type,
+                                  xpti::trace_event_data_t *parent,
+                                  xpti::trace_event_data_t *event,
+                                  uint64_t instance, const void *user_data);
+
+// Based on the documentation, every subscriber MUST implement the
+// xptiTraceInit() and xptiTraceFinish() APIs for their subscriber collector to
+// be loaded successfully.
+XPTI_CALLBACK_API void xptiTraceInit(unsigned int major_version,
+                                     unsigned int minor_version,
+                                     const char *version_str,
+                                     const char *stream_name) {
+  if (std::string_view(stream_name) == "sycl.pi.arg") {
+    GStreamID = xptiRegisterStream(stream_name);
+    xptiRegisterCallback(
+        GStreamID, (uint16_t)xpti::trace_point_type_t::function_with_args_begin,
+        tpCallback);
+    xptiRegisterCallback(
+        GStreamID, (uint16_t)xpti::trace_point_type_t::function_with_args_end,
+        tpCallback);
+
+#define _PI_API(api)                                                           \
+  ArgHandler.set##_##api([](auto &&... Args) {                                 \
+    std::cout << "---> " << #api << "("                                        \
+              << "\n";                                                         \
+    sycl::detail::pi::printArgs(Args...);                                      \
+    std::cout << ") ---> ";                                                    \
+  });
+#include <CL/sycl/detail/pi.def>
+#undef _PI_API
+  }
+}
+
+XPTI_CALLBACK_API void xptiTraceFinish(const char *stream_name) {
+  // NOP
+}
+
+XPTI_CALLBACK_API void tpCallback(uint16_t TraceType,
+                                  xpti::trace_event_data_t *Parent,
+                                  xpti::trace_event_data_t *Event,
+                                  uint64_t Instance, const void *UserData) {
+  auto Type = static_cast<xpti::trace_point_type_t>(TraceType);
+  if (Type == xpti::trace_point_type_t::function_with_args_end) {
+    // Lock while we print information
+    std::lock_guard<std::mutex> Lock(GIOMutex);
+
+    const auto *Data =
+        static_cast<const xpti::function_with_args_t *>(UserData);
+
+    ArgHandler.handle(Data->function_id, Data->args_data);
+    std::cout << *static_cast<pi_result *>(Data->ret_data) << "\n";
+  }
+}