Add dump support for empty views (#34)

* Add dump support for empty views

* Fix reported error by asan and ubsan

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* Add replay test for view without explicit size

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* make empty view test usable with other backends

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* fix: make HIP allocation range query compile

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* add TODO comments

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* update doc and add comment

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* remove todo

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* Fix GPU allocation range queries

* Limit allocation range queries

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* remove kernel_replayer namespace

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* pre-commit

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* Handle zero-byte allocations before backend checks

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* Add comments

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* Remove test

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* Update src/krepe/capture/kokkos_hooks.cpp

Co-authored-by: Trévis Morvany <63788850+tretre91@users.noreply.github.com>

* remove ci test

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

* ci: retrigger checks

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>

---------

Signed-off-by: daxmawal <jeanfrancoismanutea@gmail.com>
Co-authored-by: Trévis Morvany <63788850+tretre91@users.noreply.github.com>

Author: daxmawal

docs/hdf5-dumps.md

@@ -23,19 +23,22 @@ 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
-bytes_dumped
+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 bytes, 0 for empty or unbounded allocations
+reported_size # size reported by the Kokkos profiling allocation hook
+bytes_dumped  # 1 when the bytes dataset was written, 0 when it was skipped
+skip_reason   # present only when bytes_dumped is 0
+bytes         # byte dataset, present when bytes_dumped is 1
 ```
 
 The `/metadata` group contains user specified metadata

src/krepe/capture/CMakeLists.txt

@@ -35,7 +35,7 @@ target_compile_definitions(
 if(Kokkos_ENABLE_CUDA)
   find_package(CUDAToolkit REQUIRED)
   target_compile_definitions(krepe_kokkos_hooks PRIVATE KREPE_ENABLE_CUDA_DUMP)
-  target_link_libraries(krepe_kokkos_hooks PRIVATE CUDA::cudart)
+  target_link_libraries(krepe_kokkos_hooks PRIVATE CUDA::cudart CUDA::cuda_driver)
 endif()
 if(Kokkos_ENABLE_HIP)
   target_compile_definitions(krepe_kokkos_hooks PRIVATE KREPE_ENABLE_HIP_DUMP)

src/krepe/capture/allocation_tracker.cpp

@@ -6,12 +6,15 @@ namespace krepe {
 
 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<std::mutex> lock(mutex_);
   active_allocations_[space].insert_or_assign(ptr, std::move(record));
 }

src/krepe/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;

src/krepe/capture/kokkos_hooks.cpp

@@ -31,6 +31,19 @@
 #include <variant>
 #include <vector>
 
+#if defined(__linux__)
+#include <malloc.h>
+#endif
+
+#if defined(KREPE_ENABLE_CUDA_DUMP)
+#include <cuda.h>
+#include <cuda_runtime_api.h>
+#endif
+
+#if defined(KREPE_ENABLE_HIP_DUMP)
+#include <hip/hip_runtime_api.h>
+#endif
+
 #define KOKKOS_HOOKS_EXPORT __attribute__((visibility("default")))
 
 #if !defined(KREPE_KOKKOS_ALLOCATION_HEADER_SIZE)
@@ -89,6 +102,20 @@ bool is_internal_label(std::string_view label) {
          label.starts_with("KOKKOS_") || label.starts_with("kokkos.");
 }
 
+#if defined(KREPE_ENABLE_CUDA_DUMP)
+bool is_cuda_pointer_attribute_space(const std::string& space) {
+  return space == "Cuda" || space == "CudaUVM" || space == "CudaHostPinned";
+}
+#endif
+
+#if defined(KREPE_ENABLE_HIP_DUMP)
+bool is_hip_pointer_attribute_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 "<null>";
@@ -166,6 +193,93 @@ const void* allocation_data_pointer(const void* ptr) {
          KREPE_KOKKOS_ALLOCATION_HEADER_SIZE;
 }
 
+std::size_t host_allocation_size(const void* ptr) {
+#if defined(__linux__)
+  // The allocator may reserve more memory than Kokkos requested, so
+  // malloc_usable_size(ptr) can return a larger size than the requested one.
+  return malloc_usable_size(const_cast<void*>(ptr));
+#else
+  (void)ptr;
+  return 0;
+#endif
+}
+
+std::optional<std::uint64_t> 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<std::uintptr_t>(allocation_base);
+  const auto data = reinterpret_cast<std::uintptr_t>(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<std::uint64_t>(allocation_size - offset);
+  return available >= reported_size ? reported_size : 0;
+}
+
+std::optional<std::uint64_t> allocation_data_size(
+    const std::string& space, const void* ptr, const void* data_ptr,
+    const std::uint64_t reported_size) {
+  // Kokkos <= 5.2.0 reports the SharedAllocationHeader size instead of zero
+  // for empty Views (kokkos/kokkos#9337). Disambiguate that legacy value using
+  // the physical allocation bounds below.
+  if (reported_size != KREPE_KOKKOS_ALLOCATION_HEADER_SIZE) {
+    return reported_size;
+  }
+
+#if defined(KREPE_ENABLE_CUDA_DUMP)
+  if (is_cuda_pointer_attribute_space(space)) {
+    CUdeviceptr allocation_base       = 0;
+    std::size_t allocation_size       = 0;
+    const CUdeviceptr queried_pointer = reinterpret_cast<CUdeviceptr>(ptr);
+    const CUresult base_error         = cuPointerGetAttribute(
+        &allocation_base, CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
+        queried_pointer);
+    const CUresult size_error = cuPointerGetAttribute(
+        &allocation_size, CU_POINTER_ATTRIBUTE_RANGE_SIZE, queried_pointer);
+    if (base_error == CUDA_SUCCESS && size_error == CUDA_SUCCESS) {
+      return validated_data_size(reinterpret_cast<const void*>(allocation_base),
+                                 allocation_size, data_ptr, reported_size);
+    }
+  }
+#endif
+
+#if defined(KREPE_ENABLE_HIP_DUMP)
+  if (is_hip_pointer_attribute_space(space)) {
+    hipDeviceptr_t allocation_base       = nullptr;
+    std::size_t allocation_size          = 0;
+    const hipDeviceptr_t queried_pointer = const_cast<void*>(ptr);
+    const hipError_t base_error          = hipPointerGetAttribute(
+        &allocation_base, HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR,
+        queried_pointer);
+    const hipError_t size_error = hipPointerGetAttribute(
+        &allocation_size, HIP_POINTER_ATTRIBUTE_RANGE_SIZE, queried_pointer);
+    if (base_error == hipSuccess && size_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);
@@ -353,11 +467,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<std::uint64_t> 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());
   }
 }
 

src/krepe/capture/memory_copy.cpp

@@ -44,6 +44,14 @@ std::string allocate_staging_buffer(const ActiveAllocation& allocation,
 
 std::string copy_allocation_bytes(const ActiveAllocation& allocation,
                                   std::vector<unsigned char>& bytes) {
+  if (!allocation.record.data_size_known) {
+    return "allocation data size could not be bounded safely";
+  }
+  if (allocation.record.size == 0) {
+    bytes.clear();
+    return {};
+  }
+
   const std::string& space   = allocation.record.space;
   const bool host_accessible = is_host_accessible_space(space);
   const bool cuda_device     = is_cuda_device_space(space);

src/krepe/capture/view_dump.cpp

@@ -204,6 +204,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<unsigned char> bytes;
   const std::string skip_reason = copy_allocation_bytes(allocation, bytes);

src/krepe/replay/kernel_replayer.cpp

@@ -299,6 +299,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<std::set<std::pair<char*, std::size_t>>*>(
       allocation_set);
@@ -340,6 +343,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<hdf5_iterate_fun_t*>(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) {
@@ -667,6 +675,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;
@@ -750,17 +769,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
   }
 }

tests/CMakeLists.txt

@@ -52,3 +52,8 @@ add_subdirectory(mdrange_iteration_pattern)
 add_subdirectory(team_policy)
 add_subdirectory(non_default_execution_space)
 add_subdirectory(index_type)
+add_subdirectory(empty_view)
+
+if(Kokkos_ENABLE_CUDA OR Kokkos_ENABLE_HIP)
+  add_subdirectory(gpu_memory_spaces)
+endif()

tests/empty_view/CMakeLists.txt

@@ -0,0 +1,19 @@
+add_executable(empty_view main.cpp)
+target_link_libraries(empty_view PRIVATE krepe_warnings krepe_sanitizers Kokkos::kokkos kernel_extractor)
+
+add_executable(empty_view_replay replay_main.cpp)
+target_link_libraries(
+  empty_view_replay PRIVATE krepe_warnings krepe_sanitizers Kokkos::kokkos kernel_replayer_sanitized
+)
+
+add_replay_test(empty_view FALSE)
+
+add_executable(empty_view_without_extent without_extent_main.cpp)
+target_link_libraries(empty_view_without_extent PRIVATE krepe_warnings krepe_sanitizers Kokkos::kokkos kernel_extractor)
+
+add_executable(empty_view_without_extent_replay without_extent_replay_main.cpp)
+target_link_libraries(
+  empty_view_without_extent_replay PRIVATE krepe_warnings krepe_sanitizers Kokkos::kokkos kernel_replayer_sanitized
+)
+
+add_replay_test(empty_view_without_extent FALSE)