diff --git a/docs/hdf5-dumps.md b/docs/hdf5-dumps.md index 8fd3714..4e445c8 100644 --- a/docs/hdf5-dumps.md +++ b/docs/hdf5-dumps.md @@ -23,17 +23,18 @@ kernel_label # Kokkos label of the matched kernel kernel_id # tool-local id assigned to this kernel invocation kernel_invocation # allow dumping a specific kernel invocation active_allocations # tracked user allocations alive at dump time -active_bytes # total size, in bytes, of those active allocations +active_bytes # total bounded user-data bytes for those active allocations ``` The `/views` group contains one subgroup per active allocation. Each allocation group stores: ```text -label # Kokkos allocation label -space # Kokkos memory space name reported by the profiling hook -ptr # allocation pointer value, stored as text -p_data # user data pointer after the Kokkos allocation header -size # allocation size in bytes +label # Kokkos allocation label +space # Kokkos memory space name reported by the profiling hook +ptr # allocation pointer value, stored as text +p_data # user data pointer after the Kokkos allocation header +size # bounded user data size in bytes, 0 when it cannot be bounded +reported_size # size reported by the Kokkos profiling allocation hook bytes_dumped ``` diff --git a/src/kkf/capture/CMakeLists.txt b/src/kkf/capture/CMakeLists.txt index 99b711d..39b5eeb 100644 --- a/src/kkf/capture/CMakeLists.txt +++ b/src/kkf/capture/CMakeLists.txt @@ -35,7 +35,7 @@ target_compile_definitions( if(Kokkos_ENABLE_CUDA) find_package(CUDAToolkit REQUIRED) target_compile_definitions(kkf_kokkos_hooks PRIVATE KKF_ENABLE_CUDA_DUMP) - target_link_libraries(kkf_kokkos_hooks PRIVATE CUDA::cudart) + target_link_libraries(kkf_kokkos_hooks PRIVATE CUDA::cudart CUDA::cuda_driver) endif() if(Kokkos_ENABLE_HIP) target_compile_definitions(kkf_kokkos_hooks PRIVATE KKF_ENABLE_HIP_DUMP) diff --git a/src/kkf/capture/allocation_tracker.cpp b/src/kkf/capture/allocation_tracker.cpp index 1ad3f95..f362c04 100644 --- a/src/kkf/capture/allocation_tracker.cpp +++ b/src/kkf/capture/allocation_tracker.cpp @@ -6,12 +6,15 @@ namespace kkf { void AllocationTracker::record_allocation(std::string label, std::string space, const void* ptr, const void* p_data, - const std::uint64_t size) { + const std::uint64_t size, + const std::uint64_t reported_size, + const bool data_size_known) { if (ptr == nullptr) { return; } - AllocationRecord record{std::move(label), space, p_data, size}; + AllocationRecord record{std::move(label), space, p_data, size, + reported_size, data_size_known}; std::lock_guard lock(mutex_); active_allocations_[space].insert_or_assign(ptr, std::move(record)); } diff --git a/src/kkf/capture/allocation_tracker.hpp b/src/kkf/capture/allocation_tracker.hpp index 3f077c3..b08a0f3 100644 --- a/src/kkf/capture/allocation_tracker.hpp +++ b/src/kkf/capture/allocation_tracker.hpp @@ -14,6 +14,8 @@ struct AllocationRecord { std::string space; const void* p_data; std::uint64_t size; + std::uint64_t reported_size; + bool data_size_known; }; struct ActiveAllocation { @@ -29,7 +31,8 @@ struct AllocationSnapshot { class AllocationTracker { public: void record_allocation(std::string label, std::string space, const void* ptr, - const void* p_data, std::uint64_t size); + const void* p_data, std::uint64_t size, + std::uint64_t reported_size, bool data_size_known); void record_deallocation(std::string space, const void* ptr); AllocationSnapshot snapshot() const; diff --git a/src/kkf/capture/kokkos_hooks.cpp b/src/kkf/capture/kokkos_hooks.cpp index 1d02364..46e6329 100644 --- a/src/kkf/capture/kokkos_hooks.cpp +++ b/src/kkf/capture/kokkos_hooks.cpp @@ -30,6 +30,19 @@ #include #include +#if defined(__linux__) +#include +#endif + +#if defined(KKF_ENABLE_CUDA_DUMP) +#include +#include +#endif + +#if defined(KKF_ENABLE_HIP_DUMP) +#include +#endif + #define KOKKOS_HOOKS_EXPORT __attribute__((visibility("default"))) #if !defined(KKF_KOKKOS_ALLOCATION_HEADER_SIZE) @@ -83,6 +96,20 @@ bool is_internal_label(std::string_view label) { label.starts_with("KOKKOS_") || label.starts_with("kokkos."); } +#if defined(KKF_ENABLE_CUDA_DUMP) +bool is_cuda_range_query_space(const std::string& space) { + return space == "Cuda" || space == "CudaUVM" || space == "CudaHostPinned"; +} +#endif + +#if defined(KKF_ENABLE_HIP_DUMP) +bool is_hip_range_query_space(const std::string& space) { + return space == "HIP" || space == "HIPManaged" || space == "HIPHostPinned"; +} +#endif + +bool is_host_space(const std::string& space) { return space == "Host"; } + std::string bounded_string(const char* value, std::size_t max_size) { if (value == nullptr) { return ""; @@ -159,6 +186,76 @@ const void* allocation_data_pointer(const void* ptr) { KKF_KOKKOS_ALLOCATION_HEADER_SIZE; } +std::size_t host_allocation_size(const void* ptr) { +#if defined(__linux__) + return malloc_usable_size(const_cast(ptr)); +#else + (void)ptr; + return 0; +#endif +} + +std::optional validated_data_size( + const void* allocation_base, const std::size_t allocation_size, + const void* data_ptr, const std::uint64_t reported_size) { + const auto base = reinterpret_cast(allocation_base); + const auto data = reinterpret_cast(data_ptr); + if (data < base) { + return std::nullopt; + } + + const auto offset = data - base; + if (offset > allocation_size) { + return std::nullopt; + } + + if (offset == allocation_size) { + return 0; + } + + const auto available = static_cast(allocation_size - offset); + return available >= reported_size ? reported_size : 0; +} + +std::optional allocation_data_size( + const std::string& space, const void* ptr, const void* data_ptr, + const std::uint64_t reported_size) { +#if defined(KKF_ENABLE_CUDA_DUMP) + if (is_cuda_range_query_space(space)) { + CUdeviceptr allocation_base = 0; + std::size_t allocation_size = 0; + const CUresult error = cuMemGetAddressRange( + &allocation_base, &allocation_size, reinterpret_cast(ptr)); + if (error == CUDA_SUCCESS) { + return validated_data_size(reinterpret_cast(allocation_base), + allocation_size, data_ptr, reported_size); + } + } +#endif + +#if defined(KKF_ENABLE_HIP_DUMP) + if (is_hip_range_query_space(space)) { + void* allocation_base = nullptr; + std::size_t allocation_size = 0; + const hipError_t error = + hipMemGetAddressRange(&allocation_base, &allocation_size, ptr); + if (error == hipSuccess) { + return validated_data_size(allocation_base, allocation_size, data_ptr, + reported_size); + } + } +#endif + + if (is_host_space(space)) { + const std::size_t allocation_size = host_allocation_size(ptr); + if (allocation_size != 0) { + return validated_data_size(ptr, allocation_size, data_ptr, reported_size); + } + } + + return std::nullopt; +} + void begin_kernel(const char* label, const std::uint32_t device_id, std::uint64_t* kernel_id) { const std::uint64_t id = next_kernel_id.fetch_add(1); @@ -293,11 +390,16 @@ KOKKOS_HOOKS_EXPORT void kokkosp_allocate_data( const void* ptr, const std::uint64_t size) { const std::string allocation_label = label_or_unknown(label); const std::string allocation_space = space_name(space); - const void* p_data = ptr != nullptr ? allocation_data_pointer(ptr) : nullptr; if (should_track_allocation(label, ptr)) { + const void* p_data = allocation_data_pointer(ptr); + const std::optional data_size = + allocation_data_size(allocation_space, ptr, p_data, size); + const std::uint64_t tracked_size = data_size.value_or(0); + allocation_tracker.record_allocation(allocation_label, allocation_space, - ptr, p_data, size); + ptr, p_data, tracked_size, size, + data_size.has_value()); } } diff --git a/src/kkf/capture/memory_copy.cpp b/src/kkf/capture/memory_copy.cpp index d13e4e3..a1d52b0 100644 --- a/src/kkf/capture/memory_copy.cpp +++ b/src/kkf/capture/memory_copy.cpp @@ -52,6 +52,9 @@ std::string copy_allocation_bytes(const ActiveAllocation& allocation, if (!host_accessible && !cuda_device && !hip_device) { return "memory space is not supported for raw byte dumps"; } + if (!allocation.record.data_size_known) { + return "allocation data size could not be bounded safely"; + } #if !defined(KKF_ENABLE_CUDA_DUMP) if (cuda_device) { return "CUDA dump support was not enabled at build time"; diff --git a/src/kkf/capture/view_dump.cpp b/src/kkf/capture/view_dump.cpp index beac963..cef3d4e 100644 --- a/src/kkf/capture/view_dump.cpp +++ b/src/kkf/capture/view_dump.cpp @@ -132,6 +132,8 @@ void write_allocation_group(hid_t views_group, write_string_attribute(group.get(), "p_data", pointer_to_string(allocation.record.p_data)); write_uint64_attribute(group.get(), "size", allocation.record.size); + write_uint64_attribute(group.get(), "reported_size", + allocation.record.reported_size); std::vector bytes; const std::string skip_reason = copy_allocation_bytes(allocation, bytes); diff --git a/src/kkf/replay/kernel_replayer.cpp b/src/kkf/replay/kernel_replayer.cpp index bea19e4..b197586 100644 --- a/src/kkf/replay/kernel_replayer.cpp +++ b/src/kkf/replay/kernel_replayer.cpp @@ -233,6 +233,9 @@ herr_t get_hdf5_dataset_alloc_info(hid_t group, const char* name, std::size_t buffer_size = std::reduce(dims.begin(), dims.end(), 1, std::multiplies<>{}); + if (buffer_size == 0) { + return 0; + } auto allocations = reinterpret_cast>*>( allocation_sets); @@ -278,6 +281,11 @@ herr_t allocate_hdf5_dataset(hid_t group, const char* name, const H5L_info_t*, std::size_t buffer_size = std::reduce(dims.begin(), dims.end(), 1, std::multiplies<>{}); + if (buffer_size == 0) { + (*reinterpret_cast(allocate_fun))(label, space, + nullptr, nullptr, 0); + return 0; + } MemorySpaceType memory_space = memory_space_type_from_string(space); auto free_buffer = [memory_space](char* ptr) { @@ -373,6 +381,17 @@ ScopeGuard::ScopeGuard(int& argc, char* argv[]) { char* data, std::size_t size) { impl::MemorySpaceType space = impl::memory_space_type_from_string(memory_space); + if (size == 0) { + if (space == impl::MemorySpaceType::HOST) { + host_allocations[label] = nullptr; + } else { +#if defined(KERNEL_REPLAYER_HAS_DEVICE_SPACE) + device_allocations[label] = nullptr; +#endif + } + return; + } + impl::copy_data(space, address, data, size); if (space == impl::MemorySpaceType::HOST) { host_allocations[label] = address; @@ -454,17 +473,21 @@ void ScopeGuard::allocate(impl::MemorySpaceType memory_space, char* address, void ScopeGuard::allocate_output(std::string label, std::string_view memory_space, char* data, std::size_t size) { + if (size == 0) { + return; + } + if (impl::memory_space_type_from_string(memory_space) == impl::MemorySpaceType::HOST) { - host_output_allocations.insert( - {label, impl::regular_host_allocate(size, data)}); + host_output_allocations.insert_or_assign( + label, impl::regular_host_allocate(size, data)); } else { #if !defined(KERNEL_REPLAYER_HAS_DEVICE_SPACE) throw std::runtime_error( "Trying to access device allocations but no device space is enabled"); #else - device_output_allocations.insert( - {label, impl::regular_device_allocate(size, data)}); + device_output_allocations.insert_or_assign( + label, impl::regular_device_allocate(size, data)); #endif } } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index efcd784..c67d254 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -44,3 +44,4 @@ add_subdirectory(struct_functor) add_subdirectory(small_struct_functor) add_subdirectory(repeated_kernel_invocation) add_subdirectory(metadata) +add_subdirectory(empty_view) diff --git a/tests/empty_view/CMakeLists.txt b/tests/empty_view/CMakeLists.txt new file mode 100644 index 0000000..5038641 --- /dev/null +++ b/tests/empty_view/CMakeLists.txt @@ -0,0 +1,7 @@ +add_executable(empty_view main.cpp) +target_link_libraries(empty_view PRIVATE kkf_warnings kkf_sanitizers Kokkos::kokkos kernel_extractor) + +add_executable(empty_view_replay replay_main.cpp) +target_link_libraries(empty_view_replay PRIVATE kkf_warnings kkf_sanitizers Kokkos::kokkos kernel_replayer_sanitized) + +add_replay_test(empty_view FALSE) diff --git a/tests/empty_view/main.cpp b/tests/empty_view/main.cpp new file mode 100644 index 0000000..7cc2f7e --- /dev/null +++ b/tests/empty_view/main.cpp @@ -0,0 +1,26 @@ +#include + +#include + +int main(int argc, char* argv[]) { + Kokkos::ScopeGuard kokkos_scope(argc, argv); + + Kokkos::View values("empty_values", 0); + + Kokkos::parallel_for( + "test_kernel", Kokkos::RangePolicy<>(0, 1), + cexa::kernel_replayer::replay_functor(KOKKOS_LAMBDA(int) { + if (values.extent(0) != 0) { + values(0) = 1; + } + })); + Kokkos::fence(); + + if (values.extent(0) != 0) { + Kokkos::printf("Expected an empty view, got extent %zu\n", + values.extent(0)); + return 1; + } + + return 0; +} diff --git a/tests/empty_view/replay_main.cpp b/tests/empty_view/replay_main.cpp new file mode 100644 index 0000000..2da1c8c --- /dev/null +++ b/tests/empty_view/replay_main.cpp @@ -0,0 +1,45 @@ +#include + +#include + +#include + +int main(int argc, char* argv[]) { + cexa::kernel_replayer::ScopeGuard replay_scope(argc, argv); + Kokkos::ScopeGuard kokkos_scope(argc, argv); + + Kokkos::View values; + + Kokkos::parallel_for( + "test_kernel", Kokkos::RangePolicy<>(0, 1), + cexa::kernel_replayer::replay_functor(KOKKOS_LAMBDA(int) { + if (values.extent(0) != 0) { + values(0) = 1; + } + })); + Kokkos::fence(); + + int* in_data = static_cast( + cexa::kernel_replayer::get_allocation("empty_values")); + int* out_data = static_cast( + cexa::kernel_replayer::get_out_allocation( + "empty_values")); + + if (in_data != nullptr || out_data != nullptr) { + Kokkos::printf("Expected null replay pointers for an empty view\n"); + return 1; + } + + cexa::kernel_replayer::compare_views( + "empty_values", std::make_tuple(0), + [](auto ref_values, auto replay_values) { + if (ref_values.extent(0) != 0 || replay_values.extent(0) != 0) { + Kokkos::printf( + "Expected empty replay views, got extents %zu and %zu\n", + ref_values.extent(0), replay_values.extent(0)); + std::exit(1); + } + }); + + return 0; +}