mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-07-18 09:47:34 +00:00
113 lines
3.7 KiB
C++
113 lines
3.7 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// Licensed under the MIT License.
|
|
|
|
// In-process CUPTI kernel timer for the mscclpp LL benchmark.
|
|
// name, and exposes an extern "C" ABI so it can be driven from Python via ctypes
|
|
// (no cuda-python CUPTI bindings exist in this env, but libcupti.so is loadable).
|
|
//
|
|
// This is near-zero host perturbation (out-of-band buffer callbacks), unlike
|
|
// torch.profiler's in-process tracing which serialized the LL dispatch recv-spin
|
|
// and inflated one rank's device time into the millisecond range. It matches
|
|
// ep_bench's methodology exactly: start() after warmup, stop() after the timed
|
|
// loop, get_avg_us("dispatch"/"combine") buckets by mangled-name substring.
|
|
//
|
|
// CONCURRENT_KERNEL records ordinary LL kernel launches without serializing
|
|
// activity. The record carries the raw mangled name; the caller matches
|
|
// "dispatch"/"combine" substrings.
|
|
//
|
|
// Build (host-only C++, links libcupti):
|
|
// g++ -O2 -fPIC -shared cupti_kernel_timer.cpp -o libcupti_kernel_timer.so \
|
|
// -I<cuda>/targets/sbsa-linux/include -L<cuda>/targets/sbsa-linux/lib -lcupti
|
|
#include <cupti.h>
|
|
|
|
#include <cstdint>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <map>
|
|
#include <mutex>
|
|
#include <string>
|
|
|
|
namespace {
|
|
|
|
struct KernelStat {
|
|
uint64_t total_ns = 0;
|
|
uint64_t count = 0;
|
|
};
|
|
|
|
std::map<std::string, KernelStat> g_stats;
|
|
std::mutex g_mutex;
|
|
|
|
constexpr size_t kBufSize = 8 * 1024 * 1024; // 8 MB, matches ep_bench
|
|
|
|
void CUPTIAPI bufferRequested(uint8_t** buffer, size_t* size, size_t* maxNumRecords) {
|
|
// 8-byte aligned; aligned_alloc requires size to be a multiple of alignment.
|
|
*buffer = static_cast<uint8_t*>(aligned_alloc(8, kBufSize));
|
|
*size = kBufSize;
|
|
*maxNumRecords = 0;
|
|
}
|
|
|
|
void CUPTIAPI bufferCompleted(CUcontext, uint32_t, uint8_t* buffer, size_t, size_t validSize) {
|
|
CUpti_Activity* record = nullptr;
|
|
std::lock_guard<std::mutex> lock(g_mutex);
|
|
while (cuptiActivityGetNextRecord(buffer, validSize, &record) == CUPTI_SUCCESS) {
|
|
if (record->kind == CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL) {
|
|
auto* k = reinterpret_cast<CUpti_ActivityKernel9*>(record);
|
|
if (k->name) {
|
|
auto& s = g_stats[k->name];
|
|
s.total_ns += (k->end - k->start);
|
|
s.count += 1;
|
|
}
|
|
}
|
|
}
|
|
free(buffer);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
extern "C" {
|
|
|
|
// Clear stats, register callbacks, and enable kernel activity recording after warmup.
|
|
int kt_start() {
|
|
{
|
|
std::lock_guard<std::mutex> lock(g_mutex);
|
|
g_stats.clear();
|
|
}
|
|
CUptiResult r = cuptiActivityRegisterCallbacks(bufferRequested, bufferCompleted);
|
|
if (r != CUPTI_SUCCESS) return static_cast<int>(r);
|
|
r = cuptiActivityEnable(CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL);
|
|
return static_cast<int>(r);
|
|
}
|
|
|
|
// Flush pending buffers and disable recording. Returns CUPTI result code (0=ok).
|
|
int kt_stop() {
|
|
cuptiActivityFlushAll(0);
|
|
CUptiResult r = cuptiActivityDisable(CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL);
|
|
return static_cast<int>(r);
|
|
}
|
|
|
|
// Average GPU execution time (microseconds) over every recorded kernel whose
|
|
// mangled name contains `substr`. Returns 0 if none matched.
|
|
double kt_get_avg_us(const char* substr) {
|
|
std::lock_guard<std::mutex> lock(g_mutex);
|
|
uint64_t total_ns = 0, count = 0;
|
|
for (const auto& kv : g_stats) {
|
|
if (kv.first.find(substr) != std::string::npos) {
|
|
total_ns += kv.second.total_ns;
|
|
count += kv.second.count;
|
|
}
|
|
}
|
|
return count ? static_cast<double>(total_ns) / static_cast<double>(count) / 1000.0 : 0.0;
|
|
}
|
|
|
|
// Number of recorded kernel instances whose name contains `substr`.
|
|
long kt_get_count(const char* substr) {
|
|
std::lock_guard<std::mutex> lock(g_mutex);
|
|
uint64_t count = 0;
|
|
for (const auto& kv : g_stats) {
|
|
if (kv.first.find(substr) != std::string::npos) count += kv.second.count;
|
|
}
|
|
return static_cast<long>(count);
|
|
}
|
|
|
|
} // extern "C"
|