OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Generalize support of background thread in pinned allocator (#160505)

Browse Source 
# Motivation
https://github.com/pytorch/pytorch/pull/135524 only introduces the support of background thread for CUDA, this PR intends to support it for other backend such as XPU as well.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/160505
Approved by: https://github.com/albanD
This commit is contained in:
Yu, Guangye 2025-08-13 14:18:53 +00:00 committed by PyTorch MergeBot
parent af3cabc55d
commit 8cfaf51d4e
3 changed files with 14 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
aten/src/ATen/core/CachingHostAllocator.h
View File

@ -1,6 +1,7 @@
#pragma once
#include <c10/core/Allocator.h>
#include <c10/core/AllocatorConfig.h>
#include <c10/core/Stream.h>
#include <c10/core/thread_pool.h>
#include <c10/util/flat_hash_map.h>
@ -351,7 +352,8 @@ struct CachingHostAllocatorImpl {
}
virtual bool pinned_use_background_threads() {
return false;
return c10::CachingAllocator::AcceleratorAllocatorConfig::
pinned_use_background_threads();
}
virtual void copy_data(void* dest [[maybe_unused]], const void* src [[maybe_unused]], std::size_t count [[maybe_unused]]) const {

5
aten/src/ATen/cuda/CachingHostAllocator.cpp
View File

@ -161,11 +161,6 @@ struct CUDACachingHostAllocatorImpl
return true;
}
bool pinned_use_background_threads() override {
return c10::CachingAllocator::AcceleratorAllocatorConfig::
pinned_use_background_threads();
}
EventPool::Event create_event_internal(DeviceIndex idx) {
// Leak the event pool to avoid shutdown issue.
static auto* event_pool = new EventPool();

11
test/test_xpu.py
View File

@ -607,6 +607,17 @@ if __name__ == "__main__":
z[0] = z[0] + 1.0
self.assertEqual(z, x)
def test_background_thread_for_pin_memory(self):
# Just ensure no crash
torch._C._accelerator_setAllocatorSettings("pinned_use_background_threads:True")
cpu_tensor = torch.randn(100)
pin_tensor = cpu_tensor.pin_memory()
xpu_tensor = pin_tensor.to(device="xpu", non_blocking=True)
torch.xpu.synchronize()
del pin_tensor
gc.collect()
self.assertEqual(xpu_tensor.cpu(), cpu_tensor)
instantiate_device_type_tests(TestXpu, globals(), only_for="xpu", allow_xpu=True)