pretty printing

backends/pytorch/Makefile.am

@@ -41,9 +41,7 @@ CLEANFILES = \
 BUILT_SOURCES = \
 	$(PYTORCH_STATIC_PROBES_INCL)
 
-# 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.
+# -ltorch_cpu resolves at::addGlobalCallback, at::RecordFunction::name, at::RecordFuntion::overload_name.
 DUMMY_TORCH_LIBS = \
 	dummy_libs/libtorch_cpu.so
 
@@ -99,16 +97,87 @@ CLEANFILES += \
 BUILT_SOURCES += \
 	$(BTX_PYTORCH_GENERATED)
 
-nodist_libPyTorchInterval_la_SOURCES = \
+nodist_libPYTORCHInterval_la_SOURCES = \
 	$(BTX_PYTORCH_GENERATED)
 
-libPyTorchInterval_la_SOURCES = \
+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
+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
+bt2_LTLIBRARIES = libPYTORCHInterval.la
+
+#
+# ----------------------------------------------------------------------
+# Enable pretty.print in babeltrace_thapi
+# ----------------------------------------------------------------------
+#
+
+babeltrace_pytorch_lib.rb: $(srcdir)/gen_babeltrace_pytorch_lib.rb $(top_srcdir)/utils/gen_babeltrace_lib_helper.rb $(top_srcdir)/utils/gen_probe_base.rb $(srcdir)/btx_pytorch_model.yaml
+	SRC_DIR=$(srcdir) $(RUBY) $< > $@
+
+EXTRA_DIST += \
+	gen_babeltrace_pytorch_lib.rb
+
+CLEANFILES += \
+	babeltrace_pytorch_lib.rb
+
+data_DATA = \
+	babeltrace_pytorch_lib.rb
+
+#
+# ----------------------------------------------------------------------
+# Tests
+# ----------------------------------------------------------------------
+#
+
+TRACE_COMMON = \
+	tests/basic_op_name-AA.thapi_text_pretty \
+	tests/basic_op_overload_name-AA.thapi_text_pretty \
+	tests/nested_ops_simple-ABBA.thapi_text_pretty \
+	tests/deeply_nested_ops-ABCCBA.thapi_text_pretty \
+	tests/multithreaded_ops-AABB.thapi_text_pretty \
+	tests/reentrant_op-AAAA.thapi_text_pretty
+
+BTX_PYTORCH_GENERATED_SOURCE_TEST = \
+	btx_source_pytorch_test/metababel/metababel.h \
+	btx_source_pytorch_test/metababel/btx_component.h \
+	btx_source_pytorch_test/metababel/btx_component.c \
+	btx_source_pytorch_test/metababel/btx_downstream.h \
+	btx_source_pytorch_test/metababel/btx_downstream.c \
+	btx_source_pytorch_test/btx_main.c
+
+$(abs_builddir)/btx_pytorch_model.yaml: $(srcdir)/btx_pytorch_model.yaml
+	cp $< $@
+
+$(BTX_PYTORCH_GENERATED_SOURCE_TEST) &: $(abs_builddir)/btx_pytorch_model.yaml
+	$(METABABEL) -d $(abs_builddir)/btx_pytorch_model.yaml -t SOURCE -o btx_source_pytorch_test -p pytorchtest -c test
+
+noinst_LTLIBRARIES += libtestpytorchsource.la
+nodist_libtestpytorchsource_la_SOURCES = $(BTX_PYTORCH_GENERATED_SOURCE_TEST)
+libtestpytorchsource_la_CPPFLAGS = -I$(top_srcdir)/utils -I$(top_srcdir)/utils/include -I./btx_source_pytorch_test/
+libtestpytorchsource_la_CFLAGS = -fPIC -Wall -Wextra -Wno-unused-parameter $(WERROR) $(BABELTRACE2_CFLAGS)
+
+TRACE_OUT = $(TRACE_COMMON:.thapi_text_pretty=.bt_text_pretty)
+
+TESTS = $(TRACE_COMMON)
+
+TEST_EXTENSIONS = .thapi_text_pretty
+THAPI_TEXT_PRETTY_LOG_COMPILER = $(top_builddir)/utils/test_wrapper_thapi_text_pretty.sh
+AM_THAPI_TEXT_PRETTY_LOG_FLAGS = backends/pytorch pytorch ./btx_source_pytorch_test/ .libs/libtestpytorchsource.a
+
+check_DATA = \
+	$(TRACE_COMMON) \
+	$(TRACE_OUT)
+
+EXTRA_DIST += \
+	$(TRACE_COMMON) \
+	$(TRACE_OUT)
+
+CLEANFILES += \
+	$(BTX_PYTORCH_GENERATED_SOURCE_TEST) \
+	$(abs_builddir)/btx_pytorch_model.yaml

