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..fbb7609 --- /dev/null +++ b/include/gpufl/backends/nvidia/module_load_ptx_handler.cpp @@ -0,0 +1,76 @@ +#include "gpufl/backends/nvidia/module_load_ptx_handler.hpp" + +#include + +#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> +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(cbdata); + 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. + 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_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..ff01c79 --- /dev/null +++ b/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> + 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..360f214 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,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 (" - // 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 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& 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 + // 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,7 +406,92 @@ 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, "//" + // 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(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; + } + // 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, @@ -394,8 +499,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..3395746 100644 --- a/include/gpufl/core/dictionary_manager.hpp +++ b/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& skip); + + // 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 +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 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..83f5ea3 --- /dev/null +++ b/tests/core/test_loaded_ptx.cpp @@ -0,0 +1,112 @@ +#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, 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}; + 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); + + // 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