-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
14 changed files
with
945 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,7 @@ | ||
| target_sources( | ||
| mlx | ||
| PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp | ||
| ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp | ||
| ${CMAKE_CURRENT_SOURCE_DIR}/metal.cpp | ||
| ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp | ||
| ${CMAKE_CURRENT_SOURCE_DIR}/../no_metal/event.cpp) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,119 @@ | ||
| // Copyright © 2025 Apple Inc. | ||
|
|
||
| #include "mlx/backend/webgpu/allocator.h" | ||
|
|
||
| namespace mlx::core { | ||
|
|
||
| namespace allocator { | ||
|
|
||
| Allocator& allocator() { | ||
| return webgpu::allocator(); | ||
| } | ||
|
|
||
| void* Buffer::raw_ptr() { | ||
| return static_cast<webgpu::DoubleBuffer*>(ptr_)->cpu_data(); | ||
| } | ||
|
|
||
| } // namespace allocator | ||
|
|
||
| namespace webgpu { | ||
|
|
||
| DoubleBuffer::DoubleBuffer(size_t size) | ||
| : cpu_data_(std::malloc(size + sizeof(size_t))) { | ||
| *static_cast<size_t*>(cpu_data_) = size; | ||
| } | ||
|
|
||
| DoubleBuffer::DoubleBuffer(betann::Device& device, size_t size) | ||
| : gpu_data_(device.CreateBuffer( | ||
| size, | ||
| betann::BufferUsage::Storage | betann::BufferUsage::CopySrc)) {} | ||
|
|
||
| DoubleBuffer::~DoubleBuffer() { | ||
| std::free(cpu_data_); | ||
| } | ||
|
|
||
| void DoubleBuffer::copy_to_cpu(const void* data, size_t size) { | ||
| assert(!cpu_data_); | ||
| cpu_data_ = std::malloc(size + sizeof(size_t)); | ||
| *static_cast<size_t*>(cpu_data_) = size; | ||
| std::memcpy(cpu_data(), data, size); | ||
| } | ||
|
|
||
| size_t DoubleBuffer::size() const { | ||
| if (cpu_data_) | ||
| return *static_cast<size_t*>(cpu_data_); | ||
| if (gpu_data_) | ||
| return gpu_data_.GetSize(); | ||
| return 0; | ||
| } | ||
|
|
||
| WgpuAllocator::WgpuAllocator() : device_(webgpu::device(Device::gpu)) {} | ||
|
|
||
| Buffer | ||
| WgpuAllocator::malloc(const Device& device, size_t size, bool allow_swap) { | ||
| if (device.type == Device::gpu) | ||
| return Buffer(new DoubleBuffer(webgpu::device(device), size)); | ||
| else | ||
| return Buffer(new DoubleBuffer(size)); | ||
| } | ||
|
|
||
| void WgpuAllocator::free(Buffer buffer) { | ||
| delete static_cast<DoubleBuffer*>(buffer.ptr()); | ||
| } | ||
|
|
||
| size_t WgpuAllocator::size(Buffer buffer) const { | ||
| return static_cast<DoubleBuffer*>(buffer.ptr())->size(); | ||
| } | ||
|
|
||
| void WgpuAllocator::ensure_gpu_data(Buffer& buffer) { | ||
| auto* dbuf = static_cast<DoubleBuffer*>(buffer.ptr()); | ||
| if (dbuf->gpu_data() || dbuf->size() == 0) | ||
| return; | ||
| dbuf->set_gpu_data(device_.CreateBufferFromData( | ||
| dbuf->cpu_data(), dbuf->size(), betann::BufferUsage::Storage)); | ||
| } | ||
|
|
||
| WgpuAllocator& allocator() { | ||
| static WgpuAllocator allocator_; | ||
| return allocator_; | ||
| } | ||
|
|
||
| betann::Device& device(mlx::core::Device) { | ||
| static betann::Device device; | ||
| return device; | ||
| } | ||
|
|
||
| } // namespace webgpu | ||
|
|
||
| namespace metal { | ||
|
|
||
| size_t get_active_memory() { | ||
| return 0; | ||
| } | ||
| size_t get_peak_memory() { | ||
| return 0; | ||
| } | ||
| void reset_peak_memory() {} | ||
| size_t get_cache_memory() { | ||
| return 0; | ||
| } | ||
| size_t set_memory_limit(size_t, bool) { | ||
| return 0; | ||
| } | ||
| size_t set_cache_limit(size_t) { | ||
| return 0; | ||
| } | ||
| size_t set_wired_limit(size_t) { | ||
| return 0; | ||
| } | ||
|
|
||
| std::unordered_map<std::string, std::variant<std::string, size_t>> | ||
| device_info() { | ||
| throw std::runtime_error("[webgpu::device_info] Not implemented"); | ||
| }; | ||
|
|
||
| void clear_cache() {} | ||
|
|
||
| } // namespace metal | ||
|
|
||
| } // namespace mlx::core |
Oops, something went wrong.