Pytorch backend

Makefile.am

@@ -3,7 +3,7 @@ pkgconfigdir = $(libdir)/pkgconfig
 pkgconfig_DATA = thapi.pc
 
 ACLOCAL_AMFLAGS = -I m4
-SUBDIRS = utils xprof sampling backends/opencl backends/ze backends/cuda backends/omp backends/hip backends/mpi backends/cxi backends/itt
+SUBDIRS = utils xprof sampling backends/opencl backends/ze backends/cuda backends/omp backends/hip backends/mpi backends/cxi backends/itt backends/pytorch
 
 EXTRA_DIST = autogen.sh README.md .valgrind/dlopen.supp thapi.pc.in
 

backends/pytorch/Makefile.am

@@ -0,0 +1,86 @@
+if STRICT
+  WERROR = -Werror
+else
+  WERROR =
+endif
+
+PYTORCH_STATIC_PROBES = pytorch
+
+PYTORCH_STATIC_PROBES_TP = $(PYTORCH_STATIC_PROBES:=.tp)
+
+PYTORCH_STATIC_PROBES_INCL = $(PYTORCH_STATIC_PROBES:=.h)
+
+PYTORCH_STATIC_PROBES_SRC = $(PYTORCH_STATIC_PROBES:=.c)
+
+$(PYTORCH_STATIC_PROBES_TP): %.tp: $(top_srcdir)/utils/gen_custom_probes.rb $(srcdir)/pytorch_events.yaml $(top_srcdir)/utils/gen_probe_base.rb
+	$(RUBY) $< $(srcdir)/pytorch_events.yaml lttng_ust_$* > $@
+
+%.h %.c: %.tp
+	$(LTTNG_GEN_TP) $< -o $*.c -o $*.h
+
+noinst_LTLIBRARIES = libpytorchtracepoints.la
+
+nodist_libpytorchtracepoints_la_SOURCES = \
+	$(PYTORCH_STATIC_PROBES_INCL) \
+	$(PYTORCH_STATIC_PROBES_SRC)
+
+libpytorchtracepoints_la_CPPFLAGS = -I$(top_srcdir)/utils -I$(top_srcdir)/utils/include -I./
+libpytorchtracepoints_la_CFLAGS = -fPIC -Wall -Wextra -Wno-unused-parameter -Wno-type-limits -Wno-sign-compare $(WERROR) $(LTTNG_UST_CFLAGS)
+libpytorchtracepoints_la_LDFLAGS = $(LTTNG_UST_LIBS)
+
+EXTRA_DIST = \
+	pytorch_events.yaml \
+	$(top_srcdir)/utils/gen_custom_probes.rb \
+	include
+
+CLEANFILES = \
+	$(PYTORCH_STATIC_PROBES_TP) \
+	$(PYTORCH_STATIC_PROBES_INCL) \
+	$(PYTORCH_STATIC_PROBES_SRC)
+
+BUILT_SOURCES = \
+	$(PYTORCH_STATIC_PROBES_INCL)
+
+# Only -ltorch_cpu is needed: at::addGlobalCallback and at::RecordFunction::name
+# are both defined in libtorch_cpu.so.
+DUMMY_TORCH_LIBS = \
+	dummy_libs/libtorch_cpu.so
+
+dummy_libs/libtorch_cpu.so:
+	$(MKDIR_P) dummy_libs
+	$(CXX) -shared -o $@ -x c++ /dev/null
+
+install-exec-hook:
+	-$(RM) $(DUMMY_TORCH_LIBS)
+	-rmdir dummy_libs
+
+clean-local:
+	-$(RM) $(DUMMY_TORCH_LIBS)
+	-rmdir dummy_libs
+
+bin_SCRIPTS = \
+	tracer_pytorch.sh
+
+lib_LTLIBRARIES = libTracerPytorch.la
+
+# pytorch.c is only built once, into libpytorchtracepoints.la above; pull it
+# in via LIBADD rather than compiling it again here (a second copy would
+# clash with the first at link time: duplicate tracepoint symbols).
+nodist_libTracerPytorch_la_SOURCES = \
+	$(PYTORCH_STATIC_PROBES_INCL)
+
+libTracerPytorch_la_SOURCES = \
+	tracer_pytorch.cpp
+
+libTracerPytorch_la_CPPFLAGS = -I$(top_srcdir)/utils -I$(top_srcdir)/utils/include -I$(srcdir)/include -I./
+# -Wno-missing-field-initializers: LTTng's own tracepoint.h (not our code)
+#   aggregate-initializes structs without naming every field; clang is
+#   stricter than gcc about this under -Wextra.
+# -Wno-unused-private-field: RecordFunctionCallback's fields exist only to
+#   match PyTorch's real class layout byte-for-byte (see the comment in
+#   include/ATen/record_function.h) -- the real libtorch reads them, we
+#   never do, by design.
+libTracerPytorch_la_CXXFLAGS = -std=c++17 -Wall -Wextra -Wno-unused-parameter -Wno-missing-field-initializers -Wno-unused-private-field $(WERROR) $(LTTNG_UST_CFLAGS)
+libTracerPytorch_la_LDFLAGS = $(LTTNG_UST_LIBS) -Ldummy_libs -ltorch_cpu -avoid-version -module
+libTracerPytorch_la_LIBADD = libpytorchtracepoints.la
+EXTRA_libTracerPytorch_la_DEPENDENCIES = $(DUMMY_TORCH_LIBS)

