Capture PTX that programs load at run time

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

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 <cupti_profiler_target.h>
 
@@ -130,6 +131,7 @@ void CuptiBackend::initialize(const MonitorOptions& opts) {
     // Internal handler registration
     RegisterHandler(std::make_shared<ResourceHandler>(this));
     RegisterHandler(std::make_shared<CudaCleanupHandler>(this));
+    RegisterHandler(std::make_shared<ModuleLoadPtxHandler>(this));
     RegisterHandler(std::make_shared<KernelLaunchHandler>(this));
     RegisterHandler(std::make_shared<GraphStructureHandler>(this));
     RegisterHandler(std::make_shared<MemTransferHandler>(this));

include/gpufl/backends/nvidia/module_load_ptx_handler.cpp

@@ -0,0 +1,76 @@
+#include "gpufl/backends/nvidia/module_load_ptx_handler.hpp"
+
+#include <cstddef>
+
+#include "gpufl/core/monitor.hpp"
+
+namespace gpufl {
+namespace {
+
+// 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
+
+std::vector<std::pair<CUpti_CallbackDomain, CUpti_CallbackId>>
+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_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;
+
+    // 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.
+    // 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<const CUpti_CallbackData*>(cbdata);
+    if (!cbInfo || cbInfo->callbackSite != CUPTI_API_EXIT ||
+        !cbInfo->functionParams || !cbInfo->functionReturnValue)
+        return;
+    if (*static_cast<const CUresult*>(cbInfo->functionReturnValue) != CUDA_SUCCESS)
+        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<const cuModuleLoadData_params*>(
+                cbInfo->functionParams);
+            Monitor::EnqueueLoadedPtx(p->image, kUnsizedImageLimit);
+            break;
+        }
+        case CUPTI_DRIVER_TRACE_CBID_cuModuleLoadDataEx: {
+            const auto* p = static_cast<const cuModuleLoadDataEx_params*>(
+                cbInfo->functionParams);
+            Monitor::EnqueueLoadedPtx(p->image, kUnsizedImageLimit);
+            break;
+        }
+        case CUPTI_DRIVER_TRACE_CBID_cuLinkAddData_v2: {
+            const auto* p = static_cast<const cuLinkAddData_v2_params*>(
+                cbInfo->functionParams);
+            if (p->type == CU_JIT_INPUT_PTX) Monitor::EnqueueLoadedPtx(p->data, p->size);
+            break;
+        }
+        default:
+            break;
+    }
+}
+
+}  // namespace gpufl

include/gpufl/backends/nvidia/module_load_ptx_handler.hpp

@@ -0,0 +1,35 @@
+#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.
+// Loads that name a file instead of passing the image (cuModuleLoad,
+// cuLinkAddFile, cuLibraryLoadFromFile) are not captured; cuda-oxide passes
+// 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) {}
+
+    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<std::pair<CUpti_CallbackDomain, CUpti_CallbackId>>
+    requiredCallbacks() const override;
+
+   private:
+    CuptiBackend* backend_ = nullptr;
+};
+
+}  // namespace gpufl

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,30 @@ 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 <mangled>("
-    // 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).
+    // 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<std::string> 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,
+                                   const std::unordered_set<std::string>& 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 <mangled>(" 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;
+        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\":\""
@@ -337,7 +357,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,16 +406,106 @@ 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<const char*>(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, "//"
+    // 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];
+        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;
+        }
+        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";
+    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<const char*>(image), len);
+    const size_t hash = std::hash<std::string>{}(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<std::string> 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;
+    }
+    // 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, emitted_ptx_names_);
+    GFL_LOG_DEBUG("[flushLoadedPtx] emitted ", emitted,
+                  " ptx_disassembly records from ", modules.size(),
+                  " loaded modules");
 }
 
 void DictionaryManager::flushDisassembly(Logger& logger,
                                           const std::string& session_id) {
     std::unordered_map<uint64_t, std::vector<uint8_t>> 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

include/gpufl/core/dictionary_manager.hpp

@@ -130,11 +130,33 @@ 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
+    // (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, `//` 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;
@@ -156,9 +178,30 @@ class DictionaryManager {
         pending_source_content_.clear();
         source_capture_.reset();
         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:
+    // 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 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,
+                    const std::unordered_set<std::string>& skip);
+
+    // Serializes the PTX flushes, which run from more than one flush path.
+    std::mutex ptx_flush_mu_;
+
     std::mutex mu_;
 
     std::unordered_map<std::string, uint32_t> kernel_dict_;
@@ -197,6 +240,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<std::string> pending_loaded_ptx_;
+    size_t pending_loaded_ptx_bytes_ = 0;
+    std::unordered_set<size_t> seen_loaded_ptx_;
+    std::unordered_set<std::string> 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).

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();
 }