zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Avoid invoking allocator::malloc when creating CUDA event (#2232)

Browse Source
This commit is contained in:
Cheng
2025-06-04 08:48:40 +09:00
committed by GitHub
parent 0408ba0a76
commit 5685ceb3c7
3 changed files with 33 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
mlx/backend/cuda/allocator.h
View File

@@ -34,6 +34,9 @@ class CudaAllocator : public allocator::Allocator {
// buffers there would result in dead lock.
void register_this_thread();
// Call cudaFree in the safe thread.
void cuda_free(void* buf);
size_t get_active_memory() const;
size_t get_peak_memory() const;
void reset_peak_memory();
@@ -47,8 +50,6 @@ class CudaAllocator : public allocator::Allocator {
CudaAllocator();
friend CudaAllocator& allocator();
void cuda_free(CudaBuffer* buf);
std::mutex worker_mutex_;
std::unique_ptr<Worker> worker_;
std::set<std::thread::id> allowed_threads_;