backends/pytorch/include/ATen/record_function.h

@@ -0,0 +1,96 @@
+// Minimal hand-written stand-in for PyTorch's <ATen/record_function.h>.
+//
+// Declares only the five names tracer_pytorch.cpp actually uses:
+//   at::RecordScope             - FUNCTION / BACKWARD_FUNCTION values
+//   at::ObserverContext         - empty base; we only ever return nullptr
+//   at::RecordFunction          - only .name(); we never construct one
+//                                 ourselves, only receive a reference from
+//                                 the real library, so no field layout is
+//                                 needed here -- name() resolves against the
+//                                 real out-of-line symbol in libtorch_cpu.
+//   at::RecordFunctionCallback  - constructed BY US and passed BY VALUE into
+//                                 addGlobalCallback. Its field layout below
+//                                 (order, types, and the scopes_ array sized
+//                                 by NUM_SCOPES) must match PyTorch's real
+//                                 class byte-for-byte, copied verbatim from
+//                                 the upstream header. If a future PyTorch
+//                                 release reorders/adds a field or changes
+//                                 RecordScope's member count, this header
+//                                 will still compile cleanly but
+//                                 addGlobalCallback will read the wrong
+//                                 bytes back -- silent corruption, not a
+//                                 build failure. Re-verify this layout
+//                                 against ATen/record_function.h whenever
+//                                 upstream PyTorch changes.
+//   at::addGlobalCallback       - registers the callback pair
+#pragma once
+
+#include <array>
+#include <cstdint>
+#include <memory>
+#include <unordered_set>
+
+namespace at {
+
+enum class RecordScope : uint8_t {
+  FUNCTION = 0,
+  BACKWARD_FUNCTION,
+  TORCHSCRIPT_FUNCTION,
+  KERNEL_FUNCTION_DTYPE,
+  CUSTOM_CLASS,
+  BUILD_FEATURE,
+  LITE_INTERPRETER,
+  USER_SCOPE,
+  STATIC_RUNTIME_OP,
+  STATIC_RUNTIME_MODEL,
+  NUM_SCOPES,
+};
+
+struct ObserverContext {
+  virtual ~ObserverContext() = default;
+
+protected:
+  ObserverContext() = default;
+};
+
+struct RecordFunction {
+  const char *name() const;
+};
+
+class RecordFunctionCallback {
+public:
+  using StartCallback = std::unique_ptr<ObserverContext> (*)(const RecordFunction &);
+  using EndCallback = void (*)(const RecordFunction &, ObserverContext *);
+
+  explicit RecordFunctionCallback(StartCallback start, EndCallback end = nullptr)
+      : start_(start), end_(end) {
+    scopes_.fill(true);
+  }
+
+  RecordFunctionCallback &scopes(const std::unordered_set<RecordScope> &scopes) {
+    if (!scopes.empty()) {
+      scopes_.fill(false);
+      for (auto sc : scopes) {
+        scopes_[static_cast<std::size_t>(sc)] = true;
+      }
+    } else {
+      scopes_.fill(true);
+    }
+    return *this;
+  }
+
+private:
+  StartCallback start_;
+  EndCallback end_;
+  double sampling_prob_ = 1.0;
+  std::array<bool, static_cast<std::size_t>(RecordScope::NUM_SCOPES)> scopes_ = {};
+  bool needs_inputs_ = false;
+  bool needs_outputs_ = false;
+  bool needs_ids_ = false;
+};
+
+using CallbackHandle = uint64_t;
+
+CallbackHandle addGlobalCallback(RecordFunctionCallback cb);
+
+} // namespace at

