Pytorch backend: analysis

.github/workflows/presubmit.yml

@@ -118,6 +118,12 @@ jobs:
       run: |
         python3 -m pip install --user -U pip
         python3 -m pip install --user -U ittapi
+    - name: Install PyTorch (CPU)
+      run: |
+        python3 -m pip install --user -U pip
+        python3 -m pip install --user -U torch
+    - name: Add $HOME/.local/bin to the GITHUB_PATH
+      run: |
         echo "$HOME/.local/bin" >> "$GITHUB_PATH"
     - name: Setup remote-ssh session
       uses: owenthereal/action-upterm@v2

backends/pytorch/Makefile.am

@@ -4,7 +4,7 @@ else
   WERROR =
 endif
 
-PYTORCH_STATIC_PROBES = pytorch
+PYTORCH_STATIC_PROBES = pytorch_tracepoints
 
 PYTORCH_STATIC_PROBES_TP = $(PYTORCH_STATIC_PROBES:=.tp)
 
@@ -13,7 +13,7 @@ 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_$* > $@
+	$(RUBY) $< $(srcdir)/pytorch_events.yaml lttng_ust_pytorch > $@
 
 %.h %.c: %.tp
 	$(LTTNG_GEN_TP) $< -o $*.c -o $*.h
@@ -41,8 +41,9 @@ CLEANFILES = \
 BUILT_SOURCES = \
 	$(PYTORCH_STATIC_PROBES_INCL)
 
-# Only -ltorch_cpu is needed: at::addGlobalCallback and at::RecordFunction::name
-# are both defined in libtorch_cpu.so.
+# Only -ltorch_cpu is needed: at::addGlobalCallback, at::RecordFunction::name
+# and at::RecordFuntion::overload_name symbols must be resolved to register 
+# entry/exit callbacks at load time.
 DUMMY_TORCH_LIBS = \
 	dummy_libs/libtorch_cpu.so
 
@@ -63,24 +64,51 @@ bin_SCRIPTS = \
 
 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_LDFLAGS = $(LTTNG_UST_LIBS) -Ldummy_libs -Wl,--no-as-needed,-ltorch_cpu,--as-needed -avoid-version -module
 libTracerPytorch_la_LIBADD = libpytorchtracepoints.la
 EXTRA_libTracerPytorch_la_DEPENDENCIES = $(DUMMY_TORCH_LIBS)
+
+BTX_PYTORCH_GENERATED = \
+	btx_filter_pytorch/metababel/metababel.h \
+	btx_filter_pytorch/metababel/btx_component.h \
+	btx_filter_pytorch/metababel/btx_component.c \
+	btx_filter_pytorch/metababel/btx_upstream.h \
+	btx_filter_pytorch/metababel/btx_upstream.c \
+	btx_filter_pytorch/metababel/btx_downstream.h \
+	btx_filter_pytorch/metababel/btx_downstream.c \
+	btx_filter_pytorch/btx_main.c
+
+EXTRA_DIST += \
+	$(top_srcdir)/xprof/btx_interval_model.yaml \
+	$(srcdir)/btx_pytorch_model.yaml
+
+$(BTX_PYTORCH_GENERATED) &: $(top_srcdir)/xprof/btx_interval_model.yaml $(srcdir)/btx_pytorch_model.yaml
+	$(METABABEL) -u $(srcdir)/btx_pytorch_model.yaml -d $(top_srcdir)/xprof/btx_interval_model.yaml -t FILTER -o btx_filter_pytorch -p pytorchinterval -c interval
+
+CLEANFILES += \
+	$(BTX_PYTORCH_GENERATED)
+
+BUILT_SOURCES += \
+	$(BTX_PYTORCH_GENERATED)
+
+nodist_libPyTorchInterval_la_SOURCES = \
+	$(BTX_PYTORCH_GENERATED)
+
+libPyTorchInterval_la_SOURCES = \
+	btx_pytorchinterval_callbacks.cpp
+
+libPyTorchInterval_la_CPPFLAGS = -I$(top_srcdir)/utils -I$(top_srcdir)/utils/include -I./ -I./btx_filter_pytorch
+libPyTorchInterval_la_CFLAGS = -Wall -Wextra -Wno-unused-parameter $(WERROR) $(BABELTRACE2_CFLAGS)
+libPyTorchInterval_la_CXXFLAGS = -std=c++17 -Wall -Wextra -Wno-unused-parameter $(WERROR) $(BABELTRACE2_CFLAGS)
+libPyTorchInterval_la_LDFLAGS = $(BABELTRACE2_LIBS) -avoid-version -module
+
+bt2dir = $(pkglibdir)/bt2
+bt2_LTLIBRARIES = libPyTorchInterval.la

backends/pytorch/btx_pytorch_model.yaml

