Skip to content

Commit

Permalink
Re-enabling kokkos allocation tracking
Browse files Browse the repository at this point in the history
When enabled, APEX will keep track of allocations through Kokkos
and ensure they are all freed before exit
  • Loading branch information
khuck committed Feb 28, 2024
1 parent ad74dc5 commit 806c489
Showing 1 changed file with 24 additions and 15 deletions.
39 changes: 24 additions & 15 deletions src/apex/apex_kokkos.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -29,13 +29,12 @@
#include "apex.hpp"
#include "Kokkos_Profiling_C_Interface.h"

/*
static std::mutex memory_mtx;
static std::unordered_map<void*,apex::profiler*>& memory_map() {
static std::unordered_map<void*,apex::profiler*> themap;
static std::unordered_map<void*,std::string>& memory_map() {
static std::unordered_map<void*,std::string> themap;
return themap;
}
*/

static std::stack<apex::profiler*>& timer_stack() {
static APEX_NATIVE_TLS std::stack<apex::profiler*> thestack;
return thestack;
Expand Down Expand Up @@ -147,9 +146,21 @@ void kokkosp_init_library(int loadseq, uint64_t version,
* profiling hooks.
*/
void kokkosp_finalize_library() {
memory_mtx.lock();
if (memory_map().size() == 0) {
if (apex::apex::instance()->get_node_id() == 0) {
std::cout << "No Kokkos allocation Leaks on rank 0!" << std::endl;
}
} else {
for (auto it : memory_map()) {
std::cerr << "Rank: " << apex::apex::instance()->get_node_id()
<< ", Kokkos allocation Leak: " << it.second << std::endl;
}
}
memory_mtx.unlock();
#ifndef APEX_HAVE_HPX
if (!apex::apex_options::use_mpi()) {
apex::finalize();
//apex::finalize();
}
#endif
}
Expand All @@ -174,6 +185,7 @@ void kokkosp_request_tool_settings(int num_actions,
*/
void kokkosp_begin_parallel_for(const char* name,
uint32_t devid, uint64_t* kernid) {
apex::in_apex prevent_memory_tracking;
std::stringstream ss;
ExecutionSpaceIdentifier space_id = identifier_from_devid(devid);
ss << "Kokkos::parallel_for ["
Expand All @@ -195,6 +207,7 @@ void kokkosp_begin_parallel_for(const char* name,

void kokkosp_begin_parallel_reduce(const char* name,
uint32_t devid, uint64_t* kernid) {
apex::in_apex prevent_memory_tracking;
std::stringstream ss;
ExecutionSpaceIdentifier space_id = identifier_from_devid(devid);
ss << "Kokkos::parallel_reduce ["
Expand All @@ -216,6 +229,7 @@ void kokkosp_begin_parallel_reduce(const char* name,

void kokkosp_begin_parallel_scan(const char* name,
uint32_t devid, uint64_t* kernid) {
apex::in_apex prevent_memory_tracking;
std::stringstream ss;
ExecutionSpaceIdentifier space_id = identifier_from_devid(devid);
ss << "Kokkos::parallel_scan ["
Expand Down Expand Up @@ -263,6 +277,7 @@ void kokkosp_end_parallel_scan(uint64_t kernid) {
* user.
*/
void kokkosp_push_profile_region(const char* name) {
apex::in_apex prevent_memory_tracking;
std::stringstream ss;
ss << "Kokkos region, " << name;
std::string tmp{ss.str()};
Expand Down Expand Up @@ -292,18 +307,15 @@ void kokkosp_pop_profile_region() {
*/
void kokkosp_allocate_data(SpaceHandle_t handle, const char* name,
void* ptr, uint64_t size) {
apex::in_apex prevent_memory_tracking;
APEX_UNUSED(ptr);
std::stringstream ss;
ss << "Kokkos " << handle.name << " data, " << name;
/*
std::string tmp{ss.str()};
auto p = apex::start(tmp);
std::string tmp2{ss.str()};
memory_mtx.lock();
memory_map().insert(std::pair<void*,apex::profiler*>(ptr, p));
memory_map().insert(std::pair<void*,std::string>(ptr, tmp2));
memory_mtx.unlock();
*/
ss << ": Bytes";
std::string tmp2{ss.str()};
double bytes = (double)(size);
if (apex::apex_options::use_kokkos_counters()) {
apex::sample_value(tmp2, bytes);
Expand All @@ -321,13 +333,9 @@ void kokkosp_deallocate_data(SpaceHandle handle, const char* name,
APEX_UNUSED(name);
APEX_UNUSED(ptr);
APEX_UNUSED(size);
/*
memory_mtx.lock();
auto p = memory_map()[ptr];
memory_map().erase(ptr);
memory_mtx.unlock();
apex::stop(p);
*/
}

/* This function will be called whenever a Kokkos::deep_copy function is
Expand All @@ -342,6 +350,7 @@ void kokkosp_begin_deep_copy(
SpaceHandle dst_handle, const char* dst_name, const void* dst_ptr,
SpaceHandle src_handle, const char* src_name, const void* src_ptr,
uint64_t size) {
apex::in_apex prevent_memory_tracking;
std::stringstream ss;
ss << "Kokkos deep copy: " << src_handle.name << " " << src_name
<< " -> " << dst_handle.name << " " << dst_name;
Expand Down

0 comments on commit 806c489

Please sign in to comment.