L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[CUDA] Fallback Event impl when there is no hardware cpu/gpu coherency (#3070)

Browse Source
This commit is contained in:
Cheng
2026-01-28 10:43:22 +09:00
committed by GitHub
parent b537b3685f
commit 2ac18eddb9
4 changed files with 124 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/cuda/device.cpp
+6
View File
@@ -42,6 +42,12 @@ Device::Device(int device) : device_(device) {
&concurrent_managed_access_,
cudaDevAttrConcurrentManagedAccess,
device_));
CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
&host_native_atomic_, cudaDevAttrHostNativeAtomicSupported, device_));
CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
&managed_memory_, cudaDevAttrManagedMemory, device_));
CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
&memory_pools_, cudaDevAttrMemoryPoolsSupported, device_));
}
Device::~Device() {
mlx/backend/cuda/device.h
+12
View File
@@ -164,12 +164,24 @@ class Device {
bool concurrent_managed_access() const {
return concurrent_managed_access_ == 1;
}
bool host_native_atomic() const {
return host_native_atomic_ == 1;
}
bool managed_memory() const {
return managed_memory_ == 1;
}
bool memory_pools() const {
return memory_pools_ == 1;
}
private:
int device_;
int compute_capability_major_;
int compute_capability_minor_;
int concurrent_managed_access_;
int host_native_atomic_;
int managed_memory_;
int memory_pools_;
std::string device_name_;
cublasLtHandle_t cublaslt_handle_{nullptr};
cudnnHandle_t cudnn_handle_{nullptr};
mlx/backend/cuda/event.cu
+92 -44
View File
@@ -174,62 +174,94 @@ class CopyableCudaEvent {
// AtomicEvent implementations
///////////////////////////////////////////////////////////////////////////////
__host__ __device__ void event_wait(AtomicEvent::Atomic* ac, uint64_t value) {
uint64_t current;
while ((current = ac->load()) < value) {
ac->wait(current);
__host__ __device__ void event_wait(uint32_t* ptr, uint32_t value) {
cuda::atomic_ref<uint32_t> ac(*ptr);
uint32_t current;
while ((current = ac.load()) < value) {
ac.wait(current);
}
}
__host__ __device__ void event_signal(AtomicEvent::Atomic* ac, uint64_t value) {
ac->store(value);
ac->notify_all();
__host__ __device__ void event_signal(uint32_t* ptr, uint32_t value) {
cuda::atomic_ref<uint32_t> ac(*ptr);
ac.store(value);
ac.notify_all();
}
__global__ void event_wait_kernel(AtomicEvent::Atomic* ac, uint64_t value) {
event_wait(ac, value);
__global__ void event_wait_kernel(uint32_t* ptr, uint32_t value) {
event_wait(ptr, value);
}
__global__ void event_signal_kernel(AtomicEvent::Atomic* ac, uint64_t value) {
event_signal(ac, value);
__global__ void event_signal_kernel(uint32_t* ptr, uint32_t value) {
__threadfence_system();
event_signal(ptr, value);
__threadfence_system();
}
bool supports_concurrent_managed_access() {
static bool concurrent_managed_access = []() {
auto check_gpu_coherency() {
static auto coherency = []() {
int device_count = gpu::device_count();
bool concurrent_managed_access = true;
bool host_native_atomic = true;
for (int i = 0; i < device_count; ++i) {
if (!cu::device(i).concurrent_managed_access()) {
return false;
}
auto& d = cu::device(i);
concurrent_managed_access &= d.concurrent_managed_access();
host_native_atomic &= d.host_native_atomic();
}
return true;
return std::make_tuple(concurrent_managed_access, host_native_atomic);
}();
return concurrent_managed_access;
return coherency;
}
AtomicEvent::AtomicEvent() {
if (!supports_concurrent_managed_access()) {
throw std::runtime_error(
"Device does not support synchronization in managed memory.");
void* buf;
cudaError_t (*cuda_free)(void*);
// There are 3 kinds of systems we are implementing for:
// 1. concurrentManagedAccess == true
// => use cuda::atom_ref on managed memory
// 2. hostNativeAtomicSupported == true
// => use cuda::atom_ref on pinned host memory
// 2. no hardware cpu/gpu coherency
// => use cuda::atom_ref on device memory
auto [concurrent_managed_access, host_native_atomic] = check_gpu_coherency();
if (concurrent_managed_access) {
CHECK_CUDA_ERROR(cudaMallocManaged(&buf, sizeof(uint32_t)));
cuda_free = cudaFree;
coherent_ = true;
} else if (host_native_atomic) {
CHECK_CUDA_ERROR(cudaMallocHost(&buf, sizeof(uint32_t)));
cuda_free = cudaFreeHost;
coherent_ = true;
} else {
CHECK_CUDA_ERROR(cudaMalloc(&buf, sizeof(uint32_t)));
cuda_free = cudaFree;
coherent_ = false;
}
buf_ = std::shared_ptr<void>(
buf, [cuda_free](void* buf) { CHECK_CUDA_ERROR(cuda_free(buf)); });
if (coherent_) {
*ptr() = 0;
} else {
CHECK_CUDA_ERROR(cudaMemset(buf, 0, sizeof(uint32_t)));
}
buf_ = std::shared_ptr<Buffer>(
new Buffer{allocator().malloc(sizeof(Atomic))}, [](Buffer* ptr) {
allocator().free(*ptr);
delete ptr;
});
*static_cast<uint64_t*>(buf_->raw_ptr()) = 0;
}
void AtomicEvent::wait(uint64_t value) {
void AtomicEvent::wait(uint32_t value) {
nvtx3::scoped_range r("cu::AtomicEvent::wait");
event_wait(atomic(), value);
if (coherent_) {
event_wait(ptr(), value);
} else {
while (!is_signaled(value)) {
std::this_thread::yield();
}
}
}
void AtomicEvent::wait(cudaStream_t stream, uint64_t value) {
event_wait_kernel<<<1, 1, 0, stream>>>(atomic(), value);
void AtomicEvent::wait(cudaStream_t stream, uint32_t value) {
event_wait_kernel<<<1, 1, 0, stream>>>(ptr(), value);
}
void AtomicEvent::wait(Stream s, uint64_t value) {
void AtomicEvent::wait(Stream s, uint32_t value) {
nvtx3::scoped_range r("cu::AtomicEvent::wait(s)");
if (s.device == mlx::core::Device::cpu) {
scheduler::enqueue(s, [*this, value]() mutable { wait(value); });
@@ -241,22 +273,26 @@ void AtomicEvent::wait(Stream s, uint64_t value) {
}
}
void AtomicEvent::signal(uint64_t value) {
void AtomicEvent::signal(uint32_t value) {
nvtx3::scoped_range r("cu::AtomicEvent::signal");
event_signal(atomic(), value);
if (coherent_) {
event_signal(ptr(), value);
} else {
signal(signal_stream(), value);
}
}
void AtomicEvent::signal(cudaStream_t stream, uint64_t value) {
event_signal_kernel<<<1, 1, 0, stream>>>(atomic(), value);
void AtomicEvent::signal(cudaStream_t stream, uint32_t value) {
event_signal_kernel<<<1, 1, 0, stream>>>(ptr(), value);
}
void AtomicEvent::signal(Stream s, uint64_t value) {
void AtomicEvent::signal(Stream s, uint32_t value) {
nvtx3::scoped_range r("cu::AtomicEvent::signal(s)");
if (s.device == mlx::core::Device::cpu) {
// Signal through a GPU stream so the atomic is updated in GPU - updating
// the atomic in CPU sometimes does not get GPU notified.
static CudaStream stream(device(mlx::core::Device::gpu));
scheduler::enqueue(s, [*this, value]() mutable { signal(stream, value); });
scheduler::enqueue(
s, [*this, value]() mutable { signal(signal_stream(), value); });
} else {
auto& encoder = get_command_encoder(s);
encoder.commit();
@@ -265,14 +301,26 @@ void AtomicEvent::signal(Stream s, uint64_t value) {
}
}
bool AtomicEvent::is_signaled(uint64_t value) const {
nvtx3::scoped_range r("cu::AtomicEvent::is_signaled");
return atomic()->load() >= value;
bool AtomicEvent::is_signaled(uint32_t val) const {
return value() >= val;
}
uint64_t AtomicEvent::value() const {
uint32_t AtomicEvent::value() const {
nvtx3::scoped_range r("cu::AtomicEvent::value");
return atomic()->load();
if (coherent_) {
cuda::atomic_ref<uint32_t> ac(*ptr());
return ac.load();
} else {
uint32_t val;
CHECK_CUDA_ERROR(
cudaMemcpy(&val, ptr(), sizeof(uint32_t), cudaMemcpyDeviceToHost));
return val;
}
}
const CudaStream& AtomicEvent::signal_stream() {
static CudaStream stream(device(0));
return stream;
}
} // namespace cu
mlx/backend/cuda/event.h
+14 -13
View File
@@ -54,25 +54,26 @@ class CudaEvent {
// CudaEvent so the latter should always be preferred when possible.
class AtomicEvent {
public:
using Atomic = cuda::atomic<uint64_t>;
AtomicEvent();
void wait(uint64_t value);
void wait(cudaStream_t stream, uint64_t value);
void wait(Stream s, uint64_t value);
void signal(uint64_t value);
void signal(cudaStream_t stream, uint64_t value);
void signal(Stream s, uint64_t value);
bool is_signaled(uint64_t value) const;
uint64_t value() const;
void wait(uint32_t value);
void wait(cudaStream_t stream, uint32_t value);
void wait(Stream s, uint32_t value);
void signal(uint32_t value);
void signal(cudaStream_t stream, uint32_t value);
void signal(Stream s, uint32_t value);
bool is_signaled(uint32_t value) const;
uint32_t value() const;
private:
Atomic* atomic() const {
return static_cast<AtomicEvent::Atomic*>(buf_->raw_ptr());
const CudaStream& signal_stream();
uint32_t* ptr() const {
return static_cast<uint32_t*>(buf_.get());
}
std::shared_ptr<allocator::Buffer> buf_;
bool coherent_;
std::shared_ptr<void> buf_;
};
} // namespace mlx::core::cu