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

OpenReg: Split Allocator (#144843)

Browse Source 
Split the Allocator into HostAllocator and DeviceAllocator.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/144843
Approved by: https://github.com/albanD
This commit is contained in:
Zhenbin Lin 2025-01-17 03:38:13 +00:00 committed by PyTorch MergeBot
parent 43a00d73b3
commit adbbcd87d9
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
test/cpp_extensions/open_registration_extension/pytorch_openreg/_device_daemon.py
View File

@ -38,9 +38,13 @@ class Allocator:
del self.allocated[ptr]
return True
def is_allocated(self, ptr):
class HostAllocator(Allocator):
def is_pinned_ptr(self, ptr):
return ptr in self.allocated
class DeviceAllocator(Allocator):
def tensor_from_meta(self, meta):
def create_tensor_from_data_ptr(ptr, size):
storage = torch._C._construct_storage_from_data_pointer(
@ -109,7 +113,7 @@ class Driver:
# Allocated memory belongs to which device
self.memory_belong = {}
self.host_allocator = Allocator()
self.host_allocator = HostAllocator()
self.event_belong = {}
self.devices = []
@ -197,7 +201,7 @@ class Driver:
@register(registry)
def isPinnedPtr(self, ptr):
return self.host_allocator.is_allocated(ptr)
return self.host_allocator.is_pinned_ptr(ptr)
@register(registry)
def hostMalloc(self, size):
@ -274,7 +278,7 @@ class Driver:
class _Executor:
def __init__(self, id):
self.id = id
self.allocator = Allocator()
self.allocator = DeviceAllocator()
self.stream = 0
self.event_incr_id = 0
self.events = {}