diff --git a/projects/rocprofiler-sdk/source/docs/conceptual/range_replay/index.md b/projects/rocprofiler-sdk/source/docs/conceptual/range_replay/index.md index c2fdb97a27df..cca8c9b94024 100644 --- a/projects/rocprofiler-sdk/source/docs/conceptual/range_replay/index.md +++ b/projects/rocprofiler-sdk/source/docs/conceptual/range_replay/index.md @@ -91,8 +91,8 @@ experimental/range_replay.h public payload struct, status enum, begin/ +-- hsa/queue.cpp recording hook, foreign-dispatch detection +-- hsa/async_copy.cpp device-write detection | - +-- kernel_replay/ reused unchanged: memory_tracker, memory_snapshot, - local_context + +-- kernel_replay/ shared: memory_tracker, GPU/pinned snapshot planner, + batched blit restore, local_context ``` `range_state.cpp` and `digest.cpp` are deliberately free of GPU dependencies. The eligibility diff --git a/projects/rocprofiler-sdk/source/docs/conceptual/range_replay/range_replay_soundness.md b/projects/rocprofiler-sdk/source/docs/conceptual/range_replay/range_replay_soundness.md index b793181316cb..b20cbae73c26 100644 --- a/projects/rocprofiler-sdk/source/docs/conceptual/range_replay/range_replay_soundness.md +++ b/projects/rocprofiler-sdk/source/docs/conceptual/range_replay/range_replay_soundness.md @@ -150,9 +150,10 @@ restored after the last pass so the application resumes with its own results — final replayed pass left behind. Skipping it would silently hand the application the replay's output. -A failed `restore()` is fatal rather than a decline. A partial host-to-device copy leaves device -memory in a state that is neither the application's nor the snapshot's, and continuing would hand -that corruption to the application. This matches kernel replay's convention for the same situation. +A failed `restore()` is fatal rather than a decline. A partial host-to-device copy or failed device +blit leaves device memory in a state that is neither the application's nor the snapshot's, and +continuing would hand that corruption to the application. This matches kernel replay's convention +for the same situation. Each dispatch in a pass has its AQL barrier bit forced on, so the passes are serialized dispatch-by-dispatch. This is stricter than the application, whose packets may have been free to @@ -167,8 +168,14 @@ virtual-memory-mapped allocations, or a value the host recomputed inside the ran passes every check and is replayed, and its passes quietly describe different inputs. Setting `ROCPROF_RANGE_REPLAY_VERIFY` turns on a check for exactly that residue. The executor hashes -each snapshot region (FNV-1a) after the application's execution and again after the final replayed -pass, and reports how many regions differ as the `CLOSE` callback's `divergence_count`. +each snapshot region after the application's execution and again after the final replayed pass, and +reports how many regions differ as the `CLOSE` callback's `divergence_count`. Hashing uses the +SSE4.2 CRC32 instruction when the host provides it and falls back to FNV-1a otherwise; the digests +are compared only against digests produced by the same binary in the same process and are never +persisted, so the hash function may change between releases. The check costs one extra +device-to-host copy of the snapshot per digest plus the hash itself; with the CRC32 path the copy +dominates, and on hosts without SSE4.2 the fallback hash is slower than the copy and dominates +instead. Regions are ordered by device address before hashing, because the allocation inventory itself is unordered and the two snapshots must compare positionally. A region-count mismatch is reported as diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue.cpp index e2ca7033a296..fa263350115f 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/queue.cpp @@ -40,6 +40,7 @@ #include "lib/rocprofiler-sdk/hsa/signal_pool.hpp" #include "lib/rocprofiler-sdk/kernel_dispatch/profiling_time.hpp" #include "lib/rocprofiler-sdk/kernel_dispatch/tracing.hpp" +#include "lib/rocprofiler-sdk/kernel_replay/blit-copy.hpp" #include "lib/rocprofiler-sdk/kernel_replay/local_context.hpp" #include "lib/rocprofiler-sdk/kernel_replay/memory_snapshot.hpp" #include "lib/rocprofiler-sdk/kernel_replay/replay_callbacks.hpp" @@ -1003,17 +1004,40 @@ WriteInterceptor(const void* packets, // amd-vkale). replay_drain_agent_or_fatal(replay_agent); + const auto gpu_backend_available = + kernel_replay::blit::prepare(queue) == HSA_STATUS_SUCCESS; + if(!gpu_backend_available) + ROCP_WARNING << "kernel replay: internal blit kernel unavailable; trying " + "pinned-host snapshot backing"; + + // snap() holds the inventory read lock while this runs, so the blit must complete + // before it returns. + const auto capture_copy = [&](const auto& regions) { + auto pending = std::vector{}; + auto status = kernel_replay::blit::copy(queue, writer, regions, pending); + if(status != HSA_STATUS_SUCCESS || pending.empty()) return status; + return pending.back().wait(); + }; + // Save this agent's tracked device allocations so every pass runs against identical - // inputs. snap() returns ok=false if it could not capture the complete set (host memory - // pressure, a failed copy, or module-scope variables it could not enumerate). It logs - // which of those it hit. - const auto snapshot = kernel_replay::memory_snapshot::snap(replay_agent); + // inputs. Plan whole-region placement against GPU then pinned-host budgets and + // materialize at most one arena per tier before capture starts. + const auto snapshot = kernel_replay::memory_snapshot::snap( + kernel_replay::memory_snapshot::capture_context_t{ + replay_agent, + queue.get_agent().near_cpu(), + queue.get_agent().gpu_pool(), + queue.get_agent().cpu_pool(), + kernel_replay::memory_snapshot::capture_mode::automatic, + gpu_backend_available, + }, + capture_copy); // Snapshot incomplete: restoring a partial snapshot between passes would corrupt // application data, so decline replay. Close the CONFIG sequence, free our drain // signal, and run this dispatch once still under the writer lock with its original // completion signal, then return. - if(!snapshot.ok) + if(!snapshot.complete()) { LOG_FIRST_N(WARNING, 1) << "kernel replay: snapshot capture incomplete; running " "this dispatch once without replay"; @@ -1040,17 +1064,9 @@ WriteInterceptor(const void* packets, auto local_ctx_tls_guard = kernel_replay::scoped_local_context_control{context::get_active_contexts()}; - // Per-pass loop: PASS enter -> submit -> drain the async handler -> PASS exit -> ask - // the tool whether to continue -> restore device memory before the next pass. - for(uint64_t pass = 0;; ++pass) - { - const bool is_final = - !replay_plan.indefinite && (pass == replay_plan.total_passes - 1); - - auto pass_state = kernel_replay::pass_context_state_t{}; - kernel_replay::execute_pass_phase_enter( - replay_plan, pass, thr_id, internal_corr_id, ancestor_corr_id, pass_state); + auto restore_packets = std::vector{}; + auto submit_and_drain_pass = [&]() { process_packet_batch(packets_arr, 1, forward_to_writer, @@ -1059,27 +1075,63 @@ WriteInterceptor(const void* packets, // Drain this pass's async handler (separate HSA thread: reads counters, emits // records, releases signals/corr-id refs) before PASS EXIT / continue-decision / - // restore() / next submit, else we race its record delivery and reuse buffers and - // signals it still holds. This also implies GPU drain. Exactly one handler is in - // flight per pass (we drain before each submit, under the agent writer lock). - ROCP_FATAL_IF(queue.active_async_packets() > 1) << fmt::format( - "kernel replay: more than one async handler in flight during a replay pass"); + // next submit, else we race its record delivery and reuse buffers and signals it + // still holds. This also implies GPU drain. Exactly one handler is in flight per + // pass (we drain before each submit, under the agent writer lock). + ROCP_FATAL_IF(queue.active_async_packets() > 1) + << fmt::format("kernel replay: more than one async handler in flight during " + "a replay pass"); replay_drain_or_fatal(queue); + // The target kernel was queued behind the restore dependency barrier. Its + // completion proves the blit also completed, so the packet can be retired without + // another wait. + for(auto& restore : restore_packets) + ROCP_FATAL_IF(restore.retire() != HSA_STATUS_SUCCESS) + << "kernel replay: asynchronous restore blit failed"; + restore_packets.clear(); + }; + + // Per-pass loop: PASS enter -> optional restore -> submit -> drain the async handler -> + // PASS exit -> ask the tool whether to continue. + for(uint64_t pass = 0;; ++pass) + { + const bool is_final = + !replay_plan.indefinite && (pass == replay_plan.total_passes - 1); + + auto pass_state = kernel_replay::pass_context_state_t{}; + kernel_replay::execute_pass_phase_enter( + replay_plan, pass, thr_id, internal_corr_id, ancestor_corr_id, pass_state); + + if(pass == 0) + { + submit_and_drain_pass(); + } + else + { + // memory_snapshot::restore invokes this callback while holding the inventory + // read lock. Submit the target directly behind the blit and drain it before + // returning, so the lock remains held until both packets complete. + const auto batch_copy = [&](const auto& regions) { + auto status = + kernel_replay::blit::copy(queue, writer, regions, restore_packets); + if(status != HSA_STATUS_SUCCESS) return status; + + submit_and_drain_pass(); + return HSA_STATUS_SUCCESS; + }; + ROCP_FATAL_IF(!kernel_replay::memory_snapshot::restore(snapshot, batch_copy)) + << fmt::format("kernel replay: restore failed between passes (partial " + "host->device copy). Aborting rather than continuing with " + "corrupted device memory"); + } + kernel_replay::execute_pass_phase_exit(replay_plan, pass, pass_state); // Stop once the tool (or the fixed pass count) says we're done; the last executed // pass leaves device memory as the app expects, so no restore follows the break. if(!kernel_replay::should_continue_replay(replay_plan, pass_state, pass, is_final)) break; - - // Restore device memory between passes so the next pass sees identical inputs. - // A failed host->device copy leaves the snapshot only partially applied; continuing - // would submit the next pass over corrupted memory and (because the final pass - // skips restore) would also leave that corruption visible to the application. - ROCP_FATAL_IF(!kernel_replay::memory_snapshot::restore(snapshot)) << fmt::format( - "kernel replay: restore failed between passes (partial host->device copy); " - "aborting rather than continuing with corrupted device memory"); } kernel_replay::execute_config_phase_exit( diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/replay_window.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/replay_window.cpp index fd50501b7e6c..aad854166159 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/replay_window.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/replay_window.cpp @@ -140,36 +140,6 @@ replay_ring_submit(const Queue& queue, const rocprofiler_packet* packets, size_t return true; } -bool -replay_barrier_and_wait(const Queue& queue, hsa_signal_t completion) -{ - using namespace std::chrono_literals; - - auto barrier = hsa_barrier_and_packet_t{}; - barrier.header = HSA_PACKET_TYPE_BARRIER_AND << HSA_PACKET_HEADER_TYPE; - barrier.header |= 1 << HSA_PACKET_HEADER_BARRIER; - barrier.completion_signal = completion; - - auto packet = rocprofiler_packet{barrier}; - - const auto& core = queue.core_api(); - core.hsa_signal_store_screlease_fn(completion, 1); - - if(!replay_ring_submit(queue, &packet, 1)) return false; - - replay_wait_or_fatal( - [&]() { - return core.hsa_signal_wait_scacquire_fn(completion, - HSA_SIGNAL_CONDITION_EQ, - 0, - std::chrono::nanoseconds{5s}.count(), - HSA_WAIT_STATE_BLOCKED) == 0; - }, - "this queue's prior GPU work"); - - return true; -} - void set_interceptor_passthrough(bool enabled) { diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/replay_window.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/replay_window.hpp index c6958b6abadc..3d73a429280c 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/replay_window.hpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/replay_window.hpp @@ -125,11 +125,6 @@ replay_drain_agent_or_fatal(hsa_agent_t agent); bool replay_ring_submit(const Queue& queue, const rocprofiler_packet* packets, size_t count); -// Submit a barrier packet carrying `completion` and wait for it, fencing the CPU against all GPU -// work previously submitted to `queue`. Returns false if the submission itself failed. -bool -replay_barrier_and_wait(const Queue& queue, hsa_signal_t completion); - // While set on this thread, WriteInterceptor forwards packets untouched. Range replay's re-submit // path transforms the recorded packets once (via Queue::invoke_write_interceptor) and then writes // the result to the ring, which re-enters interception on this same thread; without this flag the diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/CMakeLists.txt b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/CMakeLists.txt index 014388973ef0..b96257cda45c 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/CMakeLists.txt +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/CMakeLists.txt @@ -1,15 +1,63 @@ # set(ROCPROFILER_LIB_KERNEL_REPLAY_HEADERS - kernel_replay.hpp local_context.hpp memory_snapshot.hpp memory_tracker.hpp - replay_callbacks.hpp utils.hpp) + blit-copy-kernel.hpp + blit-copy.hpp + kernel_replay.hpp + local_context.hpp + memory_snapshot.hpp + memory_tracker.hpp + replay_callbacks.hpp + snapshot-plan.hpp + snapshot-storage.hpp + utils.hpp) set(ROCPROFILER_LIB_KERNEL_REPLAY_SOURCES - kernel_replay.cpp local_context.cpp memory_snapshot.cpp memory_tracker.cpp - replay_callbacks.cpp utils.cpp) + blit-copy.cpp kernel_replay.cpp local_context.cpp memory_snapshot.cpp + memory_tracker.cpp replay_callbacks.cpp utils.cpp) target_sources( rocprofiler-sdk-object-library PRIVATE ${ROCPROFILER_LIB_KERNEL_REPLAY_SOURCES} ${ROCPROFILER_LIB_KERNEL_REPLAY_HEADERS}) +find_program( + kernel_replay_hipcc + NAMES hipcc + HINTS ${ROCM_PATH} ENV ROCM_PATH /opt/rocm + PATHS ${ROCM_PATH} ENV ROCM_PATH /opt/rocm + PATH_SUFFIXES bin llvm/bin REQUIRED) + +set(kernel_replay_blit_targets ${GPU_TARGETS}) +list(REMOVE_DUPLICATES kernel_replay_blit_targets) +set(kernel_replay_blit_code_objects) +foreach(target_id IN LISTS kernel_replay_blit_targets) + set(output_file ${CMAKE_CURRENT_BINARY_DIR}/${target_id}_kernel_replay_blit.hsaco) + add_custom_command( + OUTPUT ${output_file} + COMMAND + ${kernel_replay_hipcc} --genco --no-gpu-bundle-output + --offload-arch=${target_id} -mcode-object-version=5 -O3 -Xoffload-linker + --build-id=sha1 -o ${output_file} ${CMAKE_CURRENT_SOURCE_DIR}/blit-copy.hip + DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/blit-copy.hip + ${CMAKE_CURRENT_SOURCE_DIR}/blit-copy-kernel.hpp + COMMENT "Building ${target_id} kernel replay blit code object") + list(APPEND kernel_replay_blit_code_objects ${output_file}) +endforeach() + +add_custom_target(kernel-replay-blit-code-objects + DEPENDS ${kernel_replay_blit_code_objects}) +add_dependencies(rocprofiler-sdk-object-library kernel-replay-blit-code-objects) + +install( + FILES ${kernel_replay_blit_code_objects} + DESTINATION ${CMAKE_INSTALL_DATADIR}/rocprofiler-sdk/kernel-replay + COMPONENT core) + +target_compile_definitions( + rocprofiler-sdk-object-library + PRIVATE + ROCPROFILER_KERNEL_REPLAY_BLIT_BUILD_DIR="${CMAKE_CURRENT_BINARY_DIR}" + ROCPROFILER_KERNEL_REPLAY_BLIT_INSTALL_DIR="${CMAKE_INSTALL_FULL_DATADIR}/rocprofiler-sdk/kernel-replay" + ) + if(ROCPROFILER_BUILD_TESTS) add_subdirectory(tests) endif() diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy-kernel.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy-kernel.hpp new file mode 100644 index 000000000000..0a965a791e70 --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy-kernel.hpp @@ -0,0 +1,70 @@ +// MIT License +// +// Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +#pragma once + +#include + +namespace rocprofiler +{ +namespace kernel_replay +{ +namespace blit +{ +namespace kernel_abi +{ +inline constexpr auto bytes_per_item = std::uint64_t{16}; +inline constexpr auto workgroup_size = std::uint16_t{1024}; + +constexpr std::uint64_t +items_for_size(std::uint64_t size) +{ + return size / bytes_per_item + static_cast(size % bytes_per_item != 0); +} + +constexpr std::uint64_t +first_local_item(std::uint64_t thread, std::uint64_t stride, std::uint64_t logical_base) +{ + const auto base_mod = logical_base % stride; + return (thread >= base_mod) ? (thread - base_mod) : (stride - (base_mod - thread)); +} + +struct copy_descriptor_t +{ + std::uint64_t size = 0; + std::uint64_t source_address = 0; + std::uint64_t dest_address = 0; +}; + +struct kernel_args_t +{ + std::uint64_t descriptors_address = 0; + std::uint64_t descriptor_count = 0; + std::uint64_t launched_items = 0; +}; + +static_assert(sizeof(copy_descriptor_t) == 24); +static_assert(sizeof(kernel_args_t) == 24); +} // namespace kernel_abi +} // namespace blit +} // namespace kernel_replay +} // namespace rocprofiler diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy.cpp new file mode 100644 index 000000000000..19e0113df3f3 --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy.cpp @@ -0,0 +1,531 @@ +// MIT License +// +// Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +#include "lib/rocprofiler-sdk/kernel_replay/blit-copy.hpp" + +#include "lib/common/logging.hpp" +#include "lib/common/synchronized.hpp" +#include "lib/rocprofiler-sdk/hsa/agent_cache.hpp" +#include "lib/rocprofiler-sdk/hsa/hsa.hpp" +#include "lib/rocprofiler-sdk/hsa/queue.hpp" +#include "lib/rocprofiler-sdk/hsa/rocprofiler_packet.hpp" +#include "lib/rocprofiler-sdk/kernel_replay/blit-copy-kernel.hpp" + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace rocprofiler +{ +namespace kernel_replay +{ +namespace blit +{ +namespace +{ +constexpr auto kernel_name = "kernel_replay_blit"; + +using kernel_abi::copy_descriptor_t; +using kernel_abi::items_for_size; +using kernel_abi::kernel_args_t; +using kernel_abi::workgroup_size; + +struct kernel_info_t +{ + uint64_t object = 0; + uint32_t private_segment_size = 0; + uint32_t group_segment_size = 0; + uint32_t kernarg_size = 0; +}; + +struct kernel_state_t +{ + int file = -1; + hsa_code_object_reader_t reader{}; + hsa_executable_t executable{}; + kernel_info_t kernel{}; + hsa_signal_t completion{}; + void* kernarg = nullptr; + size_t kernarg_capacity = 0; + bool kernarg_on_device = false; + std::vector retired_kernargs{}; + std::atomic resource_in_use{false}; +}; + +using state_map_t = std::unordered_map; + +common::Synchronized& +states() +{ + // Keep one code object and reusable packet allocation per agent for process lifetime. HSA may + // be gone during static destruction, so deliberately do not destroy these objects there. + static auto* value = new common::Synchronized{}; + return *value; +} + +std::string +find_code_object(const hsa::AgentCache& agent) +{ + const auto filename = fmt::format("{}_kernel_replay_blit.hsaco", agent.name()); + for(const auto* directory : + {ROCPROFILER_KERNEL_REPLAY_BLIT_BUILD_DIR, ROCPROFILER_KERNEL_REPLAY_BLIT_INSTALL_DIR}) + { + auto path = fmt::format("{}/{}", directory, filename); + if(access(path.c_str(), R_OK) == 0) return path; + } + return {}; +} + +hsa_status_t +get_symbol_info(hsa_executable_symbol_t symbol, hsa_executable_symbol_info_t attribute, void* value) +{ + return hsa::get_core_table()->hsa_executable_symbol_get_info_fn(symbol, attribute, value); +} + +hsa_status_t +initialize_kernel(hsa_executable_t executable, + hsa_agent_t agent, + const char* name, + kernel_info_t& info) +{ + auto* core = hsa::get_core_table(); + auto symbol = hsa_executable_symbol_t{}; + auto status = core->hsa_executable_get_symbol_by_name_fn(executable, name, &agent, &symbol); + if(status != HSA_STATUS_SUCCESS) + { + const auto descriptor_name = fmt::format("{}.kd", name); + status = core->hsa_executable_get_symbol_by_name_fn( + executable, descriptor_name.c_str(), &agent, &symbol); + } + if(status != HSA_STATUS_SUCCESS) return status; + + status = get_symbol_info(symbol, HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_OBJECT, &info.object); + if(status != HSA_STATUS_SUCCESS) return status; + + status = get_symbol_info( + symbol, HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_PRIVATE_SEGMENT_SIZE, &info.private_segment_size); + if(status != HSA_STATUS_SUCCESS) return status; + status = get_symbol_info( + symbol, HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_GROUP_SEGMENT_SIZE, &info.group_segment_size); + if(status != HSA_STATUS_SUCCESS) return status; + status = get_symbol_info( + symbol, HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_KERNARG_SEGMENT_SIZE, &info.kernarg_size); + if(status != HSA_STATUS_SUCCESS) return status; + return (info.kernarg_size >= sizeof(void*)) ? HSA_STATUS_SUCCESS + : HSA_STATUS_ERROR_INVALID_ARGUMENT; +} + +hsa_status_t +initialize_state(const hsa::AgentCache& agent, kernel_state_t& state) +{ + auto* core = hsa::get_core_table(); + auto* ext = hsa::get_amd_ext_table(); + if(!core || !ext || !ext->hsa_amd_signal_create_fn) return HSA_STATUS_ERROR; + + const auto path = find_code_object(agent); + if(path.empty()) + { + ROCP_ERROR << fmt::format("kernel replay blit: no code object for agent '{}'", + agent.name()); + return HSA_STATUS_ERROR; + } + + state.file = open(path.c_str(), O_RDONLY); + if(state.file < 0) return HSA_STATUS_ERROR; + + auto status = core->hsa_code_object_reader_create_from_file_fn(state.file, &state.reader); + if(status != HSA_STATUS_SUCCESS) return status; + + status = core->hsa_executable_create_alt_fn( + HSA_PROFILE_FULL, HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT, nullptr, &state.executable); + if(status != HSA_STATUS_SUCCESS) return status; + + status = core->hsa_executable_load_agent_code_object_fn( + state.executable, agent.get_hsa_agent(), state.reader, nullptr, nullptr); + if(status != HSA_STATUS_SUCCESS) return status; + + status = core->hsa_executable_freeze_fn(state.executable, nullptr); + if(status != HSA_STATUS_SUCCESS) return status; + + const auto hsa_agent = agent.get_hsa_agent(); + status = initialize_kernel(state.executable, hsa_agent, kernel_name, state.kernel); + if(status != HSA_STATUS_SUCCESS) return status; + + status = ext->hsa_amd_signal_create_fn(1, 0, nullptr, 0, &state.completion); + if(status != HSA_STATUS_SUCCESS) return status; + + ROCP_INFO << fmt::format("kernel replay blit: loaded '{}' for agent '{}'", path, agent.name()); + return HSA_STATUS_SUCCESS; +} + +kernel_state_t* +get_state(const hsa::AgentCache& agent) +{ + auto* state = states().rlock([&](const auto& map) -> kernel_state_t* { + auto itr = map.find(agent.get_hsa_agent().handle); + return (itr == map.end()) ? nullptr : itr->second; + }); + if(state) return state; + + return states().wlock([&](auto& map) -> kernel_state_t* { + auto& slot = map[agent.get_hsa_agent().handle]; + if(slot) return slot; + + auto candidate = std::make_unique(); + if(initialize_state(agent, *candidate) != HSA_STATUS_SUCCESS) return nullptr; + slot = candidate.release(); + return slot; + }); +} + +// Small device allocations are carved out of larger blocks, so hostBaseAddress is set even when +// the CPU cannot map the block. The block's host-access flag is what decides it. +bool +host_writable(void* ptr) +{ +#if HSA_AMD_INTERFACE_VERSION_MAJOR > 1 || \ + (HSA_AMD_INTERFACE_VERSION_MAJOR == 1 && HSA_AMD_INTERFACE_VERSION_MINOR >= 18) + auto* ext = hsa::get_amd_ext_table(); + if(!ext || !ext->hsa_amd_pointer_info_fn) return false; + auto info = hsa_amd_pointer_info_t{}; + info.size = sizeof(info); + if(ext->hsa_amd_pointer_info_fn(ptr, &info, nullptr, nullptr, nullptr) != HSA_STATUS_SUCCESS) + return false; + return (info.alloc_flags & HSA_AMD_POINTER_INFO_ALLOC_FLAG_HOST_ACCESS) != 0; +#else + (void) ptr; + return false; +#endif +} + +// Every wave reads the blit's arguments and copy descriptors, and reading them from host memory +// makes each restore depend on PCIe read latency. Keep them in device memory the CPU can write, +// and fall back to the host kernarg pool when it cannot. Returns the buffer and whether it is in +// device memory. +std::pair +allocate_kernarg(const hsa::AgentCache& agent, size_t size) +{ + auto* ext = hsa::get_amd_ext_table(); + void* ptr = nullptr; + if(ext->hsa_amd_memory_pool_allocate_fn( + agent.gpu_pool(), size, HSA_AMD_MEMORY_POOL_EXECUTABLE_FLAG, &ptr) == HSA_STATUS_SUCCESS) + { + if(host_writable(ptr)) return {ptr, true}; + ext->hsa_amd_memory_pool_free_fn(ptr); + } + + ptr = nullptr; + if(ext->hsa_amd_memory_pool_allocate_fn(agent.kernarg_pool(), size, 0, &ptr) != + HSA_STATUS_SUCCESS) + return {nullptr, false}; + const auto hsa_agent = agent.get_hsa_agent(); + if(ext->hsa_amd_agents_allow_access_fn(1, &hsa_agent, nullptr, ptr) != HSA_STATUS_SUCCESS) + { + ext->hsa_amd_memory_pool_free_fn(ptr); + return {nullptr, false}; + } + return {ptr, false}; +} + +// Writes to device memory cross PCIe and may still be in flight when the packet is published, so +// fence them and read the last byte back before the doorbell. +void +write_kernarg(void* dest, const std::vector& image, bool on_device) +{ + std::memcpy(dest, image.data(), image.size()); + if(!on_device || image.empty()) return; + + auto* last = static_cast(dest) + image.size() - 1; + std::atomic_thread_fence(std::memory_order_seq_cst); + *last = image.back(); + std::atomic_thread_fence(std::memory_order_seq_cst); + [[maybe_unused]] const auto landed = *last; +} +} // namespace + +hsa_status_t +prepare(const hsa::Queue& queue) +{ + return get_state(queue.get_agent()) ? HSA_STATUS_SUCCESS : HSA_STATUS_ERROR; +} + +struct packet_info::implementation +{ + std::array packets{}; + kernel_state_t* state = nullptr; + hsa_signal_t completion{}; + void* kernarg = nullptr; + bool submitted = false; + bool completed = false; + hsa_status_t status = HSA_STATUS_SUCCESS; + + ~implementation() + { + if(state) state->resource_in_use.store(false, std::memory_order_release); + } +}; + +packet_info::packet_info() = default; + +packet_info::packet_info(std::unique_ptr impl) +: m_impl{std::move(impl)} +{} + +packet_info::~packet_info() +{ + if(!m_impl) return; + if(m_impl->submitted && !m_impl->completed) wait(); +} + +packet_info::packet_info(packet_info&&) noexcept = default; + +packet_info& +packet_info::operator=(packet_info&& rhs) noexcept +{ + if(this == &rhs) return *this; + if(m_impl && m_impl->submitted && !m_impl->completed) wait(); + m_impl = std::move(rhs.m_impl); + return *this; +} + +const void* +packet_info::get_raw() +{ + if(!m_impl) return nullptr; + m_impl->submitted = true; + return m_impl->packets.data(); +} + +uint64_t +packet_info::packet_count() const +{ + return m_impl ? m_impl->packets.size() : 0; +} + +hsa_status_t +packet_info::wait() +{ + if(!m_impl) return HSA_STATUS_ERROR_INVALID_ARGUMENT; + if(m_impl->completed) return m_impl->status; + if(!m_impl->submitted) return HSA_STATUS_ERROR_INVALID_ARGUMENT; + + using namespace std::chrono_literals; + auto value = hsa_signal_value_t{1}; + for(int i = 0; i < 12 && value >= 1; ++i) + { + value = hsa::get_core_table()->hsa_signal_wait_scacquire_fn( + m_impl->completion, + HSA_SIGNAL_CONDITION_LT, + 1, + std::chrono::nanoseconds{5s}.count(), + HSA_WAIT_STATE_BLOCKED); + } + + m_impl->status = (value == 0) ? HSA_STATUS_SUCCESS : HSA_STATUS_ERROR; + m_impl->completed = true; + return m_impl->status; +} + +hsa_status_t +packet_info::retire() +{ + if(!m_impl || !m_impl->submitted) return HSA_STATUS_ERROR_INVALID_ARGUMENT; + if(m_impl->completed) return m_impl->status; + + const auto value = hsa::get_core_table()->hsa_signal_load_scacquire_fn(m_impl->completion); + m_impl->status = (value == 0) ? HSA_STATUS_SUCCESS : HSA_STATUS_ERROR; + m_impl->completed = true; + return m_impl->status; +} + +std::optional +create(const hsa::Queue& queue, const std::vector& regions) +{ + if(regions.empty()) return std::nullopt; + const auto& agent = queue.get_agent(); + + auto* state = get_state(agent); + auto* core = hsa::get_core_table(); + auto* ext = hsa::get_amd_ext_table(); + if(!state || !core || !ext || !core->hsa_signal_store_screlease_fn) return std::nullopt; + + auto descriptor_data = std::vector{}; + descriptor_data.reserve(regions.size()); + auto total_items = uint64_t{0}; + for(const auto& region : regions) + { + if(region.size == 0 || !region.source || !region.dest) return std::nullopt; + const auto item_count = items_for_size(region.size); + if(item_count > std::numeric_limits::max() - total_items) return std::nullopt; + descriptor_data.emplace_back(copy_descriptor_t{ + region.size, + reinterpret_cast(region.source), + reinterpret_cast(region.dest), + }); + total_items += item_count; + } + + constexpr auto max_workgroups = uint64_t{std::numeric_limits::max()} / workgroup_size; + const auto requested_workgroups = + total_items / workgroup_size + static_cast(total_items % workgroup_size != 0); + constexpr auto blocks_per_cu = uint64_t{2}; + const auto cu_count = agent.get_rocp_agent()->cu_count; + if(cu_count == 0) return std::nullopt; + const auto occupancy_workgroups = static_cast(cu_count) * blocks_per_cu; + const auto launched_workgroups = + std::min({requested_workgroups, max_workgroups, occupancy_workgroups}); + if(launched_workgroups == 0) return std::nullopt; + const auto launched_items = launched_workgroups * workgroup_size; + + const auto& kernel = state->kernel; + if(kernel.kernarg_size < sizeof(kernel_args_t)) return std::nullopt; + const auto descriptor_offset = (kernel.kernarg_size + 15) & ~size_t{15}; + if(regions.size() > std::numeric_limits::max() || + regions.size() > + (std::numeric_limits::max() - descriptor_offset) / sizeof(copy_descriptor_t)) + return std::nullopt; + const auto allocation_size = descriptor_offset + regions.size() * sizeof(copy_descriptor_t); + + auto impl = std::make_unique(); + auto expected = false; + if(!state->resource_in_use.compare_exchange_strong(expected, true, std::memory_order_acq_rel)) + return std::nullopt; + + impl->state = state; + impl->completion = state->completion; + + if(allocation_size > state->kernarg_capacity) + { + auto new_capacity = std::max(size_t{4096}, state->kernarg_capacity); + while(new_capacity < allocation_size) + { + if(new_capacity > std::numeric_limits::max() / 2) + { + new_capacity = allocation_size; + break; + } + new_capacity *= 2; + } + + const auto [new_kernarg, on_device] = allocate_kernarg(agent, new_capacity); + if(!new_kernarg) return std::nullopt; + ROCP_INFO << fmt::format("kernel replay blit: {}B argument buffer in {} memory for agent " + "'{}'", + new_capacity, + on_device ? "device" : "host", + agent.name()); + + if(state->kernarg) state->retired_kernargs.emplace_back(state->kernarg); + state->kernarg = new_kernarg; + state->kernarg_capacity = new_capacity; + state->kernarg_on_device = on_device; + } + impl->kernarg = state->kernarg; + + auto image = std::vector{}; + try + { + image.resize(allocation_size); + } catch(const std::bad_alloc&) + { + return std::nullopt; + } + const auto descriptor_address = + reinterpret_cast(static_cast(impl->kernarg) + descriptor_offset); + const auto kernel_args = kernel_args_t{ + descriptor_address, + static_cast(descriptor_data.size()), + launched_items, + }; + std::memcpy(image.data(), &kernel_args, sizeof(kernel_args)); + std::memcpy(image.data() + descriptor_offset, + descriptor_data.data(), + descriptor_data.size() * sizeof(descriptor_data[0])); + + core->hsa_signal_store_screlease_fn(impl->completion, 1); + write_kernarg(impl->kernarg, image, state->kernarg_on_device); + + const auto grid_size_x = static_cast(launched_items); + + auto packet = hsa_kernel_dispatch_packet_t{}; + packet.header = HSA_PACKET_TYPE_KERNEL_DISPATCH << HSA_PACKET_HEADER_TYPE; + packet.header |= 1 << HSA_PACKET_HEADER_BARRIER; + packet.header |= HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE; + packet.header |= HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE; + packet.setup = 1; + packet.workgroup_size_x = workgroup_size; + packet.workgroup_size_y = 1; + packet.workgroup_size_z = 1; + packet.grid_size_x = grid_size_x; + packet.grid_size_y = 1; + packet.grid_size_z = 1; + packet.private_segment_size = kernel.private_segment_size; + packet.group_segment_size = kernel.group_segment_size; + packet.kernel_object = kernel.object; + packet.kernarg_address = impl->kernarg; + packet.completion_signal = impl->completion; + impl->packets[0] = hsa::rocprofiler_packet{packet}; + + auto barrier = hsa_barrier_and_packet_t{}; + barrier.header = HSA_PACKET_TYPE_BARRIER_AND << HSA_PACKET_HEADER_TYPE; + barrier.header |= 1 << HSA_PACKET_HEADER_BARRIER; + barrier.header |= HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE; + barrier.header |= HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE; + barrier.dep_signal[0] = impl->completion; + impl->packets[1] = hsa::rocprofiler_packet{barrier}; + + return packet_info{std::move(impl)}; +} + +hsa_status_t +copy(const hsa::Queue& queue, + hsa_amd_queue_intercept_packet_writer writer, + const std::vector& regions, + std::vector& pending) +{ + if(!writer) return HSA_STATUS_ERROR_INVALID_ARGUMENT; + if(regions.empty()) return HSA_STATUS_SUCCESS; + if(!pending.empty()) return HSA_STATUS_ERROR; + + auto packet = create(queue, regions); + if(!packet) return HSA_STATUS_ERROR; + writer(packet->get_raw(), packet->packet_count()); + pending.emplace_back(std::move(*packet)); + return HSA_STATUS_SUCCESS; +} +} // namespace blit +} // namespace kernel_replay +} // namespace rocprofiler diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy.hip b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy.hip new file mode 100644 index 000000000000..993d0d4831d4 --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy.hip @@ -0,0 +1,89 @@ +// MIT License +// +// Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +#include "blit-copy-kernel.hpp" + +#include + +#include +#include + +using rocprofiler::kernel_replay::blit::kernel_abi::bytes_per_item; +using rocprofiler::kernel_replay::blit::kernel_abi::copy_descriptor_t; +using rocprofiler::kernel_replay::blit::kernel_abi::first_local_item; +using rocprofiler::kernel_replay::blit::kernel_abi::items_for_size; +using rocprofiler::kernel_replay::blit::kernel_abi::kernel_args_t; +using rocprofiler::kernel_replay::blit::kernel_abi::workgroup_size; + +using copy_vector_t = std::uint32_t __attribute__((ext_vector_type(4))); +using global_byte_t = std::byte __attribute__((address_space(1))); +using global_descriptor_t = copy_descriptor_t __attribute__((address_space(1))); + +static_assert(bytes_per_item == sizeof(copy_vector_t)); + +__device__ __forceinline__ void +copy_item(const copy_descriptor_t& copy, std::uint64_t item_offset) +{ + const auto byte_offset = item_offset * bytes_per_item; + const auto* src_bytes = reinterpret_cast(copy.source_address); + auto* dst_bytes = reinterpret_cast(copy.dest_address); + + if(copy.size - byte_offset >= bytes_per_item) + { + auto value = copy_vector_t{}; + __builtin_memcpy(&value, src_bytes + byte_offset, sizeof(value)); + __builtin_memcpy(dst_bytes + byte_offset, &value, sizeof(value)); + return; + } + + for(auto i = byte_offset; i < copy.size; ++i) + dst_bytes[i] = src_bytes[i]; +} + +extern "C" __global__ __launch_bounds__(workgroup_size) void kernel_replay_blit(kernel_args_t args) +{ + const auto* descriptors = + reinterpret_cast(args.descriptors_address); + const auto thread = static_cast(threadIdx.x) + + static_cast(workgroup_size) * blockIdx.x; + const auto stride = args.launched_items; + + auto logical_base = std::uint64_t{0}; + for(auto region = std::uint64_t{0}; region < args.descriptor_count; ++region) + { + auto copy = copy_descriptor_t{}; + __builtin_memcpy(©, descriptors + region, sizeof(copy)); + const auto item_count = items_for_size(copy.size); + // Thread `t` owns compact logical items `t + k * stride`. Rotate that assignment by this + // region's prefix to obtain its first physical 16-byte item in the discontiguous + // allocation. + auto item_offset = first_local_item(thread, stride, logical_base); + + while(item_offset < item_count) + { + copy_item(copy, item_offset); + if(item_count - item_offset <= stride) break; + item_offset += stride; + } + logical_base += item_count; + } +} diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy.hpp new file mode 100644 index 000000000000..2cafcebafdf3 --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/blit-copy.hpp @@ -0,0 +1,93 @@ +// MIT License +// +// Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +#pragma once + +#include +#include + +#include +#include +#include +#include +#include + +namespace rocprofiler +{ +namespace hsa +{ +class AgentCache; +class Queue; +} // namespace hsa + +namespace kernel_replay +{ +namespace blit +{ +struct copy_region_t +{ + size_t size = 0; + const void* source = nullptr; + void* dest = nullptr; +}; + +hsa_status_t +prepare(const hsa::Queue& queue); + +class packet_info +{ +public: + packet_info(); + ~packet_info(); + + packet_info(const packet_info&) = delete; + packet_info& operator=(const packet_info&) = delete; + packet_info(packet_info&&) noexcept; + packet_info& operator=(packet_info&&) noexcept; + + const void* get_raw(); + uint64_t packet_count() const; + hsa_status_t wait(); + hsa_status_t retire(); + +private: + struct implementation; + explicit packet_info(std::unique_ptr impl); + + std::unique_ptr m_impl; + + friend std::optional create(const hsa::Queue&, const std::vector&); +}; + +std::optional +create(const hsa::Queue& queue, const std::vector& regions); + +// Submit all GPU-backed regions through one mixed-size raw AQL blit packet. The completion signal +// and kernarg storage are reused from per-agent state after the dependent target kernel drains. +hsa_status_t +copy(const hsa::Queue& queue, + hsa_amd_queue_intercept_packet_writer writer, + const std::vector& regions, + std::vector& pending); +} // namespace blit +} // namespace kernel_replay +} // namespace rocprofiler diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/memory_snapshot.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/memory_snapshot.cpp index 5acce99160d2..ebbcf17ecd6a 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/memory_snapshot.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/memory_snapshot.cpp @@ -23,18 +23,32 @@ #include "lib/rocprofiler-sdk/kernel_replay/memory_snapshot.hpp" #include "lib/common/logging.hpp" +#include "lib/common/synchronized.hpp" +#include "lib/rocprofiler-sdk/agent.hpp" #include "lib/rocprofiler-sdk/code_object/code_object.hpp" #include "lib/rocprofiler-sdk/code_object/hsa/code_object.hpp" +#include "lib/rocprofiler-sdk/hsa/agent_cache.hpp" #include "lib/rocprofiler-sdk/hsa/hsa.hpp" #include "lib/rocprofiler-sdk/kernel_replay/memory_tracker.hpp" +#include "lib/rocprofiler-sdk/kernel_replay/snapshot-plan.hpp" +#include "lib/rocprofiler-sdk/kernel_replay/snapshot-storage.hpp" #include #include +#include +#include +#include #include +#include +#include +#include #include #include +#include #include +#include +#include #include namespace rocprofiler @@ -45,6 +59,8 @@ namespace memory_snapshot { namespace { +using blit::copy_region_t; + hsa_status_t dma_copy(void* dst, const void* src, size_t n) { @@ -53,23 +69,6 @@ dma_copy(void* dst, const void* src, size_t n) return core->hsa_memory_copy_fn(dst, src, n); } -/// @brief Run @p copy under the tracker read lock, but only while [@p gpu_addr, +@p size) is still -/// a live tracked allocation of >= @p size bytes, so a concurrent free (write lock) can't retire it -/// mid-copy. Direction is caller-supplied: snap reads device->host, restore writes host->device. -/// @return @p copy's status, or @c std::nullopt if the region was freed/shrunk since it was -/// recorded. -template -std::optional -with_inventory_check(void* gpu_addr, size_t size, CopyFn&& copy) -{ - return memory_tracker::inventory().rlock( - [&](const memory_tracker::tracked_map_t& map) -> std::optional { - auto itr = map.find(gpu_addr); - if(itr == map.end() || itr->second.size < size) return std::nullopt; - return copy(); - }); -} - // A module-scope variable (__device__ / __constant__ global) discovered in a loaded executable. struct module_variable_t { @@ -128,6 +127,26 @@ collect_module_variable(hsa_executable_t, hsa_agent_t, hsa_executable_symbol_t s return HSA_STATUS_SUCCESS; } + // HIP emits a one-byte compilation-unit marker into every code object. It is linker/runtime + // metadata rather than application state, and snapshotting it would make the replay-owned blit + // code object recursively add another restore region. + if(size == 1) + { + constexpr auto hip_cuid_prefix = std::string_view{"__hip_cuid_"}; + uint32_t name_length = 0; + if(core->hsa_executable_symbol_get_info_fn(symbol, + HSA_EXECUTABLE_SYMBOL_INFO_NAME_LENGTH, + &name_length) == HSA_STATUS_SUCCESS && + name_length >= hip_cuid_prefix.size()) + { + auto name = std::vector(name_length + 1, '\0'); + if(core->hsa_executable_symbol_get_info_fn( + symbol, HSA_EXECUTABLE_SYMBOL_INFO_NAME, name.data()) == HSA_STATUS_SUCCESS && + std::string_view{name.data()}.substr(0, hip_cuid_prefix.size()) == hip_cuid_prefix) + return HSA_STATUS_SUCCESS; + } + } + if(addr == 0 || size == 0) return HSA_STATUS_SUCCESS; // A variable above the cap is skipped, which means a kernel's writes to it leak across replay @@ -178,12 +197,415 @@ discover_module_variables(hsa_agent_t agent) }); return scan; } + +// At most one retained high-water arena per pool. A larger returned arena replaces and frees the +// old one, and smaller returns are freed. A snapshot that needs more than the retained arena frees +// it before allocating, so retained memory never exceeds the largest snapshot so far. +using retained_backing_t = std::unordered_map; + +common::Synchronized& +retained_backing() +{ + // Keep backing allocations for process lifetime. HSA may already be unavailable during static + // destruction, so do not attempt to free them there. + static auto* value = new common::Synchronized{}; + return *value; +} + +void +free_pool_block(block_t block) noexcept +{ + if(!block) return; + if(auto* ext = hsa::get_amd_ext_table(); ext && ext->hsa_amd_memory_pool_free_fn) + ext->hsa_amd_memory_pool_free_fn(block.data); +} + +// Returns the pool's retained arena when it can hold `size`. One that is too small is freed +// instead, so the larger allocation that follows can use its memory. +block_t +take_retained_block(hsa_amd_memory_pool_t pool, size_t size) noexcept +{ + if(pool.handle == 0 || size == 0) return {}; + auto taken = block_t{}; + auto evicted = block_t{}; + try + { + retained_backing().wlock([&](auto& retained) { + auto itr = retained.find(pool.handle); + if(itr == retained.end() || !itr->second) return; + auto& out = itr->second.size >= size ? taken : evicted; + out = std::exchange(itr->second, block_t{}); + }); + } catch(...) + { + return {}; + } + free_pool_block(evicted); + return taken; +} + +size_t +retained_capacity(hsa_amd_memory_pool_t pool) noexcept +{ + if(pool.handle == 0) return 0; + try + { + return retained_backing().rlock([&](const auto& retained) { + auto itr = retained.find(pool.handle); + return itr == retained.end() ? size_t{0} : itr->second.size; + }); + } catch(...) + { + return 0; + } +} + +void +retain_block(hsa_amd_memory_pool_t pool, block_t block) noexcept +{ + if(pool.handle == 0 || !block) + { + free_pool_block(block); + return; + } + + // Map insertion may allocate. Destruction must remain noexcept, so retain nothing if the map + // cannot grow and release the arena directly. + try + { + retained_backing().wlock([&](auto& retained) { + auto& current = retained[pool.handle]; + if(!current || block.size > current.size) std::swap(current, block); + }); + } catch(...) + {} + free_pool_block(block); +} + +constexpr size_t arena_alignment = 16; + +std::optional +pool_allocation_max(hsa_amd_memory_pool_t pool) +{ + auto* ext = hsa::get_amd_ext_table(); + if(pool.handle == 0 || !ext || !ext->hsa_amd_memory_pool_get_info_fn) return std::nullopt; + + bool allowed = false; + if(ext->hsa_amd_memory_pool_get_info_fn( + pool, HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_ALLOWED, &allowed) != HSA_STATUS_SUCCESS || + !allowed) + return std::nullopt; + + size_t maximum = 0; + if(ext->hsa_amd_memory_pool_get_info_fn( + pool, HSA_AMD_MEMORY_POOL_INFO_ALLOC_MAX_SIZE, &maximum) != HSA_STATUS_SUCCESS) + return std::nullopt; + return maximum; +} + +bool +pool_accepts(hsa_amd_memory_pool_t pool, size_t size) +{ + auto maximum = pool_allocation_max(pool); + return maximum && size <= *maximum; +} + +std::optional +pool_size(hsa_amd_memory_pool_t pool) +{ + auto* ext = hsa::get_amd_ext_table(); + if(pool.handle == 0 || !ext || !ext->hsa_amd_memory_pool_get_info_fn) return std::nullopt; + + size_t total = 0; + if(ext->hsa_amd_memory_pool_get_info_fn(pool, HSA_AMD_MEMORY_POOL_INFO_SIZE, &total) != + HSA_STATUS_SUCCESS) + return std::nullopt; + return total; +} + +std::optional +gpu_available_bytes(hsa_agent_t agent) +{ + auto* core = hsa::get_core_table(); + if(!core || !core->hsa_agent_get_info_fn) return std::nullopt; + + uint64_t available = 0; + const auto status = core->hsa_agent_get_info_fn( + agent, static_cast(HSA_AMD_AGENT_INFO_MEMORY_AVAIL), &available); + if(status != HSA_STATUS_SUCCESS) return std::nullopt; + return available; +} + +std::optional +host_available_bytes() +{ + auto input = std::ifstream{"/proc/meminfo"}; + if(!input) return std::nullopt; + + auto key = std::string{}; + auto unit = std::string{}; + uint64_t value_kib = 0; + while(input >> key >> value_kib >> unit) + { + if(key != "MemAvailable:") continue; + if(unit != "kB" || value_kib > std::numeric_limits::max() / 1024) + return std::nullopt; + return value_kib * 1024; + } + return std::nullopt; +} + +size_t +gpu_plan_budget(const capture_context_t& ctx) +{ + if(ctx.mode == capture_mode::force_pinned || !ctx.gpu_backend_available) return 0; + auto maximum = pool_allocation_max(ctx.gpu_pool); + if(!maximum) return 0; + + // Hold at most half of the device's memory, and leave the application at least 5% of what is + // free. The retained arena counts as available because a larger snapshot frees it first. + auto budget = static_cast(*maximum); + if(auto total = pool_size(ctx.gpu_pool)) budget = std::min(budget, *total / 2); + if(auto available = gpu_available_bytes(ctx.agent)) + { + const auto retained = static_cast(retained_capacity(ctx.gpu_pool)); + budget = std::min(budget, retained + (*available - *available / 20)); + } + return static_cast(budget); +} + +size_t +host_plan_budget(const capture_context_t& ctx) +{ + if(ctx.mode == capture_mode::force_gpu) return 0; + auto maximum = pool_allocation_max(ctx.pinned_pool); + if(!maximum) return 0; + // Pinned memory cannot be swapped out, so use at most half of what the host has available. The + // retained arena counts as available, as on the GPU. + if(auto available = host_available_bytes()) + return std::min(*maximum, + retained_capacity(ctx.pinned_pool) + static_cast(*available / 2)); + return *maximum; +} + +// Allocates one tier's backing from an HSA pool. The returned block releases itself according to +// its retain mode. +struct hsa_pool_allocator +{ + using handle_type = owned_block_t; + + storage_kind kind = storage_kind::none; + retain_mode retain = retain_mode::per_replay; + hsa_amd_memory_pool_t pool = {}; + hsa_agent_t agent = {}; + hsa_agent_t peer_agent = {}; + bool available = false; + + std::optional allocate(size_t size) const + { + if(!available || pool.handle == 0 || size == 0) return std::nullopt; + + // A retained arena that is large enough consumes no fresh capacity. Consult it before + // asking whether a new allocation of this size is currently possible. + if(retain == retain_mode::across_replays) + { + if(auto retained = take_retained_block(pool, size); retained) + return owned_block_t{retained, kind, pool, retain}; + } + + if(!pool_accepts(pool, size)) return std::nullopt; + + auto* ext = hsa::get_amd_ext_table(); + if(!ext || !ext->hsa_amd_memory_pool_allocate_fn || !ext->hsa_amd_agents_allow_access_fn || + !ext->hsa_amd_memory_pool_free_fn) + return std::nullopt; + + void* ptr = nullptr; + auto status = ext->hsa_amd_memory_pool_allocate_fn(pool, size, 0, &ptr); + if(status != HSA_STATUS_SUCCESS || !ptr) return std::nullopt; + + if(peer_agent.handle != 0) + { + status = ext->hsa_amd_agents_allow_access_fn(1, &peer_agent, nullptr, ptr); + if(status != HSA_STATUS_SUCCESS) + { + ext->hsa_amd_memory_pool_free_fn(ptr); + return std::nullopt; + } + } + + status = ext->hsa_amd_agents_allow_access_fn(1, &agent, nullptr, ptr); + if(status == HSA_STATUS_SUCCESS) + return owned_block_t{block_t{size, ptr}, kind, pool, retain}; + + ext->hsa_amd_memory_pool_free_fn(ptr); + return std::nullopt; + } +}; + +using hsa_tier = + storage::arena>; +using snapshot_storage = storage::fallback; + +snapshot_storage +make_snapshot_storage(const capture_context_t& ctx, storage_budgets_t budgets) +{ + const auto gpu = hsa_pool_allocator{.kind = storage_kind::gpu_local, + .retain = retain_mode::across_replays, + .pool = ctx.gpu_pool, + .agent = ctx.agent, + .available = ctx.gpu_backend_available}; + const auto pinned = hsa_pool_allocator{.kind = storage_kind::pinned_host, + .retain = retain_mode::across_replays, + .pool = ctx.pinned_pool, + .agent = ctx.agent, + .peer_agent = ctx.cpu_agent, + .available = true}; + return snapshot_storage{ + hsa_tier{{}, {budgets.gpu_bytes, arena_alignment, gpu}}, + hsa_tier{{}, {budgets.host_bytes, arena_alignment, pinned}}, + storage::null_tier{}, + }; +} + +bool +region_is_live(const snapshot_region_t& region, const memory_tracker::tracked_map_t& inventory) +{ + if(region.liveness == liveness_kind::module_variable) return true; + if(region.liveness == liveness_kind::retired) return false; + auto itr = inventory.find(region.live_address); + return itr != inventory.end() && itr->second.size >= region.size; +} + +copy_region_t +capture_copy(const snapshot_t& snapshot, const snapshot_extent_t& extent) +{ + const auto& region = snapshot.regions.at(extent.region_index); + const auto& segment = snapshot.segments.at(extent.segment_index); + auto* backing = static_cast(segment.backing.block().data); + return copy_region_t{ + extent.size, + static_cast(region.live_address) + extent.region_offset, + backing + extent.segment_offset, + }; +} + +copy_region_t +restore_copy(const snapshot_t& snapshot, const snapshot_extent_t& extent) +{ + const auto& region = snapshot.regions.at(extent.region_index); + const auto& segment = snapshot.segments.at(extent.segment_index); + auto* backing = static_cast(segment.backing.block().data); + return copy_region_t{ + extent.size, + backing + extent.segment_offset, + static_cast(region.live_address) + extent.region_offset, + }; +} } // namespace -device_snapshot_t +owned_block_t::owned_block_t(block_t block, + storage_kind kind, + hsa_amd_memory_pool_t pool, + retain_mode retain) +: m_block{block} +, m_kind{kind} +, m_pool{pool} +, m_retain{retain} +{} + +owned_block_t::~owned_block_t() { reset(); } + +owned_block_t::owned_block_t(owned_block_t&& rhs) noexcept +: m_block{std::exchange(rhs.m_block, block_t{})} +, m_kind{std::exchange(rhs.m_kind, storage_kind::none)} +, m_pool{std::exchange(rhs.m_pool, hsa_amd_memory_pool_t{.handle = 0})} +, m_retain{std::exchange(rhs.m_retain, retain_mode::per_replay)} +{} + +owned_block_t& +owned_block_t::operator=(owned_block_t&& rhs) noexcept +{ + if(this == &rhs) return *this; + reset(); + m_block = std::exchange(rhs.m_block, block_t{}); + m_kind = std::exchange(rhs.m_kind, storage_kind::none); + m_pool = std::exchange(rhs.m_pool, hsa_amd_memory_pool_t{.handle = 0}); + m_retain = std::exchange(rhs.m_retain, retain_mode::per_replay); + return *this; +} + +void +owned_block_t::reset() noexcept +{ + if(!m_block) return; + if(m_retain == retain_mode::across_replays) + retain_block(m_pool, m_block); + else if(auto* ext = hsa::get_amd_ext_table(); ext && ext->hsa_amd_memory_pool_free_fn) + ext->hsa_amd_memory_pool_free_fn(m_block.data); + + m_block = {}; + m_kind = storage_kind::none; + m_pool = hsa_amd_memory_pool_t{.handle = 0}; + m_retain = retain_mode::per_replay; +} + +size_t +snapshot_t::footprint_bytes() const +{ + size_t total = 0; + for(const auto& region : regions) + if(region.liveness != liveness_kind::retired) total += region.size; + return total; +} + +storage_kind +snapshot_t::storage() const +{ + if(segments.empty()) return storage_kind::none; + const auto first = segments.front().backing.kind(); + for(const auto& segment : segments) + if(segment.backing.kind() != first) return storage_kind::mixed; + return first; +} + +snapshot_t snap(hsa_agent_t agent) { - device_snapshot_t out{}; + auto cache = agent::get_agent_cache(agent); + if(!cache) + { + auto out = snapshot_t{}; + out.status = capture_status::unavailable; + return out; + } + + return snap(capture_context_t{ + agent, + cache->near_cpu(), + cache->gpu_pool(), + cache->cpu_pool(), + capture_mode::force_pinned, + false, + }); +} + +snapshot_t +snap(const capture_context_t& ctx) +{ + return snap(ctx, storage_budgets_t{gpu_plan_budget(ctx), host_plan_budget(ctx)}); +} + +snapshot_t +snap(const capture_context_t& ctx, const batch_copy_fn_t& batch_copy) +{ + return snap(ctx, storage_budgets_t{gpu_plan_budget(ctx), host_plan_budget(ctx)}, batch_copy); +} + +snapshot_t +snap(const capture_context_t& ctx, storage_budgets_t budgets, const batch_copy_fn_t& batch_copy) +{ + auto out = snapshot_t{}; // Note: trackable allocations carrying HSA_AMD_MEMORY_POOL_EXECUTABLE_FLAG are recorded in // memory_tracker::unsupported_executable() and omitted from the main inventory. Declining @@ -200,15 +622,14 @@ snap(hsa_agent_t agent) auto scan = module_variable_scan_t{}; try { - inventory = memory_tracker::snap_inventory(agent); - scan = discover_module_variables(agent); - - out.blocks.reserve(inventory.size() + scan.found.size()); + inventory = memory_tracker::snap_inventory(ctx.agent); + scan = discover_module_variables(ctx.agent); + out.regions.reserve(inventory.size() + scan.found.size()); } catch(const std::bad_alloc&) { LOG_FIRST_N(WARNING, 1) << "kernel-replay snapshot: out of memory reserving metadata; " "declining replay for this dispatch"; - out.ok = false; + out.status = capture_status::unavailable; return out; } @@ -219,159 +640,287 @@ snap(hsa_agent_t agent) LOG_FIRST_N(WARNING, 1) << "kernel-replay snapshot: could not enumerate module-scope " "variables for every loaded executable; declining replay for " "this dispatch"; - out.ok = false; + out.status = capture_status::invalid_plan; return out; } - const auto& module_vars = scan.found; - - /// @brief Capture one region (device->host) into the snapshot. - /// @retval false The snapshot is incomplete because a host allocation or the DMA copy failed. - /// The caller must decline replay rather than restore partial state. - /// @retval true The region was captured. For tracker regions it may instead have been dropped - /// after being freed or shrunk since snap_inventory(). That drop is safe because restore() runs - /// the same liveness check, so a region gone now could not have been restored anyway. - auto capture = - [&](void* gpu_addr, size_t size, std::string_view what, bool from_tracker) -> bool { - if(size == 0) return true; - - auto blk = mem_block_t{}; - blk.gpu_addr = gpu_addr; - blk.from_tracker = from_tracker; - try - { - blk.host_copy.resize(size); - } catch(const std::bad_alloc&) - { - ROCP_WARNING << fmt::format("kernel-replay snapshot: host allocation of {} bytes " - "failed for {} {} (memory pressure)", - size, - what, - gpu_addr); - return false; - } + for(const auto& [ptr, size] : inventory) + { + if(size == 0) continue; + out.regions.emplace_back(snapshot_region_t{size, ptr, liveness_kind::tracked_allocation}); + } - // snap_inventory() released the tracker lock before returning. A host thread calling - // hsa_amd_memory_pool_free / hsa_memory_free can retire this allocation while we - // read it, because the alloc/free wrappers are not covered by the per-agent replay lock. - // The with_inventory_check helper re-checks liveness and copies under the read lock, the - // same guard restore() applies to the write direction. A nullopt result means the region - // was freed or shrunk after snap_inventory(), so we drop it rather than read retired device - // memory. Module variables live in the loaded executable, not the tracker. They are always - // present, so we copy them directly. - const auto st = - from_tracker - ? with_inventory_check( - gpu_addr, - size, - [&] { return dma_copy(blk.host_copy.data(), gpu_addr, size); }) - : std::optional{dma_copy(blk.host_copy.data(), gpu_addr, size)}; - - if(!st) - { - ROCP_INFO << fmt::format("kernel-replay snapshot: {} {} ({}B) was retired before " - "capture, dropping from snapshot", - what, - gpu_addr, - size); - return true; - } + for(const auto& var : scan.found) + { + if(var.size == 0) continue; + out.regions.emplace_back( + snapshot_region_t{var.size, var.gpu_addr, liveness_kind::module_variable}); + } + + if(out.regions.empty()) + { + out.status = capture_status::complete; + return out; + } + + // Address order makes placement under memory pressure reproducible from run to run. + std::sort(out.regions.begin(), out.regions.end(), [](const auto& lhs, const auto& rhs) { + return std::less<>{}(lhs.live_address, rhs.live_address); + }); - if(*st != HSA_STATUS_SUCCESS) + auto chain = make_snapshot_storage(ctx, budgets); + auto placements = std::optional>{}; + try + { + auto requests = std::vector{}; + requests.reserve(out.regions.size()); + for(size_t i = 0; i < out.regions.size(); ++i) + requests.push_back(planning::request_t{i, 0, out.regions[i].size}); + + out.segments.reserve(snapshot_storage::tier_count); + placements = planning::plan_and_commit(chain, requests); + if(placements) { - ROCP_WARNING << fmt::format( - "kernel-replay snapshot: device->host copy failed for {} {} ({}B)", - what, - gpu_addr, - size); - return false; + chain.release_backings([&](size_t tier, size_t size, owned_block_t backing) { + out.segments.push_back(storage_segment_t{size, tier, std::move(backing)}); + }); } + } catch(const std::bad_alloc&) + { + LOG_FIRST_N(WARNING, 1) << "kernel-replay snapshot: out of memory planning snapshot " + "storage; declining replay for this dispatch"; + out.regions.clear(); + out.segments.clear(); + out.status = capture_status::unavailable; + return out; + } - out.blocks.push_back(std::move(blk)); - return true; - }; + if(!placements) + { + ROCP_WARNING << fmt::format( + "kernel-replay snapshot: could not place {} bytes within {} GPU " + "and {} pinned-host bytes on agent {}; declining replay", + out.footprint_bytes(), + budgets.gpu_bytes, + budgets.host_bytes, + ctx.agent.handle); + out.regions.clear(); + out.status = capture_status::unavailable; + return out; + } - for(const auto& [ptr, size] : inventory) + try { - if(!capture(ptr, size, "region", /*from_tracker=*/true)) + auto segment_for_tier = std::array, snapshot_storage::tier_count>{}; + for(size_t i = 0; i < out.segments.size(); ++i) + segment_for_tier[out.segments[i].tier] = i; + + out.extents.reserve(placements->size()); + for(const auto& placement : *placements) { - out.ok = false; - return out; + const auto segment = segment_for_tier[placement.tier]; + if(!segment) + { + out.regions.clear(); + out.segments.clear(); + out.status = capture_status::invalid_plan; + return out; + } + out.extents.push_back(snapshot_extent_t{placement.size, + placement.region_index, + placement.region_offset, + *segment, + placement.backing_offset}); } + } catch(const std::bad_alloc&) + { + out.regions.clear(); + out.segments.clear(); + out.extents.clear(); + out.status = capture_status::unavailable; + return out; } - // Module-scope variables (__device__ / __constant__ globals) live in the loaded executable's - // data segment, not in the allocation tracker, so capture them here too. Restored via the same - // per-block host->device copy as tracked allocations (see restore()). - for(const auto& var : module_vars) + auto gpu_copies = std::vector{}; + try { - if(!capture(var.gpu_addr, var.size, "module variable", /*from_tracker=*/false)) + gpu_copies.reserve(out.extents.size()); + } catch(const std::bad_alloc&) + { + out.regions.clear(); + out.segments.clear(); + out.extents.clear(); + out.status = capture_status::unavailable; + return out; + } + + // One inventory read lock covers every copy, including the batch, so no tracked region can be + // freed while it is being read. + size_t captured_regions = 0; + const auto status = memory_tracker::inventory().rlock([&](const auto& map) { + for(size_t region_index = 0; region_index < out.regions.size(); ++region_index) { - out.ok = false; - return out; + auto& region = out.regions[region_index]; + if(!region_is_live(region, map)) + { + ROCP_INFO << fmt::format("kernel-replay snapshot: region {} ({}B) was retired " + "before capture, dropping from snapshot", + region.live_address, + region.size); + region.liveness = liveness_kind::retired; + continue; + } + for(const auto& extent : out.extents) + { + if(extent.region_index != region_index) continue; + const auto copy = capture_copy(out, extent); + const auto kind = out.segments.at(extent.segment_index).backing.kind(); + if(batch_copy && kind == storage_kind::gpu_local) + { + gpu_copies.emplace_back(copy); + continue; + } + const auto copy_status = dma_copy(copy.dest, copy.source, copy.size); + if(copy_status != HSA_STATUS_SUCCESS) return copy_status; + } + ++captured_regions; } + return gpu_copies.empty() ? HSA_STATUS_SUCCESS : batch_copy(gpu_copies); + }); + + if(status != HSA_STATUS_SUCCESS) + { + ROCP_WARNING << fmt::format("kernel-replay snapshot: capture copy failed for agent {}", + ctx.agent.handle); + out.regions.clear(); + out.extents.clear(); + out.segments.clear(); + out.status = capture_status::copy_failed; + return out; } - ROCP_INFO << fmt::format("kernel-replay snapshot: captured {} regions (tracked allocations + " - "module variables) for agent {}", - out.blocks.size(), - agent.handle); + out.status = capture_status::complete; + + size_t gpu_bytes = 0; + size_t host_bytes = 0; + for(const auto& segment : out.segments) + { + if(segment.backing.kind() == storage_kind::gpu_local) + gpu_bytes += segment.size; + else if(segment.backing.kind() == storage_kind::pinned_host) + host_bytes += segment.size; + } + ROCP_INFO << fmt::format("kernel-replay snapshot: captured {} regions using {} GPU bytes and " + "{} pinned-host bytes for agent {}", + captured_regions, + gpu_bytes, + host_bytes, + ctx.agent.handle); return out; } bool -restore(const device_snapshot_t& snapshot) +restore(const snapshot_t& snapshot) { - size_t restored = 0; - for(const auto& blk : snapshot.blocks) - { - hsa_status_t status = HSA_STATUS_SUCCESS; + if(!snapshot.complete()) return false; + if(snapshot.extents.empty()) return true; - if(blk.from_tracker) + size_t restored = 0; + const auto status = memory_tracker::inventory().rlock([&](const auto& inventory) { + for(const auto& extent : snapshot.extents) { - // Re-check liveness and copy under the read lock so a concurrent free cannot race the - // check and the write (see with_inventory_check). A nullopt result means the region is - // no longer a live allocation of at least its size. It was freed or reallocated after - // snap, so skip it (benign -- not a restore failure). - const auto st = with_inventory_check(blk.gpu_addr, blk.host_copy.size(), [&] { - return dma_copy(blk.gpu_addr, blk.host_copy.data(), blk.host_copy.size()); - }); - - if(!st) + const auto& region = snapshot.regions.at(extent.region_index); + if(!region_is_live(region, inventory)) { ROCP_WARNING << fmt::format( "kernel-replay restore: skipping region {} ({}B) that is no longer live", - blk.gpu_addr, - blk.host_copy.size()); + region.live_address, + region.size); continue; } - status = *st; - } - else - { - // Module variable: lives in the loaded executable, always present, so restore directly. - status = dma_copy(blk.gpu_addr, blk.host_copy.data(), blk.host_copy.size()); + + const auto copy = restore_copy(snapshot, extent); + const auto st = dma_copy(copy.dest, copy.source, copy.size); + if(st != HSA_STATUS_SUCCESS) return st; + ++restored; } + return HSA_STATUS_SUCCESS; + }); - if(status != HSA_STATUS_SUCCESS) + if(status != HSA_STATUS_SUCCESS) + { + ROCP_ERROR << fmt::format("kernel-replay restore: copy failed after {}/{} extents", + restored, + snapshot.extents.size()); + return false; + } + ROCP_INFO << fmt::format( + "kernel-replay restore: restored {}/{} extents", restored, snapshot.extents.size()); + return true; +} + +bool +restore(const snapshot_t& snapshot, const batch_copy_fn_t& batch_copy) +{ + if(!snapshot.complete()) return false; + + auto gpu_regions = std::vector{}; + try + { + gpu_regions.reserve(snapshot.extents.size()); + } catch(const std::bad_alloc&) + { + return false; + } + size_t restored = 0; + size_t skipped = 0; + + const auto status = memory_tracker::inventory().rlock([&](const auto& map) { + for(const auto& extent : snapshot.extents) { - // A live region failed to copy. Further passes would observe mutated state, and the - // final pass (which deliberately skips restore) would leave that corruption visible to - // the application. Abort: a partial restore cannot be undone. - ROCP_ERROR << fmt::format( - "kernel-replay restore: host->device copy failed for region {} ({}B); aborting " - "restore after {}/{} regions", - blk.gpu_addr, - blk.host_copy.size(), - restored, - snapshot.blocks.size()); - return false; + const auto& region = snapshot.regions.at(extent.region_index); + if(!region_is_live(region, map)) + { + ROCP_WARNING << fmt::format( + "kernel-replay restore: skipping region {} ({}B) that is no longer live", + region.live_address, + region.size); + ++skipped; + continue; + } + + auto copy = restore_copy(snapshot, extent); + const auto kind = snapshot.segments.at(extent.segment_index).backing.kind(); + if(kind == storage_kind::gpu_local) + { + gpu_regions.emplace_back(copy); + } + else + { + auto copy_status = dma_copy(copy.dest, copy.source, copy.size); + if(copy_status != HSA_STATUS_SUCCESS) return copy_status; + ++restored; + } } - ++restored; + + auto copy_status = batch_copy(gpu_regions); + if(copy_status != HSA_STATUS_SUCCESS) return copy_status; + restored += gpu_regions.size(); + return HSA_STATUS_SUCCESS; + }); + + if(status != HSA_STATUS_SUCCESS) + { + ROCP_ERROR << fmt::format("kernel-replay restore: batch copy failed after {}/{} regions", + restored, + snapshot.extents.size()); + return false; } - ROCP_INFO << fmt::format( - "kernel-replay restore: restored {}/{} regions", restored, snapshot.blocks.size()); + ROCP_INFO << fmt::format("kernel-replay restore: restored {}/{} regions ({} skipped)", + restored, + snapshot.extents.size(), + skipped); return true; } } // namespace memory_snapshot diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/memory_snapshot.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/memory_snapshot.hpp index a82c6414e920..26556d6660f0 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/memory_snapshot.hpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/memory_snapshot.hpp @@ -22,63 +22,179 @@ #pragma once +#include "lib/rocprofiler-sdk/kernel_replay/blit-copy.hpp" + #include +#include #include +#include +#include #include namespace rocprofiler { namespace kernel_replay { -// Minimal save/restore of device memory for kernel replay. -// -// snap(agent) copies every tracked device allocation owned by `agent` into host memory; restore() -// copies it back. This keeps each replay pass running against identical inputs, scoped to one agent -// so concurrent replays on other agents are unaffected. The implementation is deliberately simple: -// a full copy of each region (no dirty-page diffing) of directly-allocated device memory. +// Save/restore of one agent's tracked device memory for kernel replay. A pure planning pass assigns +// each complete region to GPU-local or pinned-host storage, then capture materializes at most one +// arena per tier before reading application memory. namespace memory_snapshot { -// Saved copy of a single device allocation. -struct mem_block_t +enum class storage_kind : uint8_t +{ + none, + gpu_local, + pinned_host, + mixed, +}; + +enum class liveness_kind : uint8_t +{ + tracked_allocation, + module_variable, + retired, +}; + +enum class capture_status : uint8_t +{ + complete, + unavailable, + invalid_plan, + copy_failed, +}; + +enum class capture_mode : uint8_t +{ + automatic, + force_gpu, + force_pinned, +}; + +struct block_t +{ + size_t size = 0; + void* data = nullptr; + + explicit operator bool() const { return data != nullptr; } +}; + +// What happens to a snapshot's backing once the replay that captured it finishes. +enum class retain_mode : uint8_t +{ + per_replay, // freed back to its pool + across_replays, // kept as its pool's high-water arena for later replays +}; + +// RAII handle for one materialized HSA arena. Releasing it frees the arena or keeps it for later +// replays, according to its retain mode. +class owned_block_t +{ +public: + owned_block_t() = default; + owned_block_t(block_t, storage_kind, hsa_amd_memory_pool_t, retain_mode); + ~owned_block_t(); + + owned_block_t(const owned_block_t&) = delete; + owned_block_t& operator=(const owned_block_t&) = delete; + owned_block_t(owned_block_t&& rhs) noexcept; + owned_block_t& operator=(owned_block_t&& rhs) noexcept; + + block_t block() const { return m_block; } + storage_kind kind() const { return m_kind; } + explicit operator bool() const { return static_cast(m_block); } + + void reset() noexcept; + +private: + block_t m_block{}; + storage_kind m_kind = storage_kind::none; + hsa_amd_memory_pool_t m_pool{.handle = 0}; + retain_mode m_retain = retain_mode::per_replay; +}; + +struct snapshot_region_t +{ + size_t size = 0; + void* live_address = nullptr; + liveness_kind liveness = liveness_kind::tracked_allocation; +}; + +struct storage_segment_t +{ + size_t size = 0; + size_t tier = 0; + owned_block_t backing{}; +}; + +struct snapshot_extent_t +{ + size_t size = 0; + + size_t region_index = 0; + size_t region_offset = 0; + + size_t segment_index = 0; + size_t segment_offset = 0; +}; + +struct capture_context_t +{ + hsa_agent_t agent{}; + hsa_agent_t cpu_agent{}; + hsa_amd_memory_pool_t gpu_pool{}; + hsa_amd_memory_pool_t pinned_pool{}; + capture_mode mode = capture_mode::automatic; + bool gpu_backend_available = true; +}; + +struct storage_budgets_t { - void* gpu_addr = nullptr; // live device allocation base pointer - std::vector host_copy; // pre-kernel contents held in host memory - // true = from the allocation tracker; re-check liveness before restoring (it can be freed). - // false = module-scope variable in a loaded executable; always live, so restore - // unconditionally. - bool from_tracker = false; + size_t gpu_bytes = 0; + size_t host_bytes = 0; }; -// A captured set of device allocations (host-side copies) for a single agent. -struct device_snapshot_t +// A captured set of device allocations for a single agent. +struct snapshot_t { - std::vector blocks; - // false => capture was incomplete: host memory pressure, a failed device->host copy, or HSA - // would not enumerate the module-scope variables of every loaded executable. The snapshot must - // not be used to restore -- a partial restore corrupts application data -- so the caller should - // decline replay and run the dispatch once instead. - bool ok = true; + std::vector regions{}; + std::vector segments{}; + std::vector extents{}; + capture_status status{capture_status::unavailable}; - bool empty() const { return blocks.empty(); } + bool complete() const { return status == capture_status::complete; } + size_t footprint_bytes() const; + storage_kind storage() const; }; -// Copy (device->host) every tracked allocation owned by `agent`, plus every module-scope variable -// (__device__ / __constant__ global) visible to `agent` in the loaded executables. On success the -// returned snapshot has ok==true. It returns early with ok==false, so the caller can decline replay -// rather than restore a partial snapshot, when a region cannot be captured (host memory pressure -- -// the host buffer allocation fails -- or a failed copy) or when the module-scope variables of a -// loaded executable cannot be enumerated. Host memory pressure is never fatal: a dispatch that -// cannot be snapshotted runs once instead of aborting the application. -device_snapshot_t +using batch_copy_fn_t = std::function&)>; + +// Plan whole-region placement against queried GPU and host budgets, materialize at most one arena +// per tier, and capture only after every required arena exists. +snapshot_t +snap(const capture_context_t&); + +// GPU-local storage is captured through one batch blit callback, which must finish the copy before +// it returns. Pinned-host storage is captured synchronously. +snapshot_t +snap(const capture_context_t&, const batch_copy_fn_t& batch_copy); + +snapshot_t +snap(const capture_context_t&, storage_budgets_t, const batch_copy_fn_t& batch_copy = {}); + +// Direct-test convenience: force the nearest CPU agent's pinned HSA pool. +snapshot_t snap(hsa_agent_t agent); -// Copy each saved region host->device. Returns true when every live region was restored. -// A region that was freed after snap is skipped (benign). A failed host->device DMA copy -// returns false immediately: the snapshot is then only partially applied, and the caller must -// abort replay rather than submit another pass over corrupted device memory. +// Copy each saved region back to its live device allocation. A region freed after snap is skipped. +// A failed copy returns false immediately because the snapshot is then only partially applied. +bool +restore(const snapshot_t& snapshot); + +// GPU-local storage restores through one batch blit callback. Pinned-host storage restores +// synchronously and invokes the callback with no copy regions so the caller submits the target. bool -restore(const device_snapshot_t& snapshot); +restore(const snapshot_t& snapshot, const batch_copy_fn_t& batch_copy); } // namespace memory_snapshot } // namespace kernel_replay } // namespace rocprofiler diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/snapshot-plan.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/snapshot-plan.hpp new file mode 100644 index 000000000000..820246cbeb98 --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/snapshot-plan.hpp @@ -0,0 +1,136 @@ +// MIT License +// +// Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace rocprofiler +{ +namespace kernel_replay +{ +namespace memory_snapshot +{ +namespace planning +{ +// A piece of one region to place. Lower priority values are placed first. Callers that need +// non-contiguous placement, such as hot or dirty pages, emit several requests for one region. +struct request_t +{ + size_t region_index = 0; + size_t region_offset = 0; + size_t size = 0; + uint32_t priority = 0; +}; + +// The bytes one tier took from the front of a request, and where they start in its backing. +struct reservation_t +{ + size_t size = 0; + size_t backing_offset = 0; +}; + +struct placement_t +{ + size_t region_index = 0; + size_t region_offset = 0; + size_t size = 0; + size_t tier = 0; + size_t backing_offset = 0; +}; + +// Admission policies return how many bytes from the front of a request one tier takes, at most +// `available`. Anything they leave is offered to the next tier. + +// The request fits whole or moves on to the next tier. +struct whole_region +{ + size_t admit(const request_t& request, size_t available) const + { + return request.size <= available ? request.size : 0; + } +}; + +// The largest prefix that fits. The split point is rounded down to a granule in region +// coordinates, so both sides stay aligned for DMA and for the blit's 16-byte items. +struct prefix_chunking +{ + size_t granularity = 4096; + + size_t admit(const request_t& request, size_t available) const + { + if(request.size <= available) return request.size; + const auto granule = std::max(granularity, 1); + const auto split = (request.region_offset + available) / granule * granule; + return split > request.region_offset ? split - request.region_offset : 0; + } +}; + +// Places every request through `storage` in priority order. On failure the storage is reset. +template +std::optional> +plan(Storage& storage, std::vector requests) +{ + std::sort(requests.begin(), requests.end(), [](const request_t& lhs, const request_t& rhs) { + return std::tie(lhs.priority, lhs.region_index, lhs.region_offset) < + std::tie(rhs.priority, rhs.region_index, rhs.region_offset); + }); + + auto placements = std::vector{}; + placements.reserve(requests.size()); + for(const auto& request : requests) + { + if(!storage.place(request, placements)) + { + storage.reset(); + return std::nullopt; + } + } + return placements; +} + +// Tier capacities are advisory until backing exists. If a tier cannot materialize what it was +// given, it is disabled and every request is planned again. +template +std::optional> +plan_and_commit(Storage& storage, const std::vector& requests) +{ + for(size_t attempt = 0; attempt < Storage::tier_count; ++attempt) + { + auto placements = plan(storage, requests); + if(!placements) return std::nullopt; + + const auto failed = storage.commit(); + if(!failed) return placements; + storage.disable(*failed); + } + return std::nullopt; +} +} // namespace planning +} // namespace memory_snapshot +} // namespace kernel_replay +} // namespace rocprofiler diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/snapshot-storage.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/snapshot-storage.hpp new file mode 100644 index 000000000000..091d33f0fab5 --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/snapshot-storage.hpp @@ -0,0 +1,220 @@ +// MIT License +// +// Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +#pragma once + +#include "lib/rocprofiler-sdk/kernel_replay/snapshot-plan.hpp" + +#include +#include +#include +#include +#include +#include +#include + +namespace rocprofiler +{ +namespace kernel_replay +{ +namespace memory_snapshot +{ +namespace storage +{ +// Packs ranges back to back, so the backing only has to cover the high-water mark. `used` never +// exceeds `capacity`, which keeps the offset arithmetic free of overflow. +template +struct bump_suballocator +{ + using handle_type = typename BackingAllocator::handle_type; + + size_t capacity = 0; + size_t alignment = 1; + BackingAllocator backing; + size_t used = 0; + std::optional handle = {}; + + size_t padding() const + { + const auto remainder = used % alignment; + return remainder == 0 ? 0 : alignment - remainder; + } + + size_t available() const + { + const auto free_bytes = capacity - used; + return free_bytes > padding() ? free_bytes - padding() : 0; + } + + size_t reserve(size_t size) + { + const auto offset = used + padding(); + used = offset + size; + return offset; + } + + bool commit() + { + if(used == 0) return true; + handle = backing.allocate(used); + return handle.has_value(); + } + + void reset() + { + handle.reset(); + used = 0; + } + + size_t backing_size() const { return used; } + std::optional release_backing() { return std::exchange(handle, std::nullopt); } +}; + +// One tier. The admission policy decides how much of a request to take; the suballocator decides +// where it goes and obtains the backing at commit. +template +struct arena +{ + static constexpr bool terminal = false; + + Admission admission; + Suballocator suballocator; + bool enabled = true; + + std::optional reserve(const planning::request_t& request) + { + if(!enabled) return std::nullopt; + const auto available = suballocator.available(); + // Clamped so a misbehaving policy cannot overrun the arena. + const auto size = std::min({admission.admit(request, available), request.size, available}); + if(size == 0) return std::nullopt; + return planning::reservation_t{size, suballocator.reserve(size)}; + } + + bool commit() { return suballocator.commit(); } + void reset() { suballocator.reset(); } + void disable() { enabled = false; } +}; + +// Anything that reaches this tier has nowhere to go, so planning fails. +struct null_tier +{ + static constexpr bool terminal = true; + + std::optional reserve(const planning::request_t&) + { + return std::nullopt; + } + + bool commit() { return true; } + void reset() {} + void disable() {} +}; + +// Offers what is left of each request to the tiers in order, ending in the null tier. +template +struct fallback +{ + static constexpr size_t tier_count = sizeof...(Tiers); + static_assert(std::tuple_element_t>::terminal, + "the last tier must be the null tier"); + + std::tuple tiers; + + explicit fallback(Tiers... t) + : tiers{std::move(t)...} + {} + + // A tier is asked again until it declines, so it can take one request in several pieces. + // Returns false when bytes are left over at the null tier. + bool place(planning::request_t request, std::vector& placements) + { + for_each([&](auto& tier, size_t index) { + while(request.size > 0) + { + const auto taken = tier.reserve(request); + if(!taken) break; + placements.push_back(planning::placement_t{request.region_index, + request.region_offset, + taken->size, + index, + taken->backing_offset}); + request.region_offset += taken->size; + request.size -= taken->size; + } + }); + return request.size == 0; + } + + // On failure, releases every tier and returns the index of the one that could not + // materialize its backing. + std::optional commit() + { + auto failed = std::optional{}; + for_each([&](auto& tier, size_t index) { + if(!failed && !tier.commit()) failed = index; + }); + if(failed) reset(); + return failed; + } + + void reset() + { + for_each([](auto& tier, size_t) { tier.reset(); }); + } + + void disable(size_t index) + { + for_each([index](auto& tier, size_t tier_index) { + if(tier_index == index) tier.disable(); + }); + } + + // Calls `fn(tier_index, size, handle)` for every committed tier, moving the backing out of the + // storage. + template + void release_backings(Fn&& fn) + { + for_each([&](auto& tier, size_t index) { + if constexpr(!std::decay_t::terminal) + { + if(auto handle = tier.suballocator.release_backing()) + fn(index, tier.suballocator.backing_size(), std::move(*handle)); + } + }); + } + + template + void for_each(Fn&& fn) + { + std::apply( + [&](auto&... tier) { + size_t index = 0; + (fn(tier, index++), ...); + }, + tiers); + } +}; +} // namespace storage +} // namespace memory_snapshot +} // namespace kernel_replay +} // namespace rocprofiler diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/CMakeLists.txt b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/CMakeLists.txt index 1c8ebbd73005..416301315d4f 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/CMakeLists.txt +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/CMakeLists.txt @@ -106,6 +106,44 @@ rocprofiler_add_unit_test( "LD_LIBRARY_PATH=${PROJECT_BINARY_DIR}/${CMAKE_INSTALL_LIBDIR}:${CMAKE_LIBRARY_OUTPUT_DIRECTORY}:${CMAKE_INSTALL_PREFIX}/lib:${ROCM_PATH}/lib:$ENV{LD_LIBRARY_PATH}" FAIL_REGULAR_EXPRESSION "${ROCPROFILER_DEFAULT_FAIL_REGEX}") +# Compact logical-item distribution is pure host arithmetic shared with the blit kernel. +# This target checks cross-descriptor thread assignment without requiring a GPU. +add_executable(kernel-replay-blit-copy-test) +target_sources(kernel-replay-blit-copy-test PRIVATE blit_copy.cpp) + +target_link_libraries( + kernel-replay-blit-copy-test + PRIVATE rocprofiler-sdk::rocprofiler-sdk-common-library + rocprofiler-sdk::rocprofiler-sdk-static-library GTest::gtest + GTest::gtest_main) + +rocprofiler_add_unit_test( + TARGET kernel-replay-blit-copy-test + SOURCES blit_copy.cpp + TIMEOUT 60 + ENVIRONMENT + "LD_LIBRARY_PATH=${PROJECT_BINARY_DIR}/${CMAKE_INSTALL_LIBDIR}:${CMAKE_LIBRARY_OUTPUT_DIRECTORY}:${CMAKE_INSTALL_PREFIX}/lib:${ROCM_PATH}/lib:$ENV{LD_LIBRARY_PATH}" + FAIL_REGULAR_EXPRESSION "${ROCPROFILER_DEFAULT_FAIL_REGEX}") + +# Planning, variadic fallback, commit/rollback, and backing allocator composition are +# deliberately GPU-independent. +add_executable(kernel-replay-snapshot-storage-test) +target_sources(kernel-replay-snapshot-storage-test PRIVATE snapshot_storage.cpp) + +target_link_libraries( + kernel-replay-snapshot-storage-test + PRIVATE rocprofiler-sdk::rocprofiler-sdk-common-library + rocprofiler-sdk::rocprofiler-sdk-static-library GTest::gtest + GTest::gtest_main) + +rocprofiler_add_unit_test( + TARGET kernel-replay-snapshot-storage-test + SOURCES snapshot_storage.cpp + TIMEOUT 60 + ENVIRONMENT + "LD_LIBRARY_PATH=${PROJECT_BINARY_DIR}/${CMAKE_INSTALL_LIBDIR}:${CMAKE_LIBRARY_OUTPUT_DIRECTORY}:${CMAKE_INSTALL_PREFIX}/lib:${ROCM_PATH}/lib:$ENV{LD_LIBRARY_PATH}" + FAIL_REGULAR_EXPRESSION "${ROCPROFILER_DEFAULT_FAIL_REGEX}") + # Layout and C/C++ agreement checks for the public kernel-replay header. These are # properties of the header alone: no GPU, no HSA, and no link against the SDK, so the test # runs anywhere. diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/blit_copy.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/blit_copy.cpp new file mode 100644 index 000000000000..b5e9a9e7b5ae --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/blit_copy.cpp @@ -0,0 +1,118 @@ +// MIT License +// +// Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +#include "lib/rocprofiler-sdk/kernel_replay/blit-copy-kernel.hpp" +#include "lib/rocprofiler-sdk/kernel_replay/blit-copy.hpp" + +#include + +#include +#include +#include +#include +#include + +namespace +{ +namespace abi = rocprofiler::kernel_replay::blit::kernel_abi; + +TEST(kernel_replay_blit_copy, descriptors_use_size_source_dest_order) +{ + EXPECT_EQ(offsetof(rocprofiler::kernel_replay::blit::copy_region_t, size), 0); + EXPECT_LT(offsetof(rocprofiler::kernel_replay::blit::copy_region_t, size), + offsetof(rocprofiler::kernel_replay::blit::copy_region_t, source)); + EXPECT_LT(offsetof(rocprofiler::kernel_replay::blit::copy_region_t, source), + offsetof(rocprofiler::kernel_replay::blit::copy_region_t, dest)); + + EXPECT_EQ(offsetof(abi::copy_descriptor_t, size), 0); + EXPECT_EQ(offsetof(abi::copy_descriptor_t, source_address), sizeof(std::uint64_t)); + EXPECT_EQ(offsetof(abi::copy_descriptor_t, dest_address), 2 * sizeof(std::uint64_t)); +} + +TEST(kernel_replay_blit_copy, compact_tid_assignment_crosses_descriptor_boundaries) +{ + // Logical regions [0, 20), [20, 80), [80, 100), [100, 145), [145, 200) + // model the discontiguous 20/60/20/45/55-item example. With 100 workers, + // every thread must own logical item `tid` and `tid + 100`, even when those + // items belong to different descriptors. + constexpr auto stride = std::uint64_t{100}; + constexpr auto item_count = std::array{20, 60, 20, 45, 55}; + constexpr auto total = std::uint64_t{200}; + + auto owner = std::array{}; + auto assignments = std::array>, stride>{}; + owner.fill(-1); + + for(auto thread = std::uint64_t{0}; thread < stride; ++thread) + { + auto logical_base = std::uint64_t{0}; + for(auto region = size_t{0}; region < item_count.size(); ++region) + { + auto local = abi::first_local_item(thread, stride, logical_base); + while(local < item_count[region]) + { + const auto logical = logical_base + local; + ASSERT_LT(logical, total); + EXPECT_EQ(owner.at(logical), -1) + << "logical item " << logical << " assigned more than once"; + owner.at(logical) = static_cast(thread); + assignments.at(thread).emplace_back(region, local); + + if(item_count[region] - local <= stride) break; + local += stride; + } + logical_base += item_count[region]; + } + } + + for(auto logical = std::uint64_t{0}; logical < total; ++logical) + EXPECT_EQ(owner.at(logical), static_cast(logical % stride)) + << "logical item " << logical; + + for(auto thread = std::uint64_t{0}; thread < stride; ++thread) + EXPECT_EQ(assignments.at(thread).size(), 2) << "thread " << thread; + + EXPECT_EQ(assignments.at(0), (std::vector>{{0, 0}, {3, 0}})); + EXPECT_EQ(assignments.at(19), + (std::vector>{{0, 19}, {3, 19}})); + EXPECT_EQ(assignments.at(20), (std::vector>{{1, 0}, {3, 20}})); + EXPECT_EQ(assignments.at(44), + (std::vector>{{1, 24}, {3, 44}})); + EXPECT_EQ(assignments.at(45), (std::vector>{{1, 25}, {4, 0}})); + EXPECT_EQ(assignments.at(79), + (std::vector>{{1, 59}, {4, 34}})); + EXPECT_EQ(assignments.at(80), (std::vector>{{2, 0}, {4, 35}})); + EXPECT_EQ(assignments.at(99), + (std::vector>{{2, 19}, {4, 54}})); +} + +TEST(kernel_replay_blit_copy, partial_region_uses_one_tail_item) +{ + EXPECT_EQ(abi::items_for_size(0), 0); + EXPECT_EQ(abi::items_for_size(1), 1); + EXPECT_EQ(abi::items_for_size(15), 1); + EXPECT_EQ(abi::items_for_size(16), 1); + EXPECT_EQ(abi::items_for_size(17), 2); + EXPECT_EQ(abi::items_for_size(31), 2); + EXPECT_EQ(abi::items_for_size(32), 2); +} +} // namespace diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/snap_bandwidth.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/snap_bandwidth.cpp index 2fc466300647..8cd212cdce9a 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/snap_bandwidth.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/snap_bandwidth.cpp @@ -27,6 +27,7 @@ #include "snap_kernels.hpp" #include "lib/rocprofiler-sdk/agent.hpp" +#include "lib/rocprofiler-sdk/hsa/agent_cache.hpp" #include "lib/rocprofiler-sdk/hsa/hsa.hpp" #include "lib/rocprofiler-sdk/kernel_replay/memory_snapshot.hpp" #include "lib/rocprofiler-sdk/kernel_replay/memory_tracker.hpp" @@ -43,6 +44,7 @@ #include #include #include +#include #include using namespace rocprofiler; @@ -110,6 +112,21 @@ gpu_agent() return hsa_agent_t{.handle = 0}; } +std::optional +capture_context(hsa_agent_t agent, msnp::capture_mode mode) +{ + auto cache = rocprofiler::agent::get_agent_cache(agent); + if(!cache) return std::nullopt; + return msnp::capture_context_t{ + agent, + cache->near_cpu(), + cache->gpu_pool(), + cache->cpu_pool(), + mode, + true, + }; +} + void sync_ok() { @@ -128,12 +145,9 @@ read_device(const float* d, int n) } size_t -snapshot_footprint_bytes(const msnp::device_snapshot_t& snap) +snapshot_footprint_bytes(const msnp::snapshot_t& snap) { - size_t total = 0; - for(const auto& block : snap.blocks) - total += block.host_copy.size(); - return total; + return snap.footprint_bytes(); } // Conservative host-link floor (GB/s). Intentionally low for heterogeneous CI including @@ -167,14 +181,20 @@ struct snap_restore_timing_t }; snap_restore_timing_t -measure_snap_restore_once(hsa_agent_t agent, float* buffer, int n_elems) +measure_snap_restore_once(hsa_agent_t agent, + float* buffer, + int n_elems, + msnp::capture_mode mode = msnp::capture_mode::force_pinned) { using clock = std::chrono::steady_clock; + const auto ctx = capture_context(agent, mode); + if(!ctx) return {}; + auto snap_start = clock::now(); - auto snapshot = msnp::snap(agent); + auto snapshot = msnp::snap(*ctx); auto snap_end = clock::now(); - if(!snapshot.ok) return {}; + if(!snapshot.complete()) return {}; kernel_launch::add(buffer, 1.0f, n_elems); sync_ok(); @@ -198,12 +218,16 @@ measure_snap_restore_once(hsa_agent_t agent, float* buffer, int n_elems) // sample, and it still catches a real regression: code that genuinely got slower slows every // iteration down, the best one included. snap_restore_timing_t -measure_snap_restore_best(hsa_agent_t agent, float* buffer, int n_elems, int iterations) +measure_snap_restore_best(hsa_agent_t agent, + float* buffer, + int n_elems, + int iterations, + msnp::capture_mode mode = msnp::capture_mode::force_pinned) { snap_restore_timing_t best{}; for(int i = 0; i < iterations; ++i) { - auto t = measure_snap_restore_once(agent, buffer, n_elems); + auto t = measure_snap_restore_once(agent, buffer, n_elems, mode); if(i == 0 || t.snap_seconds < best.snap_seconds) best.snap_seconds = t.snap_seconds; if(i == 0 || t.restore_seconds < best.restore_seconds) best.restore_seconds = t.restore_seconds; @@ -238,11 +262,12 @@ log_timing(const char* label, const snap_restore_timing_t& t) // setup dominates, so the computed rate reports fixed cost rather than link bandwidth. Those cases // still run, and still assert the wall-time ceiling below, which is what catches a blow-up. void -run_bandwidth_test(size_t ballast_bytes, - const char* label, - int warmup_iterations = 1, - int measure_iterations = 3, - bool check_bandwidth_floor = true) +run_bandwidth_test(size_t ballast_bytes, + const char* label, + int warmup_iterations = 1, + int measure_iterations = 3, + bool check_bandwidth_floor = true, + msnp::capture_mode mode = msnp::capture_mode::force_pinned) { const auto agent = gpu_agent(); ASSERT_NE(agent.handle, 0U); @@ -257,11 +282,11 @@ run_bandwidth_test(size_t ballast_bytes, // Warmup: prime caches and allocator state before timed iterations. for(int w = 0; w < warmup_iterations; ++w) { - const auto warmup = measure_snap_restore_once(agent, buffer, n_elems); + const auto warmup = measure_snap_restore_once(agent, buffer, n_elems, mode); ASSERT_GT(warmup.footprint_bytes, 0U) << "snap/restore warmup failed"; } - const auto timing = measure_snap_restore_best(agent, buffer, n_elems, measure_iterations); + const auto timing = measure_snap_restore_best(agent, buffer, n_elems, measure_iterations, mode); ASSERT_GT(timing.footprint_bytes, 0U) << "snap/restore measurement failed"; log_timing(label, timing); @@ -314,6 +339,17 @@ TEST(kernel_replay_snapshot, snap_bandwidth_meets_floor_with_128mb_ballast) run_bandwidth_test(128U * 1024U * 1024U, "ballast_128mb"); } +TEST(kernel_replay_snapshot, gpu_arena_bandwidth_with_128mb_ballast) +{ + if(!ensure_live_tracking()) GTEST_SKIP() << "could not activate rocprofiler / no HIP GPU"; + run_bandwidth_test(128U * 1024U * 1024U, + "gpu_arena_128mb", + 1, + 3, + /*check_bandwidth_floor=*/true, + msnp::capture_mode::force_gpu); +} + // Snap/restore cost should grow roughly with footprint (bandwidth-bound), not super-linearly. TEST(kernel_replay_snapshot, snap_cost_scales_sublinearly_with_footprint) { @@ -423,7 +459,7 @@ TEST(kernel_replay_snapshot, snap_restore_repeatable_from_one_snapshot) const auto snap_start = std::chrono::steady_clock::now(); auto snapshot = msnp::snap(agent); const auto snap_end = std::chrono::steady_clock::now(); - ASSERT_TRUE(snapshot.ok); + ASSERT_TRUE(snapshot.complete()); ASSERT_GT(snapshot_footprint_bytes(snapshot), 0U); // What this test exists to prove is that one snapshot image stays valid across repeated diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/snap_restore.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/snap_restore.cpp index 22e90c6b25fb..97672d1c3d6d 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/snap_restore.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/snap_restore.cpp @@ -42,8 +42,10 @@ #include #include +#include #include #include +#include #include using namespace rocprofiler; @@ -124,6 +126,33 @@ inventory_contains(void* p, hsa_agent_t agent) return inv.find(p) != inv.end(); } +std::optional +capture_context(hsa_agent_t agent, msnp::capture_mode mode, bool gpu_backend_available = true) +{ + auto cache = rocprofiler::agent::get_agent_cache(agent); + if(!cache) return std::nullopt; + return msnp::capture_context_t{ + agent, + cache->near_cpu(), + cache->gpu_pool(), + cache->cpu_pool(), + mode, + gpu_backend_available, + }; +} + +void +record_for_snapshot_test(void* ptr, size_t size, hsa_agent_t agent) +{ + mt::inventory().wlock([&](auto& map) { map[ptr] = mt::alloc_info_t{size, agent}; }); +} + +void +unrecord_for_snapshot_test(void* ptr) +{ + mt::inventory().wlock([&](auto& map) { map.erase(ptr); }); +} + std::vector read_device(const float* d, int n) { @@ -250,6 +279,277 @@ TEST(kernel_replay_snapshot, restore_reverts_device_memory) ASSERT_EQ(hipFree(buffer), hipSuccess); } +TEST(kernel_replay_snapshot, forced_gpu_arena_restores_device_memory) +{ + if(!ensure_live_tracking()) GTEST_SKIP() << "could not activate rocprofiler / no HIP GPU"; + + constexpr int n = 4096; + float* buffer = nullptr; + ASSERT_EQ(hipMalloc(&buffer, n * sizeof(float)), hipSuccess); + const auto allocation = mt::query_alloc(buffer); + ASSERT_TRUE(allocation.trackable); + const auto agent = allocation.agent; + ASSERT_NE(agent.handle, 0U); + const auto ctx = capture_context(agent, msnp::capture_mode::force_gpu); + ASSERT_TRUE(ctx.has_value()); + record_for_snapshot_test(buffer, n * sizeof(float), agent); + launch_fill(buffer, 7.0f, n); + + auto snapshot = msnp::snap(*ctx); + ASSERT_TRUE(snapshot.complete()); + ASSERT_EQ(snapshot.storage(), msnp::storage_kind::gpu_local); + + launch_fill(buffer, 19.0f, n); + bool batch_called = false; + ASSERT_TRUE(msnp::restore(snapshot, [&](const auto& gpu_regions) { + batch_called = true; + EXPECT_FALSE(gpu_regions.empty()); + for(const auto& copy : gpu_regions) + { + EXPECT_GT(copy.size, 0); + EXPECT_NE(copy.source, nullptr); + EXPECT_NE(copy.dest, nullptr); + const auto status = hsa_memory_copy(copy.dest, copy.source, copy.size); + if(status != HSA_STATUS_SUCCESS) return status; + } + return HSA_STATUS_SUCCESS; + })); + EXPECT_TRUE(batch_called); + const auto restored = read_device(buffer, n); + for(const auto value : restored) + EXPECT_FLOAT_EQ(value, 7.0f); + + unrecord_for_snapshot_test(buffer); + ASSERT_EQ(hipFree(buffer), hipSuccess); +} + +TEST(kernel_replay_snapshot, forced_gpu_arena_captures_through_batch_copy) +{ + if(!ensure_live_tracking()) GTEST_SKIP() << "could not activate rocprofiler / no HIP GPU"; + + constexpr int n = 4096; + float* buffer = nullptr; + ASSERT_EQ(hipMalloc(&buffer, n * sizeof(float)), hipSuccess); + const auto allocation = mt::query_alloc(buffer); + ASSERT_TRUE(allocation.trackable); + const auto agent = allocation.agent; + ASSERT_NE(agent.handle, 0U); + const auto ctx = capture_context(agent, msnp::capture_mode::force_gpu); + ASSERT_TRUE(ctx.has_value()); + record_for_snapshot_test(buffer, n * sizeof(float), agent); + launch_fill(buffer, 7.0f, n); + + bool batch_called = false; + auto snapshot = msnp::snap(*ctx, [&](const auto& gpu_regions) { + batch_called = true; + EXPECT_FALSE(gpu_regions.empty()); + for(const auto& copy : gpu_regions) + { + const auto status = hsa_memory_copy(copy.dest, copy.source, copy.size); + if(status != HSA_STATUS_SUCCESS) return status; + } + return HSA_STATUS_SUCCESS; + }); + ASSERT_TRUE(snapshot.complete()); + ASSERT_EQ(snapshot.storage(), msnp::storage_kind::gpu_local); + EXPECT_TRUE(batch_called); + + launch_fill(buffer, 19.0f, n); + ASSERT_TRUE(msnp::restore(snapshot)); + const auto restored = read_device(buffer, n); + for(const auto value : restored) + EXPECT_FLOAT_EQ(value, 7.0f); + + unrecord_for_snapshot_test(buffer); + ASSERT_EQ(hipFree(buffer), hipSuccess); +} + +TEST(kernel_replay_snapshot, unavailable_gpu_backend_falls_back_to_one_pinned_arena) +{ + if(!ensure_live_tracking()) GTEST_SKIP() << "could not activate rocprofiler / no HIP GPU"; + + constexpr int n = 4096; + float* buffer = nullptr; + ASSERT_EQ(hipMalloc(&buffer, n * sizeof(float)), hipSuccess); + const auto allocation = mt::query_alloc(buffer); + ASSERT_TRUE(allocation.trackable); + const auto agent = allocation.agent; + ASSERT_NE(agent.handle, 0U); + const auto ctx = + capture_context(agent, msnp::capture_mode::automatic, /*gpu_backend_available=*/false); + ASSERT_TRUE(ctx.has_value()); + record_for_snapshot_test(buffer, n * sizeof(float), agent); + launch_fill(buffer, 11.0f, n); + + auto snapshot = msnp::snap(*ctx); + ASSERT_TRUE(snapshot.complete()); + ASSERT_EQ(snapshot.storage(), msnp::storage_kind::pinned_host); + ASSERT_EQ(snapshot.segments.size(), 1); + + launch_fill(buffer, 23.0f, n); + bool continuation_called = false; + ASSERT_TRUE(msnp::restore(snapshot, [&](const auto& gpu_regions) { + EXPECT_TRUE(gpu_regions.empty()); + continuation_called = true; + return HSA_STATUS_SUCCESS; + })); + EXPECT_TRUE(continuation_called); + const auto restored = read_device(buffer, n); + for(const auto value : restored) + EXPECT_FLOAT_EQ(value, 11.0f); + + unrecord_for_snapshot_test(buffer); + ASSERT_EQ(hipFree(buffer), hipSuccess); +} + +TEST(kernel_replay_snapshot, failed_gpu_materialization_replans_all_regions_for_pinned_host) +{ + if(!ensure_live_tracking()) GTEST_SKIP() << "could not activate rocprofiler / no HIP GPU"; + + constexpr int n = 4096; + float* buffer = nullptr; + ASSERT_EQ(hipMalloc(&buffer, n * sizeof(float)), hipSuccess); + const auto allocation = mt::query_alloc(buffer); + ASSERT_TRUE(allocation.trackable); + const auto agent = allocation.agent; + record_for_snapshot_test(buffer, n * sizeof(float), agent); + launch_fill(buffer, 13.0f, n); + + auto ctx = capture_context(agent, msnp::capture_mode::automatic); + ASSERT_TRUE(ctx.has_value()); + ctx->gpu_pool.handle = 0; + + auto snapshot = + msnp::snap(*ctx, msnp::storage_budgets_t{64U * 1024U * 1024U, 512U * 1024U * 1024U}); + ASSERT_TRUE(snapshot.complete()); + EXPECT_EQ(snapshot.storage(), msnp::storage_kind::pinned_host); + ASSERT_EQ(snapshot.segments.size(), 1); + + launch_fill(buffer, 31.0f, n); + ASSERT_TRUE(msnp::restore(snapshot)); + const auto restored = read_device(buffer, n); + for(const auto value : restored) + EXPECT_FLOAT_EQ(value, 13.0f); + + unrecord_for_snapshot_test(buffer); + ASSERT_EQ(hipFree(buffer), hipSuccess); +} + +TEST(kernel_replay_snapshot, planned_regions_materialize_into_gpu_and_pinned_arenas) +{ + if(!ensure_live_tracking()) GTEST_SKIP() << "could not activate rocprofiler / no HIP GPU"; + + constexpr int n = 2 * 1024 * 1024; + constexpr size_t bytes = static_cast(n) * sizeof(float); + std::array buffers{}; + const std::array original = {3.0f, 5.0f, 7.0f}; + + hsa_agent_t agent{.handle = 0}; + for(size_t i = 0; i < buffers.size(); ++i) + { + ASSERT_EQ(hipMalloc(&buffers[i], bytes), hipSuccess); + const auto allocation = mt::query_alloc(buffers[i]); + ASSERT_TRUE(allocation.trackable); + if(agent.handle == 0) agent = allocation.agent; + ASSERT_EQ(allocation.agent.handle, agent.handle); + record_for_snapshot_test(buffers[i], bytes, agent); + launch_fill(buffers[i], original[i], n); + } + + auto ctx = capture_context(agent, msnp::capture_mode::automatic); + ASSERT_TRUE(ctx.has_value()); + + auto snapshot = + msnp::snap(*ctx, msnp::storage_budgets_t{10U * 1024U * 1024U, 512U * 1024U * 1024U}); + ASSERT_TRUE(snapshot.complete()); + ASSERT_EQ(snapshot.storage(), msnp::storage_kind::mixed); + ASSERT_EQ(snapshot.segments.size(), 2); + + size_t gpu_extents = 0; + size_t host_extents = 0; + std::array pinned_buffers{}; + for(const auto& extent : snapshot.extents) + { + const auto kind = snapshot.segments.at(extent.segment_index).backing.kind(); + if(kind == msnp::storage_kind::gpu_local) + ++gpu_extents; + else if(kind == msnp::storage_kind::pinned_host) + { + ++host_extents; + const auto* live = snapshot.regions.at(extent.region_index).live_address; + for(size_t i = 0; i < buffers.size(); ++i) + if(live == buffers[i]) pinned_buffers[i] = true; + } + } + EXPECT_GT(gpu_extents, 0); + EXPECT_GT(host_extents, 0); + EXPECT_TRUE( + std::any_of(pinned_buffers.begin(), pinned_buffers.end(), [](bool v) { return v; })); + + for(auto* buffer : buffers) + launch_fill(buffer, 29.0f, n); + + size_t callback_count = 0; + ASSERT_TRUE(msnp::restore(snapshot, [&](const auto& gpu_regions) { + ++callback_count; + EXPECT_EQ(gpu_regions.size(), gpu_extents); + // Pinned extents must be restored synchronously before the GPU batch/target continuation. + for(size_t i = 0; i < buffers.size(); ++i) + { + if(!pinned_buffers[i]) continue; + const auto sample = read_device(buffers[i], 64); + for(const auto value : sample) + EXPECT_FLOAT_EQ(value, original[i]); + } + for(const auto& copy : gpu_regions) + { + const auto status = hsa_memory_copy(copy.dest, copy.source, copy.size); + if(status != HSA_STATUS_SUCCESS) return status; + } + return HSA_STATUS_SUCCESS; + })); + EXPECT_EQ(callback_count, 1); + + for(size_t i = 0; i < buffers.size(); ++i) + { + const auto restored = read_device(buffers[i], n); + for(const auto value : restored) + ASSERT_FLOAT_EQ(value, original[i]); + unrecord_for_snapshot_test(buffers[i]); + ASSERT_EQ(hipFree(buffers[i]), hipSuccess); + } +} + +TEST(kernel_replay_snapshot, unavailable_gpu_and_pinned_arenas_decline_capture) +{ + if(!ensure_live_tracking()) GTEST_SKIP() << "could not activate rocprofiler / no HIP GPU"; + + constexpr int n = 4096; + float* buffer = nullptr; + ASSERT_EQ(hipMalloc(&buffer, n * sizeof(float)), hipSuccess); + const auto allocation = mt::query_alloc(buffer); + ASSERT_TRUE(allocation.trackable); + const auto agent = allocation.agent; + record_for_snapshot_test(buffer, n * sizeof(float), agent); + + const auto snapshot = msnp::snap(msnp::capture_context_t{ + agent, + hsa_agent_t{.handle = 0}, + hsa_amd_memory_pool_t{.handle = 0}, + hsa_amd_memory_pool_t{.handle = 0}, + msnp::capture_mode::automatic, + false, + }); + + EXPECT_FALSE(snapshot.complete()); + EXPECT_EQ(snapshot.status, msnp::capture_status::unavailable); + EXPECT_TRUE(snapshot.segments.empty()); + EXPECT_TRUE(snapshot.extents.empty()); + + unrecord_for_snapshot_test(buffer); + ASSERT_EQ(hipFree(buffer), hipSuccess); +} + // Restoring between passes must stop an in-place kernel from accumulating: N saxpy passes with a // restore each should net one application (y0 + a), not N. A no-op restore would leave y0 + N*a. // The saxpy kernel also bumps a __device__ module global once per pass, so this covers both tracked diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/snapshot_storage.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/snapshot_storage.cpp new file mode 100644 index 000000000000..2eeae17b690c --- /dev/null +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/kernel_replay/tests/snapshot_storage.cpp @@ -0,0 +1,397 @@ +// MIT License +// +// Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +#include "lib/rocprofiler-sdk/kernel_replay/snapshot-plan.hpp" +#include "lib/rocprofiler-sdk/kernel_replay/snapshot-storage.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +namespace plan = rocprofiler::kernel_replay::memory_snapshot::planning; +namespace store = rocprofiler::kernel_replay::memory_snapshot::storage; + +struct mock_state +{ + bool fail = false; + size_t allocations = 0; + size_t last_size = 0; + std::weak_ptr> last = {}; +}; + +// Stands in for an HSA pool, mmap, hipFile, ... Its handles release their bytes when destroyed. +struct mock_backing +{ + using handle_type = std::shared_ptr>; + + std::shared_ptr state = std::make_shared(); + + std::optional allocate(size_t size) const + { + ++state->allocations; + state->last_size = size; + if(state->fail) return std::nullopt; + auto bytes = std::make_shared>(size); + state->last = bytes; + return bytes; + } +}; + +using bump_t = store::bump_suballocator; +using whole_t = store::arena; +using chunked_t = store::arena; + +bump_t +bump(size_t capacity, std::shared_ptr state = std::make_shared()) +{ + return bump_t{capacity, 4, mock_backing{std::move(state)}}; +} + +struct range +{ + size_t offset = 0; + size_t size = 0; +}; + +// First fit over fixed free ranges, standing in for a reused arena whose free space is fragmented. +struct first_fit_suballocator +{ + using handle_type = mock_backing::handle_type; + + std::vector holes = {}; + size_t extent = 0; + mock_backing backing = {}; + std::vector remaining = holes; + std::optional handle = {}; + + size_t available() const + { + size_t largest = 0; + for(const auto& hole : remaining) + largest = std::max(largest, hole.size); + return largest; + } + + size_t reserve(size_t size) + { + auto hole = std::find_if(remaining.begin(), + remaining.end(), + [size](const range& candidate) { return candidate.size >= size; }); + const auto offset = hole->offset; + hole->offset += size; + hole->size -= size; + return offset; + } + + bool commit() + { + handle = backing.allocate(extent); + return handle.has_value(); + } + + void reset() + { + handle.reset(); + remaining = holes; + } + + size_t backing_size() const { return extent; } + std::optional release_backing() { return std::exchange(handle, std::nullopt); } +}; + +// Takes nothing, so every request moves on to the next tier. +struct decline +{ + size_t admit(const plan::request_t&, size_t) const { return 0; } +}; + +// Discovery would record which regions must stay device-local, such as kernargs. +struct is_device_only +{ + const std::vector* device_only = nullptr; + + bool operator()(const plan::request_t& request) const + { + return device_only->at(request.region_index); + } +}; + +// Routes each request to admission A when the predicate holds and to B otherwise: a segregator +// keyed on memory properties instead of size. +template +struct a_or_b +{ + Predicate predicate = {}; + A a = {}; + B b = {}; + + size_t admit(const plan::request_t& request, size_t available) const + { + return predicate(request) ? a.admit(request, available) : b.admit(request, available); + } +}; + +// Move-only, with whatever constructor state a real allocator (hipFile, mmap, ...) needs. +struct configured_backing +{ + using handle_type = size_t; + + std::unique_ptr multiplier; + std::string name; + + std::optional allocate(size_t size) const { return size * *multiplier; } +}; + +TEST(kernel_replay_snapshot_storage, whole_regions_spill_without_sticking) +{ + auto storage = store::fallback{ + whole_t{{}, bump(6)}, + whole_t{{}, bump(16)}, + store::null_tier{}, + }; + + const auto placements = plan::plan_and_commit(storage, {{0, 0, 10}, {1, 0, 3}, {2, 0, 2}}); + ASSERT_TRUE(placements); + ASSERT_EQ(placements->size(), 3); + EXPECT_EQ(placements->at(0).tier, 1); + EXPECT_EQ(placements->at(1).tier, 0); + EXPECT_EQ(placements->at(1).backing_offset, 0); + EXPECT_EQ(placements->at(2).tier, 0); + EXPECT_EQ(placements->at(2).backing_offset, 4); +} + +TEST(kernel_replay_snapshot_storage, prefix_chunking_splits_on_a_granule) +{ + auto storage = store::fallback{ + chunked_t{{4}, bump(10)}, + whole_t{{}, bump(16)}, + store::null_tier{}, + }; + + const auto placements = plan::plan_and_commit(storage, {{0, 0, 16}}); + ASSERT_TRUE(placements); + ASSERT_EQ(placements->size(), 2); + EXPECT_EQ(placements->at(0).size, 8); + EXPECT_EQ(placements->at(0).tier, 0); + EXPECT_EQ(placements->at(1).region_offset, 8); + EXPECT_EQ(placements->at(1).size, 8); + EXPECT_EQ(placements->at(1).tier, 1); +} + +TEST(kernel_replay_snapshot_storage, split_requests_interleave_tiers_by_priority) +{ + auto storage = store::fallback{ + whole_t{{}, bump(8)}, + whole_t{{}, bump(16)}, + store::null_tier{}, + }; + + // One region split before planning: hot [0, 4) and [12, 16) at priority 0, cold [4, 12) at 1. + const auto placements = + plan::plan_and_commit(storage, {{0, 0, 4, 0}, {0, 4, 8, 1}, {0, 12, 4, 0}}); + ASSERT_TRUE(placements); + ASSERT_EQ(placements->size(), 3); + EXPECT_EQ(placements->at(0).region_offset, 0); + EXPECT_EQ(placements->at(0).tier, 0); + EXPECT_EQ(placements->at(1).region_offset, 12); + EXPECT_EQ(placements->at(1).tier, 0); + EXPECT_EQ(placements->at(1).backing_offset, 4); + EXPECT_EQ(placements->at(2).region_offset, 4); + EXPECT_EQ(placements->at(2).tier, 1); +} + +TEST(kernel_replay_snapshot_storage, fragmented_tier_takes_a_request_in_pieces) +{ + using fragmented_t = store::arena; + auto state = std::make_shared(); + auto storage = store::fallback{ + fragmented_t{{1}, {{{0, 4}, {8, 8}}, 16, mock_backing{state}}}, + whole_t{{}, bump(16)}, + store::null_tier{}, + }; + + const auto placements = plan::plan_and_commit(storage, {{0, 0, 10}}); + ASSERT_TRUE(placements); + ASSERT_EQ(placements->size(), 2); + EXPECT_EQ(placements->at(0).size, 8); + EXPECT_EQ(placements->at(0).backing_offset, 8); + EXPECT_EQ(placements->at(1).size, 2); + EXPECT_EQ(placements->at(1).backing_offset, 0); + EXPECT_EQ(placements->at(1).tier, 0); + EXPECT_EQ(state->last_size, 16); +} + +TEST(kernel_replay_snapshot_storage, memory_properties_select_admission) +{ + using gpu_admission = a_or_b; + using host_admission = a_or_b; + using gpu_t = store::arena; + using host_t = store::arena; + + const auto device_only = std::vector{false, true}; + const auto make_storage = [&device_only](size_t gpu_capacity) { + return store::fallback{ + gpu_t{{{&device_only}, {}, {2}}, bump(gpu_capacity)}, + host_t{{{&device_only}, {}, {}}, bump(16)}, + store::null_tier{}, + }; + }; + + // The device-only region gets priority 0 so it claims the GPU before ordinary memory fills it. + auto storage = make_storage(8); + const auto placements = plan::plan_and_commit(storage, {{0, 0, 6, 1}, {1, 0, 4, 0}}); + ASSERT_TRUE(placements); + ASSERT_EQ(placements->size(), 3); + EXPECT_EQ(placements->at(0).region_index, 1); + EXPECT_EQ(placements->at(0).size, 4); + EXPECT_EQ(placements->at(0).tier, 0); + EXPECT_EQ(placements->at(1).region_index, 0); + EXPECT_EQ(placements->at(1).tier, 0); + EXPECT_EQ(placements->at(2).region_offset, 4); + EXPECT_EQ(placements->at(2).tier, 1); + + auto tight = make_storage(4); + EXPECT_FALSE(plan::plan_and_commit(tight, {{1, 0, 8}})); + const auto ordinary = plan::plan_and_commit(tight, {{0, 0, 8}}); + ASSERT_TRUE(ordinary); + ASSERT_EQ(ordinary->size(), 2); + EXPECT_EQ(ordinary->at(1).tier, 1); +} + +TEST(kernel_replay_snapshot_storage, failed_commit_disables_the_tier_and_replans) +{ + auto refusing = std::make_shared(); + refusing->fail = true; + auto storage = store::fallback{ + whole_t{{}, bump(16, refusing)}, + whole_t{{}, bump(16)}, + store::null_tier{}, + }; + + const auto placements = plan::plan_and_commit(storage, {{0, 0, 6}, {1, 0, 4}}); + ASSERT_TRUE(placements); + EXPECT_EQ(refusing->allocations, 1); + ASSERT_EQ(placements->size(), 2); + EXPECT_EQ(placements->at(0).tier, 1); + EXPECT_EQ(placements->at(1).tier, 1); +} + +TEST(kernel_replay_snapshot_storage, exhausted_tiers_fail_planning_without_allocating) +{ + auto gpu = std::make_shared(); + auto host = std::make_shared(); + auto storage = store::fallback{ + whole_t{{}, bump(4, gpu)}, + whole_t{{}, bump(4, host)}, + store::null_tier{}, + }; + + EXPECT_FALSE(plan::plan_and_commit(storage, {{0, 0, 5}})); + EXPECT_EQ(gpu->allocations, 0); + EXPECT_EQ(host->allocations, 0); +} + +TEST(kernel_replay_snapshot_storage, commit_failure_releases_committed_tiers) +{ + auto first = std::make_shared(); + auto second = std::make_shared(); + second->fail = true; + auto storage = store::fallback{ + whole_t{{}, bump(8, first)}, + whole_t{{}, bump(8, second)}, + store::null_tier{}, + }; + + ASSERT_TRUE(plan::plan(storage, {{0, 0, 8}, {1, 0, 4}})); + const auto failed = storage.commit(); + ASSERT_TRUE(failed); + EXPECT_EQ(*failed, 1); + EXPECT_EQ(first->allocations, 1); + EXPECT_TRUE(first->last.expired()); +} + +TEST(kernel_replay_snapshot_storage, released_backings_outlive_the_storage) +{ + auto first = std::make_shared(); + auto second = std::make_shared(); + auto handles = std::vector>{}; + { + auto storage = store::fallback{ + whole_t{{}, bump(10, first)}, + whole_t{{}, bump(10, second)}, + store::null_tier{}, + }; + ASSERT_TRUE(plan::plan_and_commit(storage, {{0, 0, 6}})); + storage.release_backings([&](size_t tier, size_t size, mock_backing::handle_type handle) { + EXPECT_EQ(handle->size(), size); + handles.emplace_back(tier, std::move(handle)); + }); + } + + EXPECT_EQ(second->allocations, 0); + ASSERT_EQ(handles.size(), 1); + EXPECT_EQ(handles.front().first, 0); + EXPECT_FALSE(first->last.expired()); +} + +TEST(kernel_replay_snapshot_storage, backing_allocators_can_be_move_only_with_arbitrary_state) +{ + using configured_t = + store::arena>; + auto storage = store::fallback{ + configured_t{{}, {16, 1, configured_backing{std::make_unique(3), "mock"}}}, + store::null_tier{}, + }; + + ASSERT_TRUE(plan::plan_and_commit(storage, {{0, 0, 8}})); + const auto& handle = std::get<0>(storage.tiers).suballocator.handle; + ASSERT_TRUE(handle.has_value()); + EXPECT_EQ(*handle, 24); +} + +TEST(kernel_replay_snapshot_storage, near_max_capacity_never_wraps) +{ + constexpr auto largest = std::numeric_limits::max(); + auto storage = store::fallback{ + whole_t{{}, bump_t{largest, 16, mock_backing{}}}, + whole_t{{}, bump(1024)}, + store::null_tier{}, + }; + + // The second request's aligned offset would pass SIZE_MAX, so the first tier declines it. + const auto placements = plan::plan(storage, {{0, 0, largest - 7}, {1, 0, 16}}); + ASSERT_TRUE(placements); + ASSERT_EQ(placements->size(), 2); + EXPECT_EQ(placements->at(0).tier, 0); + EXPECT_EQ(placements->at(1).tier, 1); +} +} // namespace diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/digest.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/digest.cpp index b7daf093b237..0839138edfed 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/digest.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/digest.cpp @@ -23,6 +23,11 @@ #include "lib/rocprofiler-sdk/range_replay/digest.hpp" #include +#include + +#if defined(__x86_64__) +# include +#endif namespace rocprofiler { @@ -30,6 +35,54 @@ namespace range_replay { namespace digest { +namespace +{ +// The digest is only ever compared against other digests produced by the same binary in the +// same process, and is never persisted, so the hash function may change freely between +// releases. The snapshot regions hashed here can be hundreds of megabytes, which makes +// throughput matter: the byte-at-a-time FNV-1a below is a dependent-multiply chain, while the +// SSE4.2 CRC32 instruction consumes eight bytes per instruction, and four independent lanes +// keep the multiplier pipeline fed. SSE4.2 is probed at runtime; the FNV-1a loop remains the +// fallback. +#if defined(__x86_64__) +__attribute__((target("sse4.2"))) uint64_t +crc32x4(const uint8_t* data, size_t size) +{ + constexpr uint64_t offset_basis = 14695981039346656037ULL; + uint64_t h0 = offset_basis, h1 = 1, h2 = 2, h3 = 3; + size_t i = 0; + for(; i + 32 <= size; i += 32) + { + uint64_t v0, v1, v2, v3; + std::memcpy(&v0, data + i, 8); + std::memcpy(&v1, data + i + 8, 8); + std::memcpy(&v2, data + i + 16, 8); + std::memcpy(&v3, data + i + 24, 8); + h0 = _mm_crc32_u64(h0, v0); + h1 = _mm_crc32_u64(h1, v1); + h2 = _mm_crc32_u64(h2, v2); + h3 = _mm_crc32_u64(h3, v3); + } + for(; i + 8 <= size; i += 8) + { + uint64_t v; + std::memcpy(&v, data + i, 8); + h0 = _mm_crc32_u64(h0, v); + } + for(; i < size; ++i) + h0 = _mm_crc32_u8(static_cast(h0), data[i]); + return h0 ^ h1 ^ h2 ^ h3; +} + +bool +have_sse42() +{ + static const bool v = __builtin_cpu_supports("sse4.2"); + return v; +} +#endif +} // namespace + uint64_t hash_bytes(const void* data, size_t size) { @@ -38,6 +91,10 @@ hash_bytes(const void* data, size_t size) if(data == nullptr) return offset_basis; +#if defined(__x86_64__) + if(have_sse42()) return crc32x4(static_cast(data), size); +#endif + const auto* bytes = static_cast(data); uint64_t hash = offset_basis; for(size_t i = 0; i < size; ++i) @@ -48,6 +105,21 @@ hash_bytes(const void* data, size_t size) return hash; } +region_digests_t +hash_regions(const std::vector& regions) +{ + auto ordered = regions; + std::sort(ordered.begin(), ordered.end(), [](const auto& lhs, const auto& rhs) { + return lhs.address < rhs.address; + }); + + auto out = region_digests_t{}; + out.reserve(ordered.size()); + for(const auto& region : ordered) + out.emplace_back(hash_bytes(region.data, region.size)); + return out; +} + size_t count_divergent(const region_digests_t& lhs, const region_digests_t& rhs) { diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/digest.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/digest.hpp index bcb7c2896794..81070380bcf3 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/digest.hpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/digest.hpp @@ -49,14 +49,27 @@ namespace range_replay { namespace digest { -// FNV-1a over `size` bytes at `data`. Chosen for being trivially auditable and dependency-free; -// this detects divergence, it does not defend against anything. +// Hashes `size` bytes at `data` with the SSE4.2 CRC32 instruction when available and with FNV-1a +// otherwise. Both are trivially auditable and dependency-free; this detects divergence, it does +// not defend against anything. uint64_t hash_bytes(const void* data, size_t size); // Per-region digests of one point in time, positionally aligned with the snapshot's block list. using region_digests_t = std::vector; +struct region_data_t +{ + const void* address = nullptr; + const void* data = nullptr; + size_t size = 0; +}; + +// Hash regions in ascending device-address order so snapshots with different inventory enumeration +// orders can be compared positionally. +region_digests_t +hash_regions(const std::vector& regions); + // Number of regions whose digest differs. Mismatched lengths count as fully divergent: the two // digests do not describe the same set of regions, so nothing can be concluded region by region. size_t diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/executor.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/executor.cpp index 6bff4bf47cf3..fe2448c91ce4 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/executor.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/executor.cpp @@ -26,6 +26,7 @@ #include "lib/common/scope_destructor.hpp" #include "lib/rocprofiler-sdk/context/context.hpp" #include "lib/rocprofiler-sdk/hsa/agent_cache.hpp" +#include "lib/rocprofiler-sdk/hsa/hsa.hpp" #include "lib/rocprofiler-sdk/hsa/queue.hpp" #include "lib/rocprofiler-sdk/hsa/replay_window.hpp" #include "lib/rocprofiler-sdk/hsa/rocprofiler_packet.hpp" @@ -44,6 +45,7 @@ #include #include #include +#include #include #include #include @@ -56,7 +58,7 @@ namespace range_replay { namespace { -using snapshot_t = kernel_replay::memory_snapshot::device_snapshot_t; +using snapshot_t = kernel_replay::memory_snapshot::snapshot_t; // Queue used by the ring writer below. The interceptor's writer signature carries no user data, so // the target queue is passed through thread-local state for the duration of one submission. @@ -196,10 +198,20 @@ elapsed_ms(phase_clock::duration elapsed) size_t snapshot_bytes(const snapshot_t& snapshot) { - auto total = size_t{0}; - for(const auto& block : snapshot.blocks) - total += block.host_copy.size(); - return total; + return snapshot.footprint_bytes(); +} + +kernel_replay::memory_snapshot::capture_context_t +make_snapshot_context(const hsa::Queue& queue, hsa_agent_t agent) +{ + return kernel_replay::memory_snapshot::capture_context_t{ + agent, + queue.get_agent().near_cpu(), + queue.get_agent().gpu_pool(), + queue.get_agent().cpu_pool(), + kernel_replay::memory_snapshot::capture_mode::automatic, + kernel_replay::blit::prepare(queue) == HSA_STATUS_SUCCESS, + }; } } // namespace @@ -233,24 +245,62 @@ build_pass_packets(const std::vector& dispatches) return packets; } -digest::region_digests_t +std::optional snapshot_digests(const snapshot_t& snapshot) { - auto keyed = std::vector>{}; - keyed.reserve(snapshot.blocks.size()); - for(const auto& block : snapshot.blocks) - keyed.emplace_back(block.gpu_addr, - digest::hash_bytes(block.host_copy.data(), block.host_copy.size())); - - std::sort(keyed.begin(), keyed.end(), [](const auto& lhs, const auto& rhs) { - return lhs.first < rhs.first; - }); - - auto out = digest::region_digests_t{}; - out.reserve(keyed.size()); - for(const auto& [addr, hash] : keyed) - out.emplace_back(hash); - return out; + if(!snapshot.complete()) return std::nullopt; + + // Thread-local, grow-only scratch. std::vector::resize only value-initializes the elements it + // adds, so once the buffers reach the region sizes, later digests pay no zero-fill and no + // first-touch page faults. Safe to reuse without clearing because every region is covered by + // exactly one extent, so every byte is overwritten by the readback below before it is hashed. + static thread_local auto region_bytes = std::vector>{}; + try + { + region_bytes.resize(snapshot.regions.size()); + for(size_t i = 0; i < snapshot.regions.size(); ++i) + region_bytes[i].resize(snapshot.regions[i].size); + } catch(const std::bad_alloc&) + { + return std::nullopt; + } + + const auto* core = hsa::get_core_table(); + for(const auto& extent : snapshot.extents) + { + const auto& segment = snapshot.segments.at(extent.segment_index); + const auto* source = + static_cast(segment.backing.block().data) + extent.segment_offset; + auto* destination = region_bytes.at(extent.region_index).data() + extent.region_offset; + + if(segment.backing.kind() == kernel_replay::memory_snapshot::storage_kind::gpu_local) + { + if(core == nullptr || core->hsa_memory_copy_fn == nullptr || + core->hsa_memory_copy_fn(destination, source, extent.size) != HSA_STATUS_SUCCESS) + { + ROCP_WARNING << "range replay: divergence snapshot readback failed"; + return std::nullopt; + } + } + else + { + std::memcpy(destination, source, extent.size); + } + } + + auto regions = std::vector{}; + regions.reserve(snapshot.regions.size()); + for(size_t i = 0; i < snapshot.regions.size(); ++i) + { + const auto has_extent = + std::any_of(snapshot.extents.begin(), snapshot.extents.end(), [i](const auto& extent) { + return extent.region_index == i; + }); + if(has_extent) + regions.emplace_back(digest::region_data_t{ + snapshot.regions[i].live_address, region_bytes[i].data(), region_bytes[i].size()}); + } + return digest::hash_regions(regions); } bool @@ -310,11 +360,12 @@ ensure_entry_snapshot(range_context_t& ctx, hsa::replay_drain_agent_or_fatal(ctx.hsa_agent); const auto drained_at = phase_clock::now(); - ctx.snapshot = kernel_replay::memory_snapshot::snap(ctx.hsa_agent); + ctx.snapshot = + kernel_replay::memory_snapshot::snap(make_snapshot_context(queue, ctx.hsa_agent)); ctx.snapshot_taken = true; const auto snapped_at = phase_clock::now(); - if(!ctx.snapshot.ok) + if(!ctx.snapshot.complete()) { ROCP_WARNING << "range replay: snapshot capture failed (memory pressure or copy error); " "this range will not be replayed"; @@ -331,7 +382,7 @@ ensure_entry_snapshot(range_context_t& ctx, elapsed_ms(drained_at - entry_start), elapsed_ms(snapped_at - drained_at), snapshot_bytes(ctx.snapshot), - ctx.snapshot.blocks.size()); + ctx.snapshot.regions.size()); } rocprofiler_range_replay_status_t @@ -343,7 +394,7 @@ execute_range(range_context_t& ctx, uint64_t& divergence_count) if(!ctx.plan.replay_requested) return ROCPROFILER_RANGE_REPLAY_STATUS_NO_PASS_COUNT; if(ctx.record.dispatch_count() == 0 || ctx.queue == nullptr) return ROCPROFILER_RANGE_REPLAY_STATUS_NO_DISPATCH; - if(!ctx.snapshot_taken || !ctx.snapshot.ok) + if(!ctx.snapshot_taken || !ctx.snapshot.complete()) return ROCPROFILER_RANGE_REPLAY_STATUS_SNAPSHOT_FAILED; const auto& queue = *ctx.queue; @@ -368,10 +419,11 @@ execute_range(range_context_t& ctx, uint64_t& divergence_count) return ROCPROFILER_RANGE_REPLAY_STATUS_ALLOCATION_CHANGED_IN_RANGE; // The state the application must resume with, captured before the first pass overwrites it. - const auto snap_start = phase_clock::now(); - const auto exit_snapshot = kernel_replay::memory_snapshot::snap(ctx.hsa_agent); - const auto snapped_at = phase_clock::now(); - if(!exit_snapshot.ok) return ROCPROFILER_RANGE_REPLAY_STATUS_SNAPSHOT_FAILED; + const auto snap_start = phase_clock::now(); + const auto exit_snapshot = + kernel_replay::memory_snapshot::snap(make_snapshot_context(queue, ctx.hsa_agent)); + const auto snapped_at = phase_clock::now(); + if(!exit_snapshot.complete()) return ROCPROFILER_RANGE_REPLAY_STATUS_SNAPSHOT_FAILED; auto staging = kernarg_staging{}; if(!staging.reserve(queue, ctx.agent_id, ctx.record.dispatches())) @@ -398,27 +450,17 @@ execute_range(range_context_t& ctx, uint64_t& divergence_count) const auto agent_id = ctx.agent_id; const auto dispatch_count = static_cast(ctx.record.dispatch_count()); - auto passes_time = phase_clock::duration{}; - auto restores_time = phase_clock::duration{}; - uint64_t passes_run = 0; - uint64_t restores_run = 0; + auto passes_time = phase_clock::duration{}; + auto restores_time = phase_clock::duration{}; + uint64_t passes_run = 0; + uint64_t restores_run = 0; + auto pending_restore = std::vector{}; // Pass 0 was the application's own execution, so the replayed passes are numbered from 1. for(uint64_t pass = 1;; ++pass) { const bool is_final = !ctx.plan.indefinite && (pass == ctx.plan.total_passes - 1); - // Rewind to the state the range started from. A failed restore leaves device memory - // partially written; continuing would submit a pass over corrupted data and hand that - // corruption back to the application. - const auto restore_start = phase_clock::now(); - ROCP_FATAL_IF(!kernel_replay::memory_snapshot::restore(ctx.snapshot)) << fmt::format( - "range replay: restore of the range-entry snapshot failed (partial host->device copy); " - "aborting rather than continuing with corrupted device memory"); - const auto pass_start = phase_clock::now(); - restores_time += pass_start - restore_start; - ++restores_run; - auto pass_state = pass_context_state_t{}; execute_pass_phase_enter(ctx.plan, pass, @@ -429,13 +471,34 @@ execute_range(range_context_t& ctx, uint64_t& divergence_count) ctx.ancestor_corr_id, pass_state); - staging.fill(ctx.record.dispatches(), packets); - queue.invoke_write_interceptor(packets.data(), packets.size(), ring_writer); + const auto pass_start = phase_clock::now(); + const auto batch_copy = [&](const auto& regions) { + if(kernel_replay::blit::copy(queue, ring_writer, regions, pending_restore) != + HSA_STATUS_SUCCESS) + return HSA_STATUS_ERROR; + + staging.fill(ctx.record.dispatches(), packets); + queue.invoke_write_interceptor(packets.data(), packets.size(), ring_writer); + + // The target dispatch is submitted behind the restore blit. Draining it also proves + // that the blit completed, so the restore packet resources can be retired while the + // inventory read lock remains held by memory_snapshot::restore(). + hsa::replay_drain_or_fatal(queue); + for(auto& restore : pending_restore) + ROCP_FATAL_IF(restore.retire() != HSA_STATUS_SUCCESS) + << "range replay: asynchronous restore blit failed"; + pending_restore.clear(); + return HSA_STATUS_SUCCESS; + }; - // Drain this pass's async completion handlers (a separate HSA thread reads counters, emits - // records, and releases signals) before the continue-decision, the next restore, and the - // next submit. - hsa::replay_drain_or_fatal(queue); + const auto restore_start = phase_clock::now(); + ROCP_FATAL_IF(!kernel_replay::memory_snapshot::restore(ctx.snapshot, batch_copy)) + << fmt::format( + "range replay: restore of the range-entry snapshot failed (partial copy); " + "aborting rather than continuing with corrupted device memory"); + const auto pass_submitted_at = phase_clock::now(); + restores_time += pass_submitted_at - restore_start; + ++restores_run; execute_pass_phase_exit(ctx.plan, pass, agent_id, dispatch_count, pass_state); passes_time += phase_clock::now() - pass_start; @@ -447,18 +510,50 @@ execute_range(range_context_t& ctx, uint64_t& divergence_count) const auto verify_start = phase_clock::now(); if(divergence_check_enabled()) { - const auto post_snapshot = kernel_replay::memory_snapshot::snap(ctx.hsa_agent); - if(post_snapshot.ok) - divergence_count = digest::count_divergent(snapshot_digests(exit_snapshot), - snapshot_digests(post_snapshot)); + const auto post_snapshot = + kernel_replay::memory_snapshot::snap(make_snapshot_context(queue, ctx.hsa_agent)); + if(!post_snapshot.complete()) + { + ROCP_WARNING << "range replay: divergence snapshot capture failed"; + divergence_count = exit_snapshot.regions.size(); + } + else + { + const auto exit_digests = snapshot_digests(exit_snapshot); + const auto post_digests = snapshot_digests(post_snapshot); + if(!exit_digests || !post_digests) + { + ROCP_WARNING << "range replay: divergence snapshot readback failed"; + divergence_count = + std::max(exit_digests ? exit_digests->size() : exit_snapshot.regions.size(), + post_digests ? post_digests->size() : post_snapshot.regions.size()); + } + else + { + divergence_count = digest::count_divergent(*exit_digests, *post_digests); + } + } } const auto verified_at = phase_clock::now(); // Hand the application back the state its own execution produced, whether or not the replay - // reproduced it. - ROCP_FATAL_IF(!kernel_replay::memory_snapshot::restore(exit_snapshot)) << fmt::format( - "range replay: restore of the range-exit snapshot failed (partial host->device copy); " - "aborting rather than returning corrupted device memory to the application"); + // reproduced it. GPU-backed regions are restored through the blit queue and waited on before + // end() returns. + auto pending_exit_restore = std::vector{}; + const auto exit_batch_copy = [&](const auto& regions) { + if(kernel_replay::blit::copy(queue, ring_writer, regions, pending_exit_restore) != + HSA_STATUS_SUCCESS) + return HSA_STATUS_ERROR; + for(auto& restore : pending_exit_restore) + if(restore.wait() != HSA_STATUS_SUCCESS) return HSA_STATUS_ERROR; + for(auto& restore : pending_exit_restore) + if(restore.retire() != HSA_STATUS_SUCCESS) return HSA_STATUS_ERROR; + pending_exit_restore.clear(); + return HSA_STATUS_SUCCESS; + }; + ROCP_FATAL_IF(!kernel_replay::memory_snapshot::restore(exit_snapshot, exit_batch_copy)) + << fmt::format("range replay: restore of the range-exit snapshot failed (partial copy); " + "aborting rather than returning corrupted device memory to the application"); const auto window_end = phase_clock::now(); ROCP_INFO << fmt::format( @@ -470,7 +565,7 @@ execute_range(range_context_t& ctx, uint64_t& divergence_count) elapsed_ms(drained_at - window_start), elapsed_ms(snapped_at - snap_start), snapshot_bytes(exit_snapshot), - exit_snapshot.blocks.size(), + exit_snapshot.regions.size(), elapsed_ms(staged_at - snapped_at), staging.capacity(), staging.origin(), diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/executor.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/executor.hpp index f085b5ddc49f..404d340503af 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/executor.hpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/executor.hpp @@ -83,10 +83,10 @@ plan_kernarg_layout(const std::vector& dispatches); std::vector build_pass_packets(const std::vector& dispatches); -// Per-region digests of a snapshot's host copies, ordered by device address so two snapshots of the -// same regions compare positionally (the inventory itself is unordered). -digest::region_digests_t -snapshot_digests(const kernel_replay::memory_snapshot::device_snapshot_t& snapshot); +// Per-region digests of a snapshot's storage, reading GPU-backed regions back to host memory as +// needed. The optional is empty when a readback or allocation fails. +std::optional +snapshot_digests(const kernel_replay::memory_snapshot::snapshot_t& snapshot); // Capture the range-entry snapshot, if this range is bound and does not have one yet. Called from // the queue path while recording the range's first submission, before that submission reaches the diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/range_state.hpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/range_state.hpp index 06335b9b2f00..eb1bfcd74c7b 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/range_state.hpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/range_state.hpp @@ -159,8 +159,8 @@ struct range_context_t hsa_agent_t hsa_agent = {.handle = 0}; rocprofiler_agent_id_t agent_id = {.handle = 0}; - kernel_replay::memory_snapshot::device_snapshot_t snapshot = {}; - bool snapshot_taken = false; + kernel_replay::memory_snapshot::snapshot_t snapshot = {}; + bool snapshot_taken = false; // Base pointers of the tracked allocations the snapshot covers, sorted. Re-checked at close: a // range that allocated or freed device memory while open has regions the snapshot cannot // restore, so it is declined rather than replayed against partially-restored inputs. diff --git a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/tests/executor.cpp b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/tests/executor.cpp index 2064bacb9f08..257c633ee324 100644 --- a/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/tests/executor.cpp +++ b/projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/range_replay/tests/executor.cpp @@ -48,16 +48,14 @@ #include #include -#include #include -#include #include namespace rr = rocprofiler::range_replay; namespace { -using snapshot_t = rocprofiler::kernel_replay::memory_snapshot::device_snapshot_t; +using snapshot_t = rocprofiler::kernel_replay::memory_snapshot::snapshot_t; rr::recorded_dispatch_t make_dispatch(size_t kernarg_size, uint64_t kernel_object = 1) @@ -82,46 +80,10 @@ make_replayable_context() ctx.plan.replay_requested = true; ctx.plan.total_passes = 3; ctx.snapshot_taken = true; - ctx.snapshot.ok = true; + ctx.snapshot.status = rocprofiler::kernel_replay::memory_snapshot::capture_status::complete; return ctx; } -// mem_block_t::host_copy is either a std::vector or a pooled staging buffer that has to be -// acquired for a size before it is written. Fill whichever the snapshot uses. -template -void -fill_host_copy(HostCopy& host_copy, const std::string& contents) -{ - if constexpr(std::is_same_v>) - { - host_copy.assign(contents.begin(), contents.end()); - } - else - { - ASSERT_TRUE(host_copy.acquire(hsa_agent_t{.handle = 1}, contents.size())); - std::memcpy(host_copy.data(), contents.data(), contents.size()); - } -} - -snapshot_t -make_snapshot(const std::vector>& regions) -{ - auto snapshot = snapshot_t{}; - for(const auto& [addr, contents] : regions) - { - auto block = rocprofiler::kernel_replay::memory_snapshot::mem_block_t{}; - block.gpu_addr = addr; - fill_host_copy(block.host_copy, contents); - snapshot.blocks.emplace_back(std::move(block)); - } - return snapshot; -} - -// Arbitrary distinct device addresses. Never dereferenced: the digest reads host_copy and only -// sorts on gpu_addr. -void* const ADDR_LOW = reinterpret_cast(uintptr_t{0x1000}); -void* const ADDR_MID = reinterpret_cast(uintptr_t{0x2000}); -void* const ADDR_HIGH = reinterpret_cast(uintptr_t{0x3000}); } // namespace // --------------------------------------------------------------------------- @@ -355,44 +317,41 @@ TEST(range_replay_executor, building_a_pass_does_not_consume_the_recording) } } -// --------------------------------------------------------------------------- -// Snapshot digests -// --------------------------------------------------------------------------- - -// The snapshot inventory is unordered, so two snapshots of the same regions can enumerate them in -// different orders. Digests are ordered by device address so the comparison is positional; without -// that, a reordered inventory would report every region as divergent. -TEST(range_replay_executor, digests_are_ordered_by_device_address) +TEST(range_replay_executor, an_empty_snapshot_digests_to_nothing) { - const auto forward = - make_snapshot({{ADDR_LOW, "alpha"}, {ADDR_MID, "beta"}, {ADDR_HIGH, "gamma"}}); - const auto shuffled = - make_snapshot({{ADDR_HIGH, "gamma"}, {ADDR_LOW, "alpha"}, {ADDR_MID, "beta"}}); - - const auto forward_digests = rr::snapshot_digests(forward); - const auto shuffled_digests = rr::snapshot_digests(shuffled); - - ASSERT_EQ(forward_digests.size(), 3u); - EXPECT_EQ(forward_digests, shuffled_digests); - EXPECT_EQ(rocprofiler::range_replay::digest::count_divergent(forward_digests, shuffled_digests), - 0u); + auto snapshot = snapshot_t{}; + snapshot.status = rocprofiler::kernel_replay::memory_snapshot::capture_status::complete; + const auto digests = rr::snapshot_digests(snapshot); + ASSERT_TRUE(digests); + EXPECT_TRUE(digests->empty()); } -// A region whose contents changed between the application's execution and the last replayed pass is -// reported, and only that region: the count is what CLOSE hands the tool as divergence_count. -TEST(range_replay_executor, only_changed_regions_are_counted_as_divergent) +TEST(range_replay_executor, digest_regions_are_address_ordered) { - const auto before = - make_snapshot({{ADDR_LOW, "alpha"}, {ADDR_MID, "beta"}, {ADDR_HIGH, "gamma"}}); - const auto after = - make_snapshot({{ADDR_HIGH, "gamma"}, {ADDR_MID, "BETA"}, {ADDR_LOW, "alpha"}}); - - EXPECT_EQ(rocprofiler::range_replay::digest::count_divergent(rr::snapshot_digests(before), - rr::snapshot_digests(after)), - 1u); + const std::string alpha = "alpha"; + const std::string beta = "beta"; + const auto forward = std::vector{ + {reinterpret_cast(0x1000), alpha.data(), alpha.size()}, + {reinterpret_cast(0x2000), beta.data(), beta.size()}}; + const auto reversed = std::vector{ + {reinterpret_cast(0x2000), beta.data(), beta.size()}, + {reinterpret_cast(0x1000), alpha.data(), alpha.size()}}; + + EXPECT_EQ(rocprofiler::range_replay::digest::hash_regions(forward), + rocprofiler::range_replay::digest::hash_regions(reversed)); } -TEST(range_replay_executor, an_empty_snapshot_digests_to_nothing) +TEST(range_replay_executor, digest_regions_report_changed_content) { - EXPECT_TRUE(rr::snapshot_digests(snapshot_t{}).empty()); + const std::string before = "before"; + const std::string after = "after"; + const auto lhs = std::vector{ + {reinterpret_cast(0x1000), before.data(), before.size()}}; + const auto rhs = std::vector{ + {reinterpret_cast(0x1000), after.data(), after.size()}}; + + EXPECT_EQ(rocprofiler::range_replay::digest::count_divergent( + rocprofiler::range_replay::digest::hash_regions(lhs), + rocprofiler::range_replay::digest::hash_regions(rhs)), + 1u); }