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

No copy command encoder (#986)

Browse Source 
* no copy command encoder

* up layer norm test tolerances
This commit is contained in:
Awni Hannun
2024-04-11 21:15:36 -07:00
committed by GitHub
parent 91eba8e485
commit ae18326533
3 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
mlx/backend/metal/device.cpp
View File

@@ -129,7 +129,7 @@ Device::~Device() {
b.second.second->release();
}
for (auto& e : encoder_map_) {
e.second->release();
(*e.second)->release();
}
for (auto& k : kernel_map_) {
k.second->release();
@@ -200,8 +200,8 @@ void Device::commit_command_buffer(int index) {
void Device::end_encoding(int index) {
auto eit = encoder_map_.find(index);
if (eit != encoder_map_.end()) {
eit->second->endEncoding();
eit->second->release();
(*eit->second)->endEncoding();
(*eit->second)->release();
encoder_map_.erase(eit);
}
}
@@ -214,9 +214,11 @@ CommandEncoder& Device::get_command_encoder(int index) {
cb->computeCommandEncoder(MTL::DispatchTypeConcurrent);
// Increment ref count so the buffer is not garbage collected
compute_encoder->retain();
eit = encoder_map_.emplace(index, CommandEncoder{compute_encoder}).first;
eit = encoder_map_
.emplace(index, std::make_unique<CommandEncoder>(compute_encoder))
.first;
}
return eit->second;
return *(eit->second);
}
void Device::register_library(