KokkosPerfCounters.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
//
// Copyright 2026 Caleb Buahin
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

 
#ifndef OPENSWMM_ENGINE_2D_GPU_KOKKOS_PERF_COUNTERS_HPP
#define OPENSWMM_ENGINE_2D_GPU_KOKKOS_PERF_COUNTERS_HPP
 
#include <Kokkos_Core.hpp>
 
#include <algorithm>
#include <chrono>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <map>
#include <string>
#include <vector>
 
#include "../../core/PerfTimers.hpp"
 
namespace openswmm::twoD::gpu::kperf {
 
struct KernelEntry {
    std::string label;
    long n = 0;       
    double sec = 0.0; 
};
 
struct Counters {
    long n_advance = 0;
    long n_launch = 0; 
    long n_fence = 0;
    double sec_fence = 0.0;
    long n_deep_copy = 0;
    unsigned long long bytes_h2d = 0, bytes_d2h = 0, bytes_same = 0;
    std::vector<KernelEntry> kernels; 
    std::map<std::string, std::uint64_t> label_id;
};
 
// Each solver owns its counters. Only dispatches in its active call scope
// are attributed to it; calls from unrelated Kokkos clients are ignored.
inline thread_local Counters* active = nullptr;
class Scope {
public:
    explicit Scope(Counters& counters) : previous_(active) { active = &counters; }
    ~Scope() { active = previous_; }
    Scope(const Scope&) = delete;
 
private:
    Counters* previous_;
};
 
namespace detail {
 
using clock = std::chrono::steady_clock;

inline thread_local std::vector<clock::time_point> dispatch_t0;
inline thread_local std::vector<clock::time_point> fence_t0;

inline void begin_dispatch(const char* name, std::uint32_t /*devid*/, std::uint64_t* kID) {
    *kID = 0;
    if (!active)
        return;
    auto& g = *active;
    const auto it = g.label_id.find(name);
    std::uint64_t id;
    if (it == g.label_id.end()) {
        id = g.kernels.size();
        g.label_id.emplace(name, id);
        g.kernels.push_back({name, 0, 0.0});
    } else {
        id = it->second;
    }
    *kID = id;
    ++g.n_launch;
    dispatch_t0.push_back(clock::now());
}

inline void end_dispatch(std::uint64_t kID) {
    if (!active)
        return;
    auto& g = *active;
    if (dispatch_t0.empty())
        return;
    const double dt = std::chrono::duration<double>(clock::now() - dispatch_t0.back()).count();
    dispatch_t0.pop_back();
    if (kID < g.kernels.size()) {
        ++g.kernels[kID].n;
        g.kernels[kID].sec += dt;
    }
}

inline void begin_fence(const char* /*name*/, std::uint32_t /*devid*/, std::uint64_t* handle) {
    *handle = 0;
    if (!active)
        return;
    fence_t0.push_back(clock::now());
}

inline void end_fence(std::uint64_t /*handle*/) {
    if (!active)
        return;
    auto& g = *active;
    if (fence_t0.empty())
        return;
    g.sec_fence += std::chrono::duration<double>(clock::now() - fence_t0.back()).count();
    fence_t0.pop_back();
    ++g.n_fence;
}

inline void begin_deep_copy(Kokkos_Profiling_SpaceHandle dst, const char*, const void*,
                            Kokkos_Profiling_SpaceHandle src, const char*, const void*,
                            std::uint64_t size) {
    if (!active)
        return;
    auto& g = *active;
    ++g.n_deep_copy;
    const bool dst_host = std::strncmp(dst.name, "Host", 4) == 0;
    const bool src_host = std::strncmp(src.name, "Host", 4) == 0;
    if (src_host && !dst_host)
        g.bytes_h2d += size;
    else if (dst_host && !src_host)
        g.bytes_d2h += size;
    else
        g.bytes_same += size; // OMP backend: all
}
 
} // namespace detail

inline void install() {
    static const bool done = [] {
        if (!perf::enabled())
            return false;
        if (std::getenv("KOKKOS_TOOLS_LIBS") || std::getenv("KOKKOS_PROFILE_LIBRARY")) {
            std::fprintf(stderr, "[PERF-2D-KOKKOS] external Kokkos tool loaded; in-process "
                                 "counters disabled for this run.\n");
            return false;
        }
        namespace KTE = Kokkos::Tools::Experimental;
        KTE::set_begin_parallel_for_callback(detail::begin_dispatch);
        KTE::set_end_parallel_for_callback(detail::end_dispatch);
        KTE::set_begin_parallel_reduce_callback(detail::begin_dispatch);
        KTE::set_end_parallel_reduce_callback(detail::end_dispatch);
        KTE::set_begin_parallel_scan_callback(detail::begin_dispatch);
        KTE::set_end_parallel_scan_callback(detail::end_dispatch);
        KTE::set_begin_fence_callback(detail::begin_fence);
        KTE::set_end_fence_callback(detail::end_fence);
        KTE::set_begin_deep_copy_callback(detail::begin_deep_copy);
        return true;
    }();
    (void)done;
}

inline void reset() {
    if (!active)
        return;
    auto& g = *active;
    g = Counters{};
}

inline void count_advance() {
    if (!active)
        return;
    auto& g = *active;
    if (perf::enabled())
        ++g.n_advance;
}

inline void dump() {
    if (!active)
        return;
    auto& g = *active;
    if (!perf::enabled() || g.n_advance == 0)
        return;
    const double adv = static_cast<double>(g.n_advance);
    std::fprintf(stderr,
                 "[PERF-2D-KOKKOS] advances=%ld launches=%ld launches_per_advance=%.1f "
                 "fences=%ld fence_s=%.4f deep_copies=%ld "
                 "h2d_bytes=%llu d2h_bytes=%llu samespace_bytes=%llu "
                 "h2d_bytes_per_advance=%.0f d2h_bytes_per_advance=%.0f\n",
                 g.n_advance, g.n_launch, g.n_launch / adv, g.n_fence, g.sec_fence, g.n_deep_copy,
                 g.bytes_h2d, g.bytes_d2h, g.bytes_same, g.bytes_h2d / adv, g.bytes_d2h / adv);
    std::vector<const KernelEntry*> order;
    order.reserve(g.kernels.size());
    for (const auto& k : g.kernels)
        order.push_back(&k);
    std::sort(order.begin(), order.end(),
              [](const KernelEntry* a, const KernelEntry* b) { return a->sec > b->sec; });
    for (const KernelEntry* k : order)
        std::fprintf(stderr, "[PERF-2D-KOKKOS-KERNEL] name=%s n=%ld s=%.4f per_advance=%.1f\n",
                     k->label.c_str(), k->n, k->sec, k->n / adv);
}
 
} // namespace openswmm::twoD::gpu::kperf
 
#endif // OPENSWMM_ENGINE_2D_GPU_KOKKOS_PERF_COUNTERS_HPP