diff --git a/Makefile.am b/Makefile.am index f24f94b3b..79a737042 100644 --- a/Makefile.am +++ b/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 diff --git a/backends/pytorch/Makefile.am b/backends/pytorch/Makefile.am new file mode 100644 index 000000000..44ccfceea --- /dev/null +++ b/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) diff --git a/backends/pytorch/include/ATen/record_function.h b/backends/pytorch/include/ATen/record_function.h new file mode 100644 index 000000000..6003887ed --- /dev/null +++ b/backends/pytorch/include/ATen/record_function.h @@ -0,0 +1,96 @@ +// Minimal hand-written stand-in for PyTorch's . +// +// 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 +#include +#include +#include + +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 (*)(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 &scopes) { + if (!scopes.empty()) { + scopes_.fill(false); + for (auto sc : scopes) { + scopes_[static_cast(sc)] = true; + } + } else { + scopes_.fill(true); + } + return *this; + } + +private: + StartCallback start_; + EndCallback end_; + double sampling_prob_ = 1.0; + std::array(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 diff --git a/backends/pytorch/pytorch_events.yaml b/backends/pytorch/pytorch_events.yaml new file mode 100644 index 000000000..1d6807f42 --- /dev/null +++ b/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] diff --git a/backends/pytorch/tracer_pytorch.cpp b/backends/pytorch/tracer_pytorch.cpp new file mode 100644 index 000000000..ac95849ae --- /dev/null +++ b/backends/pytorch/tracer_pytorch.cpp @@ -0,0 +1,21 @@ +#include + +#include "pytorch.h" + +// ENTRY: fires BEFORE the op runs. LTTng adds time + vpid/vtid via context. +static std::unique_ptr 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})); +} diff --git a/backends/pytorch/tracer_pytorch.sh.in b/backends/pytorch/tracer_pytorch.sh.in new file mode 100644 index 000000000..696ddabcc --- /dev/null +++ b/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] [--] " + 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 diff --git a/configure.ac b/configure.ac index 275796323..661244ebd 100644 --- a/configure.ac +++ b/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]) diff --git a/xprof/xprof.rb.in b/xprof/xprof.rb.in index 62d19453a..35cf9b9cd 100755 --- a/xprof/xprof.rb.in +++ b/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')