From e0e44cb1da21d4c37fd76f35c5edcf47c0d84066 Mon Sep 17 00:00:00 2001 From: Myoungho Shin Date: Tue, 20 Oct 2026 10:12:26 -0700 Subject: [PATCH 1/3] Capture PTX that programs load at run time PTX was read only from the executable's fatbin with cuobjdump. cuda-oxide (Rust) kernels hand their PTX to cuModuleLoadData at run time, so their sessions had SASS but no PTX. A driver-API handler now copies PTX text passed to cuModuleLoadData, cuModuleLoadDataEx, cuLibraryLoadData and cuLinkAddData (binaries are ignored), and the PTX flush emits its kernels through the same per-kernel parser. Kernels cuobjdump already emitted are not repeated, and cuobjdump's own output is unchanged. --- CMakeLists.txt | 1 + .../gpufl/backends/nvidia/cupti_backend.cpp | 2 + .../nvidia/module_load_ptx_handler.cpp | 77 ++++++++++++ .../nvidia/module_load_ptx_handler.hpp | 27 +++++ include/gpufl/core/dictionary_manager.cpp | 114 ++++++++++++++++-- include/gpufl/core/dictionary_manager.hpp | 43 +++++++ include/gpufl/core/monitor.cpp | 1 + include/gpufl/core/monitor.hpp | 7 ++ include/gpufl/core/monitor_batch_manager.cpp | 4 + include/gpufl/core/monitor_batch_manager.hpp | 1 + tests/CMakeLists.txt | 1 + tests/core/test_loaded_ptx.cpp | 86 +++++++++++++ 12 files changed, 354 insertions(+), 10 deletions(-) create mode 100644 include/gpufl/backends/nvidia/module_load_ptx_handler.cpp create mode 100644 include/gpufl/backends/nvidia/module_load_ptx_handler.hpp create mode 100644 tests/core/test_loaded_ptx.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 42dc36d..4ff34ef 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -402,6 +402,7 @@ if(GPUFL_ENABLE_NVIDIA) include/gpufl/backends/nvidia/cupti_engine_selection.cpp include/gpufl/backends/nvidia/cupti_runtime_support.cpp include/gpufl/backends/nvidia/cuda_cleanup_handler.cpp + include/gpufl/backends/nvidia/module_load_ptx_handler.cpp include/gpufl/backends/nvidia/resource_handler.cpp include/gpufl/backends/nvidia/kernel_launch_handler.cpp include/gpufl/backends/nvidia/graph_structure_handler.cpp diff --git a/include/gpufl/backends/nvidia/cupti_backend.cpp b/include/gpufl/backends/nvidia/cupti_backend.cpp index 188f69c..e93a4b7 100644 --- a/include/gpufl/backends/nvidia/cupti_backend.cpp +++ b/include/gpufl/backends/nvidia/cupti_backend.cpp @@ -1,5 +1,6 @@ #include "gpufl/backends/nvidia/cupti_backend.hpp" #include "gpufl/backends/nvidia/cuda_cleanup_handler.hpp" +#include "gpufl/backends/nvidia/module_load_ptx_handler.hpp" #include @@ -130,6 +131,7 @@ void CuptiBackend::initialize(const MonitorOptions& opts) { // Internal handler registration RegisterHandler(std::make_shared(this)); RegisterHandler(std::make_shared(this)); + RegisterHandler(std::make_shared(this)); RegisterHandler(std::make_shared(this)); RegisterHandler(std::make_shared(this)); RegisterHandler(std::make_shared(this)); diff --git a/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp b/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp new file mode 100644 index 0000000..a99a397 --- /dev/null +++ b/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp @@ -0,0 +1,77 @@ +#include "gpufl/backends/nvidia/module_load_ptx_handler.hpp" + +#include + +#include "gpufl/core/monitor.hpp" + +namespace gpufl { +namespace { + +// cuModuleLoadData and cuLibraryLoadData take a NUL-terminated PTX string or a +// binary with no size; DictionaryManager stops at the NUL, at the first +// non-PTX byte, or here. +constexpr size_t kUnsizedImageLimit = 64u << 20; + +} // namespace + +std::vector> +ModuleLoadPtxHandler::requiredCallbacks() const { + return { + {CUPTI_CB_DOMAIN_DRIVER_API, CUPTI_DRIVER_TRACE_CBID_cuModuleLoadData}, + {CUPTI_CB_DOMAIN_DRIVER_API, CUPTI_DRIVER_TRACE_CBID_cuModuleLoadDataEx}, + {CUPTI_CB_DOMAIN_DRIVER_API, CUPTI_DRIVER_TRACE_CBID_cuLibraryLoadData}, + {CUPTI_CB_DOMAIN_DRIVER_API, CUPTI_DRIVER_TRACE_CBID_cuLinkAddData_v2}, + }; +} + +bool ModuleLoadPtxHandler::shouldHandle(CUpti_CallbackDomain domain, + CUpti_CallbackId cbid) const { + for (const auto& cb : requiredCallbacks()) { + if (cb.first == domain && cb.second == cbid) return true; + } + return false; +} + +void ModuleLoadPtxHandler::handle(CUpti_CallbackDomain domain, + CUpti_CallbackId cbid, const void* cbdata) { + if (!backend_ || domain != CUPTI_CB_DOMAIN_DRIVER_API) return; + // PTX is shown beside SASS, so it is kept only where cubins are. + if (!backend_->NeedsCubinCapture()) return; + + auto* cbInfo = static_cast(cbdata); + if (!cbInfo || cbInfo->callbackSite != CUPTI_API_ENTER || + !cbInfo->functionParams) + return; + + // Only the bytes are copied here; parsing waits for the PTX flush. + switch (cbid) { + case CUPTI_DRIVER_TRACE_CBID_cuModuleLoadData: { + const auto* p = static_cast( + cbInfo->functionParams); + Monitor::EnqueueLoadedPtx(p->image, kUnsizedImageLimit); + break; + } + case CUPTI_DRIVER_TRACE_CBID_cuModuleLoadDataEx: { + const auto* p = static_cast( + cbInfo->functionParams); + Monitor::EnqueueLoadedPtx(p->image, kUnsizedImageLimit); + break; + } + case CUPTI_DRIVER_TRACE_CBID_cuLibraryLoadData: { + const auto* p = static_cast( + cbInfo->functionParams); + Monitor::EnqueueLoadedPtx(p->code, kUnsizedImageLimit); + break; + } + case CUPTI_DRIVER_TRACE_CBID_cuLinkAddData_v2: { + const auto* p = static_cast( + cbInfo->functionParams); + if (p->type == CU_JIT_INPUT_PTX) Monitor::EnqueueLoadedPtx(p->data, p->size); + break; + } + default: + break; + } +} + +} // namespace gpufl diff --git a/include/gpufl/backends/nvidia/module_load_ptx_handler.hpp b/include/gpufl/backends/nvidia/module_load_ptx_handler.hpp new file mode 100644 index 0000000..dbc48ac --- /dev/null +++ b/include/gpufl/backends/nvidia/module_load_ptx_handler.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include "gpufl/backends/nvidia/cupti_backend.hpp" +#include "gpufl/backends/nvidia/cupti_common.hpp" + +namespace gpufl { + +// Keeps the PTX of modules loaded at run time. cuobjdump finds PTX only in an +// executable's fatbin; code that hands PTX text to the driver itself (cuda-oxide +// Rust kernels, NVRTC, a JIT linker) is seen only here, at the load call. +class ModuleLoadPtxHandler : public ICuptiHandler { + public: + explicit ModuleLoadPtxHandler(CuptiBackend* backend) : backend_(backend) {} + + const char* getName() const override { return "ModuleLoadPtxHandler"; } + bool shouldHandle(CUpti_CallbackDomain domain, + CUpti_CallbackId cbid) const override; + void handle(CUpti_CallbackDomain domain, CUpti_CallbackId cbid, + const void* cbdata) override; + std::vector> + requiredCallbacks() const override; + + private: + CuptiBackend* backend_ = nullptr; +}; + +} // namespace gpufl diff --git a/include/gpufl/core/dictionary_manager.cpp b/include/gpufl/core/dictionary_manager.cpp index a46cc30..198c8a7 100644 --- a/include/gpufl/core/dictionary_manager.cpp +++ b/include/gpufl/core/dictionary_manager.cpp @@ -245,9 +245,16 @@ void DictionaryManager::enqueueDisassembly(uint64_t crc, const uint8_t* data, } void DictionaryManager::flushPtx(Logger& logger, const std::string& session_id) { - if (ptx_emitted_) return; - ptx_emitted_ = true; + std::lock_guard flush_lk(ptx_flush_mu_); + if (!ptx_emitted_) { + ptx_emitted_ = true; + dumpExecutablePtx(logger, session_id); + } + drainLoadedPtx(logger, session_id); +} +void DictionaryManager::dumpExecutablePtx(Logger& logger, + const std::string& session_id) { // Resolve the running executable. PTX (the virtual ISA) is embedded in the // binary's fatbin, NOT in the SASS cubins CUPTI delivers, so cuobjdump reads // the executable directly. Covers the target's own kernels; precompiled @@ -310,17 +317,26 @@ void DictionaryManager::flushPtx(Logger& logger, const std::string& session_id) GFL_LOG_DEBUG("[flushPtx] cuobjdump produced ", out.size(), " bytes; head: ", out.substr(0, out.size() < 200 ? out.size() : 200).c_str()); - // cuobjdump -ptx output is one or more PTX modules: a - // .version/.target/.address_size header, then ".visible .entry (" - // kernels with a "{ ... }" body. Split into per-kernel blocks, prepend the - // module header so each is self-contained, and emit one ptx_disassembly per - // kernel keyed by the mangled name (joins to kernels/SASS by md5 of it). + const int emitted = emitPtxText(logger, session_id, out, false); + GFL_LOG_DEBUG("[flushPtx] emitted ", emitted, " ptx_disassembly records"); +} + +int DictionaryManager::emitPtxText(Logger& logger, const std::string& session_id, + const std::string& ptx, bool skipEmitted) { + // The text (cuobjdump -ptx output, or a module loaded at run time) is one + // or more PTX modules: a .version/.target/.address_size header, then + // ".visible .entry (" kernels with a "{ ... }" body. Split into + // per-kernel blocks, prepend the module header so each is self-contained, + // and emit one ptx_disassembly per kernel keyed by the mangled name (joins + // to kernels/SASS by md5 of it). std::string header, arch, curName, curBody; bool inFunc = false, bodyStarted = false; int braceDepth = 0, emitted = 0; auto emitFunc = [&]() { if (curName.empty() || curBody.empty()) return; + const bool firstSeen = emitted_ptx_names_.insert(curName).second; + if (!firstSeen && skipEmitted) return; std::ostringstream oss; oss << "{\"version\":1,\"type\":\"ptx_disassembly\",\"session_id\":\"" << model::jsonEscape(session_id) << "\",\"function_name\":\"" @@ -337,7 +353,7 @@ void DictionaryManager::flushPtx(Logger& logger, const std::string& session_id) } }; - std::istringstream iss(out); + std::istringstream iss(ptx); std::string line; while (std::getline(iss, line)) { if (!line.empty() && line.back() == '\r') line.pop_back(); @@ -386,7 +402,80 @@ void DictionaryManager::flushPtx(Logger& logger, const std::string& session_id) } } - GFL_LOG_DEBUG("[flushPtx] emitted ", emitted, " ptx_disassembly records"); + return emitted; +} + +size_t DictionaryManager::loadedPtxLength(const void* image, size_t limit) { + if (!image || limit == 0) return 0; + const auto* p = static_cast(image); + // Walk the preamble one byte at a time, so a binary image is rejected at + // its first non-text byte and never read past its end: whitespace and + // "//" comment lines may precede the first directive, which must be + // .version. + size_t i = 0; + while (i < limit) { + const char c = p[i]; + if (c == ' ' || c == '\t' || c == '\r' || c == '\n') { ++i; continue; } + if (c == '/' && i + 1 < limit && p[i + 1] == '/') { + while (i < limit && p[i] != '\n' && p[i] != '\0') ++i; + continue; + } + break; + } + static constexpr char kVersion[] = ".version"; + const size_t versionLen = sizeof(kVersion) - 1; + for (size_t k = 0; k < versionLen; ++k) { + if (i + k >= limit || p[i + k] != kVersion[k]) return 0; + } + size_t len = i + versionLen; + while (len < limit && p[len] != '\0') ++len; + return len; +} + +bool DictionaryManager::enqueueLoadedPtx(const void* image, size_t limit) { + // A module bigger than this, or past this total, is skipped rather than + // held in memory for the rest of the run. + static constexpr size_t kMaxModuleBytes = 64u << 20; + static constexpr size_t kMaxPendingBytes = 256u << 20; + const size_t len = loadedPtxLength( + image, limit < kMaxModuleBytes ? limit : kMaxModuleBytes); + if (len == 0) return false; + std::string ptx(static_cast(image), len); + const size_t hash = std::hash{}(ptx); + std::lock_guard lk(mu_); + if (seen_loaded_ptx_.count(hash)) return true; + if (pending_loaded_ptx_bytes_ + len > kMaxPendingBytes) { + GFL_LOG_DEBUG("[enqueueLoadedPtx] pending PTX over ", kMaxPendingBytes, + " bytes - skipping a ", len, "-byte module"); + return false; + } + seen_loaded_ptx_.insert(hash); + pending_loaded_ptx_bytes_ += len; + pending_loaded_ptx_.push_back(std::move(ptx)); + return true; +} + +void DictionaryManager::flushLoadedPtx(Logger& logger, + const std::string& session_id) { + std::lock_guard flush_lk(ptx_flush_mu_); + drainLoadedPtx(logger, session_id); +} + +void DictionaryManager::drainLoadedPtx(Logger& logger, + const std::string& session_id) { + std::vector modules; + { + std::lock_guard lk(mu_); + if (pending_loaded_ptx_.empty()) return; + modules = std::move(pending_loaded_ptx_); + pending_loaded_ptx_.clear(); + pending_loaded_ptx_bytes_ = 0; + } + int emitted = 0; + for (const auto& ptx : modules) emitted += emitPtxText(logger, session_id, ptx, true); + GFL_LOG_DEBUG("[flushLoadedPtx] emitted ", emitted, + " ptx_disassembly records from ", modules.size(), + " loaded modules"); } void DictionaryManager::flushDisassembly(Logger& logger, @@ -394,8 +483,13 @@ void DictionaryManager::flushDisassembly(Logger& logger, std::unordered_map> pending; { std::lock_guard lk(mu_); - if (pending_disasm_cubins_.empty()) return; pending = std::move(pending_disasm_cubins_); + pending_disasm_cubins_.clear(); + } + if (pending.empty()) { + // A module loaded after the last flush may still carry PTX. + flushLoadedPtx(logger, session_id); + return; } // Cubins are present => CUDA kernels exist; dump the executable's PTX once diff --git a/include/gpufl/core/dictionary_manager.hpp b/include/gpufl/core/dictionary_manager.hpp index 6baebc3..94a942c 100644 --- a/include/gpufl/core/dictionary_manager.hpp +++ b/include/gpufl/core/dictionary_manager.hpp @@ -133,6 +133,23 @@ class DictionaryManager { // capture, so this reads the executable directly. Runs once per session. void flushPtx(Logger& logger, const std::string& session_id); + // Keeps a copy of a PTX module handed to the driver at run time + // (cuModuleLoadData and friends): JIT-compiled code such as cuda-oxide + // (Rust) kernels has no fatbin PTX for cuobjdump to find. `limit` bounds + // the read; binaries (ELF cubins, fatbins) and other text are ignored. + // Returns true when the image was PTX and is now queued. + bool enqueueLoadedPtx(const void* image, size_t limit); + + // Emits one ptx_disassembly line per kernel of the PTX modules queued by + // enqueueLoadedPtx(). A kernel already emitted (by cuobjdump or an + // earlier module) is not emitted again. No-op if nothing is queued. + void flushLoadedPtx(Logger& logger, const std::string& session_id); + + // Length of the PTX text at `image` (up to its NUL or `limit`), or 0 when + // the bytes are not PTX: a module that starts with `.version` after + // optional whitespace and `//` comment lines. + static size_t loadedPtxLength(const void* image, size_t limit); + void reset() { std::lock_guard lk(mu_); kernel_dict_.clear(); @@ -156,9 +173,27 @@ class DictionaryManager { pending_source_content_.clear(); source_capture_.reset(); pending_disasm_cubins_.clear(); + pending_loaded_ptx_.clear(); + pending_loaded_ptx_bytes_ = 0; } private: + // Runs cuobjdump -ptx on the running executable and emits its kernels. + void dumpExecutablePtx(Logger& logger, const std::string& session_id); + + // Emits and clears the queued run-time PTX modules. ptx_flush_mu_ held. + void drainLoadedPtx(Logger& logger, const std::string& session_id); + + // Splits PTX text into kernels and emits one ptx_disassembly line per + // kernel, recording each name in emitted_ptx_names_. With skipEmitted, a + // kernel already recorded is not emitted again (cuobjdump's output keeps + // every copy, as before). Returns the number emitted. + int emitPtxText(Logger& logger, const std::string& session_id, + const std::string& ptx, bool skipEmitted); + + // Serializes the PTX flushes, which run from more than one flush path. + std::mutex ptx_flush_mu_; + std::mutex mu_; std::unordered_map kernel_dict_; @@ -197,6 +232,14 @@ class DictionaryManager { // PTX is dumped from the executable once per session (see flushPtx). bool ptx_emitted_ = false; + // PTX modules loaded at run time (see enqueueLoadedPtx), their total size, + // the hashes of modules already queued, and the kernels already emitted + // by either path. + std::vector pending_loaded_ptx_; + size_t pending_loaded_ptx_bytes_ = 0; + std::unordered_set seen_loaded_ptx_; + std::unordered_set emitted_ptx_names_; + // Batch ID counter for profile_sample_batch records emitted from // AMD DWARF source mapping (avoids collision with monitor batch IDs // by starting at a high offset). diff --git a/include/gpufl/core/monitor.cpp b/include/gpufl/core/monitor.cpp index 2c754b9..676b516 100644 --- a/include/gpufl/core/monitor.cpp +++ b/include/gpufl/core/monitor.cpp @@ -822,6 +822,7 @@ void Monitor::EndPerfScope(const char* name) { if (auto* b = GetBackend()) b->On IMonitorBackend* Monitor::GetBackend() { return g_state.adapter ? g_state.adapter->backend() : nullptr; } uint32_t Monitor::InternScopeName(const std::string& name) { return g_state.batches.internScopeName(name); } void Monitor::EnqueueCubinForDisassembly(uint64_t crc, const uint8_t* data, size_t size) { g_state.batches.enqueueDisassembly(crc, data, size); } +void Monitor::EnqueueLoadedPtx(const void* image, size_t limit) { g_state.batches.enqueueLoadedPtx(image, limit); } void Monitor::FlushDisassemblyNow() { g_state.batches.flushDisassembly(); } diff --git a/include/gpufl/core/monitor.hpp b/include/gpufl/core/monitor.hpp index 5c0e3d7..117bdd0 100644 --- a/include/gpufl/core/monitor.hpp +++ b/include/gpufl/core/monitor.hpp @@ -409,6 +409,13 @@ class Monitor { static void EnqueueCubinForDisassembly(uint64_t crc, const uint8_t* data, size_t size); + /** + * @brief Keep a PTX module handed to the driver at run time for the + * next PTX flush. Called from the CUPTI module-load callbacks; binaries + * and non-PTX images are ignored. Reads at most `limit` bytes. + */ + static void EnqueueLoadedPtx(const void* image, size_t limit); + /** * Disassemble and emit any enqueued cubins now, instead of waiting for the * shutdown flush. Used by the cubin worker under Windows-injection PC diff --git a/include/gpufl/core/monitor_batch_manager.cpp b/include/gpufl/core/monitor_batch_manager.cpp index 54688b6..34c1fdd 100644 --- a/include/gpufl/core/monitor_batch_manager.cpp +++ b/include/gpufl/core/monitor_batch_manager.cpp @@ -212,6 +212,10 @@ void MonitorBatchManager::enqueueDisassembly(uint64_t crc, const uint8_t* data, dictManager_.enqueueDisassembly(crc, data, size); } +bool MonitorBatchManager::enqueueLoadedPtx(const void* image, size_t limit) { + return dictManager_.enqueueLoadedPtx(image, limit); +} + void MonitorBatchManager::flushDisassembly() { if (!flushSink_.available()) return; const auto context = diff --git a/include/gpufl/core/monitor_batch_manager.hpp b/include/gpufl/core/monitor_batch_manager.hpp index fe950d0..7d76e03 100644 --- a/include/gpufl/core/monitor_batch_manager.hpp +++ b/include/gpufl/core/monitor_batch_manager.hpp @@ -46,6 +46,7 @@ class MonitorBatchManager { std::string sourceFileName(uint32_t source_file_id); void enqueueDisassembly(uint64_t crc, const uint8_t* data, size_t size); + bool enqueueLoadedPtx(const void* image, size_t limit); void flushDisassembly(); uint64_t allocateScopeInstanceId(); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 1db43ae..ea72765 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -35,6 +35,7 @@ set(GPUFL_TEST_SOURCES core/test_sampler.cpp core/test_segment_context.cpp core/test_segment_coordinator.cpp + core/test_loaded_ptx.cpp core/test_session_bootstrap.cpp core/test_source_capture_policy.cpp core/test_startup_configuration.cpp diff --git a/tests/core/test_loaded_ptx.cpp b/tests/core/test_loaded_ptx.cpp new file mode 100644 index 0000000..7b92ec8 --- /dev/null +++ b/tests/core/test_loaded_ptx.cpp @@ -0,0 +1,86 @@ +#include + +#include +#include +#include + +#include "gpufl/core/dictionary_manager.hpp" +#include "gpufl/core/logger/log_sink.hpp" +#include "gpufl/core/logger/logger.hpp" + +namespace { + +class RecordingSink final : public gpufl::ILogSink { + public: + explicit RecordingSink(std::shared_ptr> lines) + : lines_(std::move(lines)) {} + + void write(gpufl::Channel, std::string_view json) override { + lines_->emplace_back(json); + } + void close() override {} + + private: + std::shared_ptr> lines_; +}; + +// The shape cuda-oxide hands to cuModuleLoadData: LLVM's comment banner, one +// module header, unmangled Rust kernel names. +constexpr char kRustPtx[] = + "//\n" + "// Generated by LLVM NVPTX Back-End\n" + "//\n" + "\n" + ".version 8.7\n" + ".target sm_89\n" + ".address_size 64\n" + "\n" + ".visible .entry layernorm_rows(\n" + "\t.param .u64 layernorm_rows_param_0\n" + ")\n" + "{\n" + "\tret;\n" + "}\n" + ".visible .entry layernorm_rows_baseline(\n" + "\t.param .u64 layernorm_rows_baseline_param_0\n" + ")\n" + "{\n" + "\tret;\n" + "}\n"; + +TEST(LoadedPtxTest, EmitsEachKernelOfALoadedModuleOnce) { + gpufl::DictionaryManager registry; + gpufl::Logger logger; + auto lines = std::make_shared>(); + logger.addSink(std::make_unique(lines)); + + ASSERT_TRUE(registry.enqueueLoadedPtx(kRustPtx, 1u << 20)); + registry.flushLoadedPtx(logger, "s0"); + + ASSERT_EQ(lines->size(), 2u); + EXPECT_NE((*lines)[0].find("\"type\":\"ptx_disassembly\""), std::string::npos); + EXPECT_NE((*lines)[0].find("\"function_name\":\"layernorm_rows\""), std::string::npos); + EXPECT_NE((*lines)[0].find("\"arch\":\"sm_89\""), std::string::npos); + // Each kernel carries the module header, so it stands alone. + EXPECT_NE((*lines)[0].find(".version 8.7"), std::string::npos); + EXPECT_NE((*lines)[1].find("\"function_name\":\"layernorm_rows_baseline\""), + std::string::npos); + + // The same module loaded again (another context) adds nothing. + ASSERT_TRUE(registry.enqueueLoadedPtx(kRustPtx, 1u << 20)); + registry.flushLoadedPtx(logger, "s0"); + EXPECT_EQ(lines->size(), 2u); +} + +TEST(LoadedPtxTest, IgnoresBinaryImages) { + gpufl::DictionaryManager registry; + // An ELF cubin and a fatbin header: neither is PTX text. + const unsigned char cubin[] = {0x7f, 'E', 'L', 'F', 2, 1, 1, 0x33}; + const unsigned char fatbin[] = {0x50, 0xed, 0x55, 0xba, 1, 0, 0x10, 0}; + EXPECT_FALSE(registry.enqueueLoadedPtx(cubin, sizeof(cubin))); + EXPECT_FALSE(registry.enqueueLoadedPtx(fatbin, sizeof(fatbin))); + // A PTX image cut off by its size limit before .version is complete. + EXPECT_EQ(gpufl::DictionaryManager::loadedPtxLength(kRustPtx, 40), 0u); +} + +} // namespace From 32216cdd665f496481099a54fe5594726d0dec84 Mon Sep 17 00:00:00 2001 From: Myoungho Shin Date: Tue, 20 Oct 2026 10:37:57 -0700 Subject: [PATCH 2/3] Address the loaded PTX review - Keep a module's PTX only when its load returned CUDA_SUCCESS (API exit), so a rejected PTX cannot shadow a fallback with the same kernel names. - reset() clears the queued hashes, the emitted kernel names and the executable-dump flag, so a later session in the same process gets PTX for the modules it loads and runs cuobjdump again. - The one-time cuobjdump dump skips kernels the run-time path already emitted in this session; its own per-arch copies are kept. - PTX may open with a /* */ comment; one never closed is not PTX. - Note that loads by file name (cuModuleLoad, cuLinkAddFile, cuLibraryLoadFromFile) are not captured. --- .../nvidia/module_load_ptx_handler.cpp | 10 ++++-- .../nvidia/module_load_ptx_handler.hpp | 3 ++ include/gpufl/core/dictionary_manager.cpp | 34 ++++++++++++++----- include/gpufl/core/dictionary_manager.hpp | 22 ++++++++---- tests/core/test_loaded_ptx.cpp | 28 ++++++++++++++- 5 files changed, 78 insertions(+), 19 deletions(-) diff --git a/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp b/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp index a99a397..5d01016 100644 --- a/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp +++ b/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp @@ -38,9 +38,15 @@ void ModuleLoadPtxHandler::handle(CUpti_CallbackDomain domain, // PTX is shown beside SASS, so it is kept only where cubins are. if (!backend_->NeedsCubinCapture()) return; + // Keep only a module the driver accepted: a program may retry a rejected + // PTX (an unsupported .version, say) with a fallback that has the same + // kernel names, and the first PTX kept for a name is the one emitted. The + // caller's image is still valid at API exit. auto* cbInfo = static_cast(cbdata); - if (!cbInfo || cbInfo->callbackSite != CUPTI_API_ENTER || - !cbInfo->functionParams) + if (!cbInfo || cbInfo->callbackSite != CUPTI_API_EXIT || + !cbInfo->functionParams || !cbInfo->functionReturnValue) + return; + if (*static_cast(cbInfo->functionReturnValue) != CUDA_SUCCESS) return; // Only the bytes are copied here; parsing waits for the PTX flush. diff --git a/include/gpufl/backends/nvidia/module_load_ptx_handler.hpp b/include/gpufl/backends/nvidia/module_load_ptx_handler.hpp index dbc48ac..97bbb38 100644 --- a/include/gpufl/backends/nvidia/module_load_ptx_handler.hpp +++ b/include/gpufl/backends/nvidia/module_load_ptx_handler.hpp @@ -8,6 +8,9 @@ namespace gpufl { // Keeps the PTX of modules loaded at run time. cuobjdump finds PTX only in an // executable's fatbin; code that hands PTX text to the driver itself (cuda-oxide // Rust kernels, NVRTC, a JIT linker) is seen only here, at the load call. +// Loads that name a file instead of passing the image (cuModuleLoad, +// cuLinkAddFile, cuLibraryLoadFromFile) are not captured; cuda-oxide passes +// the image. class ModuleLoadPtxHandler : public ICuptiHandler { public: explicit ModuleLoadPtxHandler(CuptiBackend* backend) : backend_(backend) {} diff --git a/include/gpufl/core/dictionary_manager.cpp b/include/gpufl/core/dictionary_manager.cpp index 198c8a7..360f214 100644 --- a/include/gpufl/core/dictionary_manager.cpp +++ b/include/gpufl/core/dictionary_manager.cpp @@ -317,12 +317,17 @@ void DictionaryManager::dumpExecutablePtx(Logger& logger, GFL_LOG_DEBUG("[flushPtx] cuobjdump produced ", out.size(), " bytes; head: ", out.substr(0, out.size() < 200 ? out.size() : 200).c_str()); - const int emitted = emitPtxText(logger, session_id, out, false); + // A flush with no cubins may have emitted run-time PTX before this dump: + // those kernels are not repeated. cuobjdump's own copies of a kernel (one + // per PTX arch) are all kept, as before. + const std::unordered_set emittedBefore = emitted_ptx_names_; + const int emitted = emitPtxText(logger, session_id, out, emittedBefore); GFL_LOG_DEBUG("[flushPtx] emitted ", emitted, " ptx_disassembly records"); } int DictionaryManager::emitPtxText(Logger& logger, const std::string& session_id, - const std::string& ptx, bool skipEmitted) { + const std::string& ptx, + const std::unordered_set& skip) { // The text (cuobjdump -ptx output, or a module loaded at run time) is one // or more PTX modules: a .version/.target/.address_size header, then // ".visible .entry (" kernels with a "{ ... }" body. Split into @@ -334,9 +339,8 @@ int DictionaryManager::emitPtxText(Logger& logger, const std::string& session_id int braceDepth = 0, emitted = 0; auto emitFunc = [&]() { - if (curName.empty() || curBody.empty()) return; - const bool firstSeen = emitted_ptx_names_.insert(curName).second; - if (!firstSeen && skipEmitted) return; + if (curName.empty() || curBody.empty() || skip.count(curName)) return; + emitted_ptx_names_.insert(curName); std::ostringstream oss; oss << "{\"version\":1,\"type\":\"ptx_disassembly\",\"session_id\":\"" << model::jsonEscape(session_id) << "\",\"function_name\":\"" @@ -409,9 +413,9 @@ size_t DictionaryManager::loadedPtxLength(const void* image, size_t limit) { if (!image || limit == 0) return 0; const auto* p = static_cast(image); // Walk the preamble one byte at a time, so a binary image is rejected at - // its first non-text byte and never read past its end: whitespace and - // "//" comment lines may precede the first directive, which must be - // .version. + // its first non-text byte and never read past its end: whitespace, "//" + // comment lines and "/* */" comments may precede the first directive, + // which must be .version. size_t i = 0; while (i < limit) { const char c = p[i]; @@ -420,6 +424,15 @@ size_t DictionaryManager::loadedPtxLength(const void* image, size_t limit) { while (i < limit && p[i] != '\n' && p[i] != '\0') ++i; continue; } + if (c == '/' && i + 1 < limit && p[i + 1] == '*') { + i += 2; + while (i + 1 < limit && p[i] != '\0' && + !(p[i] == '*' && p[i + 1] == '/')) + ++i; + if (i + 1 >= limit || p[i] == '\0') return 0; // never closed + i += 2; + continue; + } break; } static constexpr char kVersion[] = ".version"; @@ -471,8 +484,11 @@ void DictionaryManager::drainLoadedPtx(Logger& logger, pending_loaded_ptx_.clear(); pending_loaded_ptx_bytes_ = 0; } + // Checked against the live set, so a kernel is emitted once across every + // module and the executable dump. int emitted = 0; - for (const auto& ptx : modules) emitted += emitPtxText(logger, session_id, ptx, true); + for (const auto& ptx : modules) + emitted += emitPtxText(logger, session_id, ptx, emitted_ptx_names_); GFL_LOG_DEBUG("[flushLoadedPtx] emitted ", emitted, " ptx_disassembly records from ", modules.size(), " loaded modules"); diff --git a/include/gpufl/core/dictionary_manager.hpp b/include/gpufl/core/dictionary_manager.hpp index 94a942c..3395746 100644 --- a/include/gpufl/core/dictionary_manager.hpp +++ b/include/gpufl/core/dictionary_manager.hpp @@ -130,7 +130,9 @@ class DictionaryManager { // Runs cuobjdump -ptx on the running executable and emits one // ptx_disassembly JSON line per kernel (joined to kernels/SASS by the same // mangled function name). PTX lives in the fatbin, not the SASS cubins we - // capture, so this reads the executable directly. Runs once per session. + // capture, so this reads the executable directly; that dump runs once per + // session. Every call then emits the PTX modules loaded at run time (see + // enqueueLoadedPtx). void flushPtx(Logger& logger, const std::string& session_id); // Keeps a copy of a PTX module handed to the driver at run time @@ -147,11 +149,14 @@ class DictionaryManager { // Length of the PTX text at `image` (up to its NUL or `limit`), or 0 when // the bytes are not PTX: a module that starts with `.version` after - // optional whitespace and `//` comment lines. + // optional whitespace, `//` comment lines and `/* */` comments. static size_t loadedPtxLength(const void* image, size_t limit); void reset() { - std::lock_guard lk(mu_); + // The PTX state belongs to the session too: a new session dumps the + // executable again and keeps the PTX of every module it loads, seen + // in an earlier session or not. Same lock order as the PTX flushes. + std::scoped_lock lk(ptx_flush_mu_, mu_); kernel_dict_.clear(); dirty_kernels_.clear(); next_kernel_id_ = 1; @@ -175,6 +180,9 @@ class DictionaryManager { pending_disasm_cubins_.clear(); pending_loaded_ptx_.clear(); pending_loaded_ptx_bytes_ = 0; + seen_loaded_ptx_.clear(); + emitted_ptx_names_.clear(); + ptx_emitted_ = false; } private: @@ -185,11 +193,11 @@ class DictionaryManager { void drainLoadedPtx(Logger& logger, const std::string& session_id); // Splits PTX text into kernels and emits one ptx_disassembly line per - // kernel, recording each name in emitted_ptx_names_. With skipEmitted, a - // kernel already recorded is not emitted again (cuobjdump's output keeps - // every copy, as before). Returns the number emitted. + // kernel whose name is not in `skip`, recording each emitted name in + // emitted_ptx_names_ (`skip` may be that set). Returns the number emitted. int emitPtxText(Logger& logger, const std::string& session_id, - const std::string& ptx, bool skipEmitted); + const std::string& ptx, + const std::unordered_set& skip); // Serializes the PTX flushes, which run from more than one flush path. std::mutex ptx_flush_mu_; diff --git a/tests/core/test_loaded_ptx.cpp b/tests/core/test_loaded_ptx.cpp index 7b92ec8..83f5ea3 100644 --- a/tests/core/test_loaded_ptx.cpp +++ b/tests/core/test_loaded_ptx.cpp @@ -72,7 +72,26 @@ TEST(LoadedPtxTest, EmitsEachKernelOfALoadedModuleOnce) { EXPECT_EQ(lines->size(), 2u); } -TEST(LoadedPtxTest, IgnoresBinaryImages) { +TEST(LoadedPtxTest, ANewSessionEmitsTheModuleAgain) { + gpufl::DictionaryManager registry; + gpufl::Logger logger; + auto lines = std::make_shared>(); + logger.addSink(std::make_unique(lines)); + + ASSERT_TRUE(registry.enqueueLoadedPtx(kRustPtx, 1u << 20)); + registry.flushLoadedPtx(logger, "s0"); + ASSERT_EQ(lines->size(), 2u); + + // Session start resets the registry; the module loaded again in the new + // session gets its PTX again. + registry.reset(); + ASSERT_TRUE(registry.enqueueLoadedPtx(kRustPtx, 1u << 20)); + registry.flushLoadedPtx(logger, "s1"); + ASSERT_EQ(lines->size(), 4u); + EXPECT_NE((*lines)[2].find("\"session_id\":\"s1\""), std::string::npos); +} + +TEST(LoadedPtxTest, RecognizesOnlyPtxText) { gpufl::DictionaryManager registry; // An ELF cubin and a fatbin header: neither is PTX text. const unsigned char cubin[] = {0x7f, 'E', 'L', 'F', 2, 1, 1, 0x33}; @@ -81,6 +100,13 @@ TEST(LoadedPtxTest, IgnoresBinaryImages) { EXPECT_FALSE(registry.enqueueLoadedPtx(fatbin, sizeof(fatbin))); // A PTX image cut off by its size limit before .version is complete. EXPECT_EQ(gpufl::DictionaryManager::loadedPtxLength(kRustPtx, 40), 0u); + + // A block comment may open the module; one never closed is not PTX. + constexpr char kBlockComment[] = "/* generated */\n.version 8.7\n"; + EXPECT_EQ(gpufl::DictionaryManager::loadedPtxLength(kBlockComment, 1u << 20), + sizeof(kBlockComment) - 1); + EXPECT_EQ(gpufl::DictionaryManager::loadedPtxLength("/* open .version 8.7\n", 1u << 20), + 0u); } } // namespace From 6a0462a841c97c609a04784984e278ba6a41baa0 Mon Sep 17 00:00:00 2001 From: Myoungho Shin Date: Tue, 20 Oct 2026 10:57:00 -0700 Subject: [PATCH 3/3] Stop capturing PTX handed to cuLibraryLoadData cuLibraryLoadData returns success before the driver compiles PTX: on driver 610.62 an unsupported .version loads with CUDA_SUCCESS and fails only at cuLibraryGetKernel / cuLibraryGetModule, with CUDA_MODULE_LOADING lazy, eager or unset. Its success therefore does not prove the PTX is the code that runs, and the API mostly carries the CUDA runtime's own fatbins. cuModuleLoadData, cuModuleLoadDataEx and cuLinkAddData reject the same PTX with CUDA_ERROR_UNSUPPORTED_PTX_VERSION before they return in all three modes, so their success is final and they stay captured. cuda-oxide, NVRTC users and Numba hand PTX to those. --- .../backends/nvidia/module_load_ptx_handler.cpp | 17 +++++------------ .../backends/nvidia/module_load_ptx_handler.hpp | 7 ++++++- 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp b/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp index 5d01016..fbb7609 100644 --- a/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp +++ b/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp @@ -7,9 +7,8 @@ namespace gpufl { namespace { -// cuModuleLoadData and cuLibraryLoadData take a NUL-terminated PTX string or a -// binary with no size; DictionaryManager stops at the NUL, at the first -// non-PTX byte, or here. +// cuModuleLoadData(Ex) take a NUL-terminated PTX string or a binary with no +// size; DictionaryManager stops at the NUL, at the first non-PTX byte, or here. constexpr size_t kUnsizedImageLimit = 64u << 20; } // namespace @@ -19,7 +18,6 @@ ModuleLoadPtxHandler::requiredCallbacks() const { return { {CUPTI_CB_DOMAIN_DRIVER_API, CUPTI_DRIVER_TRACE_CBID_cuModuleLoadData}, {CUPTI_CB_DOMAIN_DRIVER_API, CUPTI_DRIVER_TRACE_CBID_cuModuleLoadDataEx}, - {CUPTI_CB_DOMAIN_DRIVER_API, CUPTI_DRIVER_TRACE_CBID_cuLibraryLoadData}, {CUPTI_CB_DOMAIN_DRIVER_API, CUPTI_DRIVER_TRACE_CBID_cuLinkAddData_v2}, }; } @@ -40,8 +38,9 @@ void ModuleLoadPtxHandler::handle(CUpti_CallbackDomain domain, // Keep only a module the driver accepted: a program may retry a rejected // PTX (an unsupported .version, say) with a fallback that has the same - // kernel names, and the first PTX kept for a name is the one emitted. The - // caller's image is still valid at API exit. + // kernel names, and the first PTX kept for a name is the one emitted. + // These calls compile PTX before they return, lazy loading included, so + // their result is final. The caller's image is still valid at API exit. auto* cbInfo = static_cast(cbdata); if (!cbInfo || cbInfo->callbackSite != CUPTI_API_EXIT || !cbInfo->functionParams || !cbInfo->functionReturnValue) @@ -63,12 +62,6 @@ void ModuleLoadPtxHandler::handle(CUpti_CallbackDomain domain, Monitor::EnqueueLoadedPtx(p->image, kUnsizedImageLimit); break; } - case CUPTI_DRIVER_TRACE_CBID_cuLibraryLoadData: { - const auto* p = static_cast( - cbInfo->functionParams); - Monitor::EnqueueLoadedPtx(p->code, kUnsizedImageLimit); - break; - } case CUPTI_DRIVER_TRACE_CBID_cuLinkAddData_v2: { const auto* p = static_cast( cbInfo->functionParams); diff --git a/include/gpufl/backends/nvidia/module_load_ptx_handler.hpp b/include/gpufl/backends/nvidia/module_load_ptx_handler.hpp index 97bbb38..ff01c79 100644 --- a/include/gpufl/backends/nvidia/module_load_ptx_handler.hpp +++ b/include/gpufl/backends/nvidia/module_load_ptx_handler.hpp @@ -10,7 +10,12 @@ namespace gpufl { // Rust kernels, NVRTC, a JIT linker) is seen only here, at the load call. // Loads that name a file instead of passing the image (cuModuleLoad, // cuLinkAddFile, cuLibraryLoadFromFile) are not captured; cuda-oxide passes -// the image. +// the image. cuLibraryLoadData is not captured either: it returns success +// before the driver compiles PTX (an unsupported .version fails only at +// cuLibraryGetKernel or cuLibraryGetModule, eager loading included), and it +// mostly carries the CUDA runtime's own fatbins. cuda-oxide, NVRTC users and +// Numba hand PTX to cuModuleLoadData(Ex) or cuLinkAddData, which reject PTX +// the driver cannot compile. class ModuleLoadPtxHandler : public ICuptiHandler { public: explicit ModuleLoadPtxHandler(CuptiBackend* backend) : backend_(backend) {}