wehub-resource-sync
60 lines
1.9 KiB
C++
60 lines
1.9 KiB
C++
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
#include "completion_recorder.h"
|
|
|
|
#include <utility>
|
|
|
|
CompletionRecorder& CompletionRecorder::instance() {
|
|
static CompletionRecorder recorder;
|
|
return recorder;
|
|
}
|
|
|
|
void CompletionRecorder::push(std::unique_ptr<PendingCompletion> completion) {
|
|
std::lock_guard<std::mutex> lock(mutex_);
|
|
buffer_.push_back(std::move(completion));
|
|
}
|
|
|
|
std::vector<std::unique_ptr<PendingCompletion>> CompletionRecorder::drain() {
|
|
std::lock_guard<std::mutex> lock(mutex_);
|
|
std::vector<std::unique_ptr<PendingCompletion>> result;
|
|
result.swap(buffer_);
|
|
return result;
|
|
}
|
|
|
|
static void
|
|
#ifndef USE_ROCM
|
|
CUDART_CB
|
|
#endif
|
|
completion_host_callback(void* data) {
|
|
// Adopt the raw pointer back into a unique_ptr.
|
|
std::unique_ptr<PendingCompletion> completion(
|
|
static_cast<PendingCompletion*>(data));
|
|
CompletionRecorder::instance().push(std::move(completion));
|
|
}
|
|
|
|
void record_completion_on_stream(int64_t cuda_stream_ptr,
|
|
const std::string& kind, std::string payload) {
|
|
auto completion = std::make_unique<PendingCompletion>(
|
|
PendingCompletion{kind, std::move(payload)});
|
|
auto stream = reinterpret_cast<lmcache_completion_stream_t>(
|
|
static_cast<uintptr_t>(cuda_stream_ptr));
|
|
// Pass ownership through the driver as a raw pointer; the host callback
|
|
// re-adopts it. Reclaim if the launch itself fails.
|
|
PendingCompletion* raw = completion.release();
|
|
auto err = LMCACHE_COMPLETION_LAUNCH_HOST_FUNC(stream,
|
|
completion_host_callback, raw);
|
|
if (err != 0) {
|
|
delete raw;
|
|
}
|
|
}
|
|
|
|
CompletionDrainResult drain_recorded_completions() {
|
|
auto completions = CompletionRecorder::instance().drain();
|
|
CompletionDrainResult result;
|
|
result.reserve(completions.size());
|
|
for (auto& c : completions) {
|
|
result.emplace_back(std::move(c->kind), std::move(c->payload));
|
|
}
|
|
return result;
|
|
}
|