@@ -0,0 +1,53 @@
+:environment:
+  :entries:
+  - :name: hostname
+    :type: string
+:stream_classes:
+- :name: thapi_pytorch
+  :default_clock_class: {}
+  :packet_context_field_class:
+    :type: structure
+    :members:
+    - :name: cpu_id
+      :field_class:
+        :type: integer_unsigned
+        :cast_type: uint64_t
+        :field_value_range: 32
+  :event_common_context_field_class:
+    :type: structure
+    :members:
+    - :name: vpid
+      :field_class:
+        :type: integer_signed
+        :cast_type: int64_t
+        :field_value_range: 64
+    - :name: vtid
+      :field_class:
+        :type: integer_unsigned
+        :cast_type: uint64_t
+        :field_value_range: 64
+  :event_classes:
+  - :name: lttng_ust_pytorch:op_entry
+    :payload_field_class:
+      :type: structure
+      :members:
+      - :name: name
+        :field_class:
+          :cast_type: char *
+          :type: string
+      - :name: overload_name
+        :field_class:
+          :cast_type: char *
+          :type: string
+  - :name: lttng_ust_pytorch:op_exit
+    :payload_field_class:
+      :type: structure
+      :members:
+      - :name: name
+        :field_class:
+          :cast_type: char *
+          :type: string
+      - :name: overload_name
+        :field_class:
+          :cast_type: char *
+          :type: string

backends/pytorch/btx_pytorchinterval_callbacks.cpp

@@ -0,0 +1,67 @@
+#include "xprof_utils.hpp"
+#include <metababel/metababel.h>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+// PyTorch requies a per-thread LIFO stack since there is reentrace and nesting calls.
+struct data_s {
+  std::unordered_map<hpt_t, std::vector<int64_t>> entry_stack;
+};
+typedef struct data_s data_t;
+
+static void btx_initialize_component(void **usr_data) { *usr_data = new data_t; }
+
+static void btx_finalize_component(void *usr_data) { delete static_cast<data_t *>(usr_data); }
+
+// PyTorch identifies an operator by TWO strings: a schema name (e.g.
+// "aten::abs") and an overload name (e.g. "" for the default overload, "out"
+// Joining them here keeps e.g. "aten::abs" and "aten::abs.out" distinguishable
+// in the trace instead of both showing up as plain "aten::abs".
+static std::string qualified_name(const char *name, const char *overload_name) {
+  return (overload_name[0] == '\0') ? name : std::string(name) + "." + overload_name;
+}
+
+static void lttng_ust_pytorch_op_entry_callback(void *btx_handle,
+                                                void *usr_data,
+                                                int64_t ts,
+                                                const char *hostname,
+                                                int64_t vpid,
+                                                uint64_t vtid,
+                                                char * /*name*/,
+                                                char * /*overload_name*/) {
+  static_cast<data_t *>(usr_data)->entry_stack[{hostname, vpid, vtid}].push_back(ts);
+}
+
+static void lttng_ust_pytorch_op_exit_callback(void *btx_handle,
+                                               void *usr_data,
+                                               int64_t ts,
+                                               const char *hostname,
+                                               int64_t vpid,
+                                               uint64_t vtid,
+                                               char *name,
+                                               char *overload_name) {
+  auto *state = static_cast<data_t *>(usr_data);
+  auto &stack = state->entry_stack[{hostname, vpid, vtid}];
+
+  // Empty means an exit arrived with no matching entry
+  const bool err = stack.empty();
+  int64_t entry_ts = ts;
+  if (!err) {
+    entry_ts = stack.back();
+    stack.pop_back();
+  }
+
+  const std::string full_name = qualified_name(name, overload_name);
+  btx_push_message_lttng_host(btx_handle, hostname, vpid, vtid, entry_ts, BACKEND_PYTORCH,
+                              full_name.c_str(), (ts - entry_ts), err);
+}
+
+void btx_register_usr_callbacks(void *btx_handle) {
+  btx_register_callbacks_initialize_component(btx_handle, &btx_initialize_component);
+  btx_register_callbacks_finalize_component(btx_handle, &btx_finalize_component);
+
+  btx_register_callbacks_lttng_ust_pytorch_op_entry(btx_handle,
+                                                    &lttng_ust_pytorch_op_entry_callback);
+  btx_register_callbacks_lttng_ust_pytorch_op_exit(btx_handle, &lttng_ust_pytorch_op_exit_callback);
+}

backends/pytorch/include/ATen/record_function.h

@@ -1,28 +1,5 @@
-// 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
+// Minimal hand-written PyTorch's <ATen/record_function.h>.
+
 #pragma once
 
 #include <array>
