Add dump support for empty views

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
 ```

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)

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<std::mutex> lock(mutex_);
   active_allocations_[space].insert_or_assign(ptr, std::move(record));
 }

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;

src/kkf/capture/kokkos_hooks.cpp

@@ -30,6 +30,19 @@
 #include <unordered_map>
 #include <vector>
 
+#if defined(__linux__)
+#include <malloc.h>
+#endif
+
+#if defined(KKF_ENABLE_CUDA_DUMP)
+#include <cuda.h>
+#include <cuda_runtime_api.h>
+#endif
+
+#if defined(KKF_ENABLE_HIP_DUMP)
+#include <hip/hip_runtime_api.h>
+#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 "<null>";
@@ -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<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) {
+#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<CUdeviceptr>(ptr));
+    if (error == CUDA_SUCCESS) {
+      return validated_data_size(reinterpret_cast<const void*>(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<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/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";

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

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<std::set<std::pair<char*, std::size_t>>*>(
       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<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) {
@@ -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
   }
 }

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)

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)

tests/empty_view/main.cpp

@@ -0,0 +1,26 @@
+#include <Kokkos_Core.hpp>
+
+#include <kkf/extractor.hpp>
+
+int main(int argc, char* argv[]) {
+  Kokkos::ScopeGuard kokkos_scope(argc, argv);
+
+  Kokkos::View<int*> 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;
+}

tests/empty_view/replay_main.cpp

@@ -0,0 +1,45 @@
+#include <Kokkos_Core.hpp>
+
+#include <kkf/replayer.hpp>
+
+#include <cstdlib>
+
+int main(int argc, char* argv[]) {
+  cexa::kernel_replayer::ScopeGuard replay_scope(argc, argv);
+  Kokkos::ScopeGuard kokkos_scope(argc, argv);
+
+  Kokkos::View<int*> 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<int*>(
+      cexa::kernel_replayer::get_allocation<Kokkos::HostSpace>("empty_values"));
+  int* out_data = static_cast<int*>(
+      cexa::kernel_replayer::get_out_allocation<Kokkos::HostSpace>(
+          "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<int*>(
+      "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;
+}