backends/pytorch/btx_pytorchinterval_callbacks.cpp

@@ -1,4 +1,5 @@
 #include "xprof_utils.hpp"
+#include <cassert>
 #include <metababel/metababel.h>
 #include <string>
 #include <unordered_map>
@@ -14,10 +15,8 @@ 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".
+// Combine schema name and overload name (e.g. "aten::abs" + "out") so
+// different overloads don't collapse into the same trace entry.
 static std::string qualified_name(const char *name, const char *overload_name) {
   return (overload_name[0] == '\0') ? name : std::string(name) + "." + overload_name;
 }
@@ -28,8 +27,8 @@ static void lttng_ust_pytorch_op_entry_callback(void *btx_handle,
                                                 const char *hostname,
                                                 int64_t vpid,
                                                 uint64_t vtid,
-                                                char * /*name*/,
-                                                char * /*overload_name*/) {
+                                                char *name,
+                                                char *overload_name) {
   static_cast<data_t *>(usr_data)->entry_stack[{hostname, vpid, vtid}].push_back(ts);
 }
 
@@ -44,17 +43,12 @@ static void lttng_ust_pytorch_op_exit_callback(void *btx_handle,
   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();
-  }
+  assert(!stack.empty() && "op_exit with no matching op_entry");
+  const int64_t 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);
+                              qualified_name(name, overload_name).c_str(), (ts - entry_ts), false);
 }
 
 void btx_register_usr_callbacks(void *btx_handle) {

backends/pytorch/gen_babeltrace_pytorch_lib.rb

@@ -0,0 +1,2 @@
+require_relative '../../utils/gen_babeltrace_lib_helper'
+add_babeltrace_event_callbacks(nil, File.join(SRC_DIR, 'btx_pytorch_model.yaml'))

backends/pytorch/tests/README.md

@@ -0,0 +1,10 @@
+# PyTorch Backend Tests
+
+| Test | Feature | Description |
+|------|---------|-------------|
+| **basic_op_name-AA** | Basic tracing | Validates tracing of a simple PyTorch operation with empty overload name. |
+| **basic_op_overload_name-AA** | Overload name handling | Validates that operations with overload names create qualified names (e.g., `aten::add.Tensor`). |
+| **nested_ops_simple-ABBA** | LIFO stack (ABBA) | Validates LIFO stack correctly matches nested operations in ABBA order (outer→inner→inner exit→outer exit). |
+| **deeply_nested_ops-ABCCBA** | Stack depth handling | Validates LIFO stack correctly handles three levels of nesting (forward→linear→matmul). |
+| **multithreaded_ops-AABB** | Per-thread isolation | Validates each thread maintains its own independent LIFO stack without cross-thread interference. |
+| **reentrant_op-AAAA** | Reentrancy | Validates the LIFO stack correctly handles the same op re-entering itself (aten::norm.Scalar calling aten::norm.Scalar internally). |

backends/pytorch/tests/basic_op_name-AA.bt_text_pretty

@@ -0,0 +1 @@
+lttng:host: { hostname = "testnode", vpid = 0, vtid = 1, ts = 1704110400000000000, backend = 10 }, { name = "aten::relu", dur = 1000000000, err = false }

backends/pytorch/tests/basic_op_name-AA.thapi_text_pretty

@@ -0,0 +1,2 @@
+12:00:00.000000000 - testnode - vpid: 0, vtid: 1 - lttng_ust_pytorch:op_entry: { name: "aten::relu", overload_name: "" }
+12:00:01.000000000 - testnode - vpid: 0, vtid: 1 - lttng_ust_pytorch:op_exit: { name: "aten::relu", overload_name: "" }

backends/pytorch/tests/basic_op_overload_name-AA.bt_text_pretty

@@ -0,0 +1 @@
+lttng:host: { hostname = "testnode", vpid = 0, vtid = 1, ts = 1704110400000000000, backend = 10 }, { name = "aten::add.Tensor", dur = 1000000000, err = false }

backends/pytorch/tests/basic_op_overload_name-AA.thapi_text_pretty

@@ -0,0 +1,2 @@
+12:00:00.000000000 - testnode - vpid: 0, vtid: 1 - lttng_ust_pytorch:op_entry: { name: "aten::add", overload_name: "Tensor" }
+12:00:01.000000000 - testnode - vpid: 0, vtid: 1 - lttng_ust_pytorch:op_exit: { name: "aten::add", overload_name: "Tensor" }

backends/pytorch/tests/deeply_nested_ops-ABCCBA.bt_text_pretty

@@ -0,0 +1,3 @@
+lttng:host: { hostname = "testhost", vpid = 1, vtid = 1, ts = 1704110402000000000, backend = 10 }, { name = "aten::mm", dur = 1000000000, err = false }
+lttng:host: { hostname = "testhost", vpid = 1, vtid = 1, ts = 1704110401000000000, backend = 10 }, { name = "aten::matmul", dur = 3000000000, err = false }
+lttng:host: { hostname = "testhost", vpid = 1, vtid = 1, ts = 1704110400000000000, backend = 10 }, { name = "aten::addmm", dur = 5000000000, err = false }

backends/pytorch/tests/deeply_nested_ops-ABCCBA.thapi_text_pretty

@@ -0,0 +1,6 @@
+12:00:00.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_entry: { name: "aten::addmm", overload_name: "" }
+12:00:01.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_entry: { name: "aten::matmul", overload_name: "" }
+12:00:02.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_entry: { name: "aten::mm", overload_name: "" }
+12:00:03.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_exit: { name: "aten::mm", overload_name: "" }
+12:00:04.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_exit: { name: "aten::matmul", overload_name: "" }
+12:00:05.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_exit: { name: "aten::addmm", overload_name: "" }

backends/pytorch/tests/exit_without_entry-A.bt_text_pretty

@@ -0,0 +1 @@
+lttng:host: { hostname = "testhost", vpid = 1, vtid = 1, ts = 1704110400000000000, backend = 10 }, { name = "aten::sigmoid", dur = 0, err = true }

backends/pytorch/tests/exit_without_entry-A.thapi_text_pretty

@@ -0,0 +1,2 @@
+12:00:00.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_exit: { name: "aten::sigmoid", overload_name: "" }
+

backends/pytorch/tests/multithreaded_ops-AABB.bt_text_pretty

@@ -0,0 +1,2 @@
+lttng:host: { hostname = "testhost", vpid = 1, vtid = 1, ts = 1704110400000000000, backend = 10 }, { name = "aten::relu", dur = 1000000000, err = false }
+lttng:host: { hostname = "testhost", vpid = 1, vtid = 2, ts = 1704110400000000000, backend = 10 }, { name = "aten::tanh", dur = 2000000000, err = false }

backends/pytorch/tests/multithreaded_ops-AABB.thapi_text_pretty

@@ -0,0 +1,4 @@
+12:00:00.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_entry: { name: "aten::relu", overload_name: "" }
+12:00:00.000000000 - testhost - vpid: 1, vtid: 2 - lttng_ust_pytorch:op_entry: { name: "aten::tanh", overload_name: "" }
+12:00:01.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_exit: { name: "aten::relu", overload_name: "" }
+12:00:02.000000000 - testhost - vpid: 1, vtid: 2 - lttng_ust_pytorch:op_exit: { name: "aten::tanh", overload_name: "" }

backends/pytorch/tests/nested_ops_simple-ABBA.bt_text_pretty

@@ -0,0 +1,2 @@
+lttng:host: { hostname = "testhost", vpid = 1, vtid = 1, ts = 1704110401000000000, backend = 10 }, { name = "aten::matmul", dur = 1000000000, err = false }
+lttng:host: { hostname = "testhost", vpid = 1, vtid = 1, ts = 1704110400000000000, backend = 10 }, { name = "aten::linear", dur = 3000000000, err = false }

backends/pytorch/tests/nested_ops_simple-ABBA.thapi_text_pretty

@@ -0,0 +1,4 @@
+12:00:00.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_entry: { name: "aten::linear", overload_name: "" }
+12:00:01.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_entry: { name: "aten::matmul", overload_name: "" }
+12:00:02.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_exit: { name: "aten::matmul", overload_name: "" }
+12:00:03.000000000 - testhost - vpid: 1, vtid: 1 - lttng_ust_pytorch:op_exit: { name: "aten::linear", overload_name: "" }

backends/pytorch/tests/reentrant_op-AAAA.bt_text_pretty

@@ -0,0 +1,2 @@
+lttng:host: { hostname = "testnode", vpid = 0, vtid = 1, ts = 1704110401000000000, backend = 10 }, { name = "aten::norm.Scalar", dur = 1000000000, err = false }
+lttng:host: { hostname = "testnode", vpid = 0, vtid = 1, ts = 1704110400000000000, backend = 10 }, { name = "aten::norm.Scalar", dur = 3000000000, err = false }

backends/pytorch/tests/reentrant_op-AAAA.thapi_text_pretty

@@ -0,0 +1,4 @@
+12:00:00.000000000 - testnode - vpid: 0, vtid: 1 - lttng_ust_pytorch:op_entry: { name: "aten::norm", overload_name: "Scalar" }
+12:00:01.000000000 - testnode - vpid: 0, vtid: 1 - lttng_ust_pytorch:op_entry: { name: "aten::norm", overload_name: "Scalar" }
+12:00:02.000000000 - testnode - vpid: 0, vtid: 1 - lttng_ust_pytorch:op_exit: { name: "aten::norm", overload_name: "Scalar" }
+12:00:03.000000000 - testnode - vpid: 0, vtid: 1 - lttng_ust_pytorch:op_exit: { name: "aten::norm", overload_name: "Scalar" }

integration_tests/backend_pytorch.bats

@@ -1,7 +1,16 @@
 bats_require_minimum_version 1.5.0
 
-@test "PyTorch: trace contains aten::empty.memory_format" {
+@test "PyTorch: tally 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
 }
+
+@test "PyTorch: trace (-t) contains pretty-printed op_entry/op_exit events" {
+  iprof --backends pytorch -t --analysis-output ./pytorch_trace_out.txt -- \
+    python3 ./integration_tests/pytorch_example.py
+  grep "lttng_ust_pytorch:op_entry" ./pytorch_trace_out.txt
+  grep "lttng_ust_pytorch:op_exit" ./pytorch_trace_out.txt
+  grep 'name: "aten::empty' ./pytorch_trace_out.txt
+  grep 'overload_name: "memory_format' ./pytorch_trace_out.txt
+}

utils/babeltrace_thapi.in

@@ -169,6 +169,7 @@ def get_components(names)
       require 'babeltrace_cuda_lib' if $options[:'backend-names'].include?('cuda')
       require 'babeltrace_hip_lib' if $options[:'backend-names'].include?('hip')
       require 'babeltrace_itt_lib' if $options[:'backend-names'].include?('itt')
+      require 'babeltrace_pytorch_lib' if $options[:'backend-names'].include?('pytorch')
       f = lambda { |iterator, _|
         iterator.next_messages.each do |m|
           next unless m.type == :BT_MESSAGE_TYPE_EVENT