backends/pytorch/pytorch_events.yaml

@@ -0,0 +1,12 @@
+lttng_ust_pytorch:
+  events:
+  - name: op_entry
+    args:
+    - ["const char *", name]
+    fields:
+    - [ctf_string, name, name]
+  - name: op_exit
+    args:
+    - ["const char *", name]
+    fields:
+    - [ctf_string, name, name]

backends/pytorch/tracer_pytorch.cpp

@@ -0,0 +1,21 @@
+#include <ATen/record_function.h>
+
+#include "pytorch.h"
+
+// ENTRY: fires BEFORE the op runs. LTTng adds time + vpid/vtid via context.
+static std::unique_ptr<at::ObserverContext> on_entry(const at::RecordFunction &fn) {
+  tracepoint(lttng_ust_pytorch, op_entry, fn.name());
+  return nullptr;
+}
+
+// EXIT: fires AFTER the op returns.
+static void on_exit(const at::RecordFunction &fn, at::ObserverContext *) {
+  tracepoint(lttng_ust_pytorch, op_exit, fn.name());
+}
+
+// Auto-register at library load (works under LD_PRELOAD, no python changes).
+__attribute__((constructor)) static void tracer_pytorch_init() {
+  at::addGlobalCallback(
+      at::RecordFunctionCallback(&on_entry, &on_exit)
+          .scopes({at::RecordScope::FUNCTION, at::RecordScope::BACKWARD_FUNCTION}));
+}

backends/pytorch/tracer_pytorch.sh.in

@@ -0,0 +1,71 @@
+#!/bin/sh
+exec_prefix=@prefix@
+libdir=@libdir@
+prefix=@prefix@
+datarootdir=@datarootdir@
+datadir=@datadir@
+
+display_help() {
+  echo "Usage:"
+  echo " $(basename $0) [options] [--] <application> <application-arguments>"
+  echo "  --help                        Show this screen"
+  echo "  --version                     Print the version string"
+  exit 1
+}
+
+display_version() {
+  cat $datadir/thapi_version
+  exit 1
+}
+
+while true; do
+  case "$1" in
+    --help)
+      display_help
+      exit
+      ;;
+    --version)
+      display_version
+      exit
+      ;;
+    --)
+      shift
+      break
+      ;;
+    *) break ;;
+  esac
+done
+
+if [ "$#" -eq 0 ]; then
+  display_help
+fi
+
+torch_lib=$(python3 -c 'import torch, os; print(os.path.join(os.path.dirname(torch.__file__), "lib"))' 2>/dev/null)
+if [ -z "$torch_lib" ]; then
+  echo "$(basename $0): no 'torch' module found for python3, cannot locate libtorch_cpu.so/libc10.so" >&2
+  exit 1
+fi
+
+lttng-sessiond --daemonize --quiet
+lttng create thapi-pytorch-session
+lttng enable-channel --userspace --blocking-timeout=inf blocking-channel
+lttng add-context --userspace --channel=blocking-channel -t vpid -t vtid
+lttng enable-event --channel=blocking-channel --userspace lttng_ust_pytorch:*
+
+export LD_LIBRARY_PATH=$torch_lib:$LD_LIBRARY_PATH
+export LD_PRELOAD=$libdir/libTracerPytorch.so
+export LTTNG_UST_ALLOW_BLOCKING=1
+
+lttng start
+
+ctrl_c() {
+  lttng stop
+  lttng destroy
+  exit
+}
+
+trap ctrl_c INT
+
+"$@"
+lttng stop
+lttng destroy