@@ -55,6 +32,7 @@ struct ObserverContext {
 
 struct RecordFunction {
   const char *name() const;
+  const char *overload_name() const;
 };
 
 class RecordFunctionCallback {

backends/pytorch/pytorch_events.yaml

@@ -3,10 +3,14 @@ lttng_ust_pytorch:
   - name: op_entry
     args:
     - ["const char *", name]
+    - ["const char *", overload_name]
     fields:
     - [ctf_string, name, name]
+    - [ctf_string, overload_name, overload_name]
   - name: op_exit
     args:
     - ["const char *", name]
+    - ["const char *", overload_name]
     fields:
     - [ctf_string, name, name]
+    - [ctf_string, overload_name, overload_name]

backends/pytorch/tracer_pytorch.cpp

@@ -1,19 +1,15 @@
+#include "pytorch_tracepoints.h"
 #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());
+  tracepoint(lttng_ust_pytorch, op_entry, fn.name(), fn.overload_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());
+  tracepoint(lttng_ust_pytorch, op_exit, fn.name(), fn.overload_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)

integration_tests/backend_pytorch.bats

@@ -0,0 +1,7 @@
+bats_require_minimum_version 1.5.0
+
+@test "PyTorch: trace contains aten::empty.memory_format" {
+  iprof --backends pytorch --analysis-output ./pytorch_out.txt -- \
+    python3 ./integration_tests/pytorch_example.py
+  grep "aten::empty.memory_format" ./pytorch_out.txt
+}

integration_tests/pytorch_example.py

@@ -0,0 +1,2 @@
+import torch
+torch.empty(3)

utils/babeltrace_thapi.in

@@ -401,7 +401,7 @@ class BabeltraceParserThapi < OptionParserWithDefaultAndValidation
     on('-h', '--help', 'Prints this help') { print_help_and_exit(self, exit_code: 0) }
     on('-b', '--backends BACKENDS', Array, "Select which and how backends' need to handled.",
        'Format: backend_name[:backend_level],...',
-       default: ['mpi:3', 'omp:2', 'cl:1', 'ze:1', 'cuda:1', 'hip:1', 'cxi:4', 'itt:5'])
+       default: ['mpi:3', 'omp:2', 'cl:1', 'ze:1', 'cuda:1', 'hip:1', 'cxi:4', 'itt:5', 'pytorch:6'])
     on('--debug', default: false)
     on('--archive SESSION-NAME')
     on('--archive-session-found-file-path PATH')

utils/xprof_utils.hpp

@@ -23,6 +23,7 @@ enum backend_e {
   BACKEND_MPI = 7,
   BACKEND_CXI = 8,
   BACKEND_ITT = 9,
+  BACKEND_PYTORCH = 10,
 };
 typedef enum backend_e backend_t;
 typedef unsigned backend_level_t;
@@ -39,6 +40,7 @@ const std::unordered_map<std::string, backend_t> pretty_backend_name_g = {
     {"mpi", BACKEND_MPI},
     {"cxi", BACKEND_CXI},
     {"itt", BACKEND_ITT},
+    {"pytorch", BACKEND_PYTORCH},
 };
 
 const std::unordered_map<backend_t, backend_level_t> backend_levels_g = {
@@ -52,6 +54,7 @@ const std::unordered_map<backend_t, backend_level_t> backend_levels_g = {
     {BACKEND_MPI, 3},
     {BACKEND_CXI, 4},
     {BACKEND_ITT, 5},
+    {BACKEND_PYTORCH, 6},
 };
 
 typedef std::string thapi_metadata_t;
@@ -74,14 +77,18 @@ typedef std::tuple<hostname_t, process_id_t, thread_id_t> hpt_t;
 
 typedef std::tuple<hostname_t, process_id_t, thread_id_t, thapi_function_name> hpt_function_name_t;
 typedef std::tuple<thread_id_t, thapi_function_name> t_function_name_t;
-typedef std::tuple<hostname_t, process_id_t, thread_id_t, thapi_device_id, thapi_device_id,
+typedef std::tuple<hostname_t,
+                   process_id_t,
+                   thread_id_t,
+                   thapi_device_id,
+                   thapi_device_id,
                    thapi_function_name>
     hpt_device_function_name_t;
 typedef std::tuple<hostname_t, process_id_t, thapi_device_id> hp_device_t;
 
 typedef std::tuple<hostname_t, process_id_t, thapi_device_id, thapi_device_id> hp_dsd_t;
-typedef std::tuple<hostname_t, std::string> hi_t;                // host + NIC interface
-typedef std::tuple<hostname_t, std::string, std::string> hic_t;  // host + NIC interface + counter
+typedef std::tuple<hostname_t, std::string> hi_t;               // host + NIC interface
+typedef std::tuple<hostname_t, std::string, std::string> hic_t; // host + NIC interface + counter
 typedef std::tuple<long, long> sd_t;
 typedef std::tuple<thread_id_t, thapi_function_name, long> tfn_ts_t;
 typedef std::tuple<thapi_function_name, long> fn_ts_t;