configure.ac

@@ -176,6 +176,7 @@ AC_CONFIG_FILES([
 	backends/mpi/Makefile
 	backends/itt/Makefile
 	backends/cxi/Makefile
+	backends/pytorch/Makefile
 ])
 AC_CONFIG_FILES([utils/test_wrapper_thapi_text_pretty.sh], [chmod +x utils/test_wrapper_thapi_text_pretty.sh])
 AC_CONFIG_FILES([backends/opencl/tracer_opencl.sh], [chmod +x backends/opencl/tracer_opencl.sh])
@@ -186,6 +187,7 @@ AC_CONFIG_FILES([backends/omp/tracer_omp.sh], [chmod +x backends/omp/tracer_omp.
 AC_CONFIG_FILES([backends/hip/tracer_hip.sh], [chmod +x backends/hip/tracer_hip.sh])
 AC_CONFIG_FILES([backends/mpi/tracer_mpi.sh], [chmod +x backends/mpi/tracer_mpi.sh])
 AC_CONFIG_FILES([backends/itt/tracer_itt.sh], [chmod +x backends/itt/tracer_itt.sh])
+AC_CONFIG_FILES([backends/pytorch/tracer_pytorch.sh], [chmod +x backends/pytorch/tracer_pytorch.sh])
 AC_CONFIG_FILES([utils/babeltrace_thapi], [chmod +x utils/babeltrace_thapi])
 AC_CONFIG_FILES([xprof/xprof.rb], [chmod +x xprof/xprof.rb])
 AC_CONFIG_FILES([xprof/iprof:xprof/xprof_wrapper.rb.in], [chmod +x xprof/iprof])

xprof/xprof.rb.in

@@ -652,6 +652,13 @@ def enable_events_itt(channel_name, tracing_mode: 'default', profiling: true)
   LOGGER.debug("Tracing mode #{tracing_mode} similar to default for itt") unless tracing_mode == 'default'
 end
 
+def enable_events_pytorch(channel_name, tracing_mode: 'default', profiling: false)
+  lttng_enable = "lttng enable-event --userspace --session=#{lttng_session_uuid} --channel=#{channel_name}"
+  exec("#{lttng_enable} lttng_ust_pytorch:*")
+  LOGGER.debug('Profiling is ignored for pytorch') if profiling
+  LOGGER.debug("Tracing mode #{tracing_mode} similar to default for pytorch") unless tracing_mode == 'default'
+end
+
 def enable_events_metadata(channel_name, tracing_mode: 'default', profiling: true)
   lttng_enable = "lttng enable-event --userspace --session=#{lttng_session_uuid} --channel=#{channel_name}"
   LOGGER.debug('Profiling is ignored for metadata') if profiling
@@ -958,6 +965,23 @@ def all_env_tracers(usr_binary)
     h['INTEL_LIBITTNOTIFY64'] = File.join(LIBDIR, 'libittnotify.so')
   end
 
+  if OPTIONS[:'backend-names'].include?('pytorch')
+    # Supported calling: iprof -- python3 model.py (torch env is active).
+    # Unsupported calling: iprof -- venv/bin/python3 module.py (non-active torch env)
+    # TODO: We need env vars to support the second calling convention.
+    torch_file = exec('python3 -c "import torch; print(torch.__file__)"', debug: false).strip
+    torch_lib = torch_file.empty? ? '' : File.join(File.dirname(torch_file), 'lib')
+    if torch_lib.empty?
+      LOGGER.warn('No torch module found for python3, pytorch backend will not be enabled')
+    else
+      backends << 'pytorch'
+      # LibTracerPytorch.so calls at::addGlobalCallback() and RecordFunction::name()
+      # both implemented in libtorch_cpu.so. The lib is reachable via LD_LIBRARY_PATH.
+      h[%w[LD_LIBRARY_PATH prepend]] << torch_lib
+      h[%w[LD_PRELOAD prepend]] << File.join(LIBDIR, 'libTracerPytorch.so')
+    end
+  end
+
   # Sample
   if SamplingDaemon.active?
     LOGGER.debug('Sampling Enabled')