| import torch |
| import os |
| import weakref |
| import multiprocessing |
| from multiprocessing.reduction import ForkingPickler |
| import sys |
| try: |
| # Early load resource_sharer to prevent a partially initialized instance |
| # from being inherited in a forked child process. The reduce_storage method |
| # requires this module indirectly through DupFd(). The built-in mp.Queue |
| # class pickles arguments in a background thread which may overlap with the |
| # fork. |
| import multiprocessing.resource_sharer |
| except ImportError: |
| pass |
| |
| |
| class StorageRef(object): |
| # An object with a cdata field which may be set to None. We subclass object |
| # instead of using a dict() to support weak references. |
| |
| def __init__(self, ptr): |
| self.cdata = ptr |
| |
| def __del__(self): |
| torch.Storage._free_weak_ref(self.cdata) |
| |
| |
| # mapping from handles to StorageRef objects |
| shared_cache = weakref.WeakValueDictionary() |
| |
| |
| def rebuild_event(handle): |
| return torch.cuda.Event(_handle=handle) |
| |
| |
| def reduce_event(event): |
| return (rebuild_event, (event.ipc_handle(),)) |
| |
| |
| def rebuild_tensor(cls, storage, metadata): |
| storage_offset, size, stride = metadata |
| return torch._utils._rebuild_tensor(storage, storage_offset, size, stride) |
| |
| |
| def rebuild_cuda_tensor(tensor_cls, tensor_size, tensor_stride, tensor_offset, |
| storage_cls, storage_device, storage_handle, storage_size): |
| |
| storage = storage_from_cache(storage_cls, storage_handle) |
| if storage is None: |
| torch.cuda._lazy_init() |
| storage = storage_cls._new_shared_cuda(storage_device, storage_handle, storage_size) |
| shared_cache[storage_handle] = storage._weak_ref(StorageRef) |
| |
| return torch._utils._rebuild_tensor(storage, tensor_offset, tensor_size, tensor_stride) |
| |
| |
| def reduce_tensor(tensor): |
| storage = tensor.storage() |
| |
| # Note [CUDA IPC and the caching allocator] |
| # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ |
| # When you send a CUDA tensor over IPC, you might expect that you will |
| # get out the same storage from the other end. However, the CUDA caching |
| # allocator makes it difficult to preserve this invariant. Consider |
| # the following situation: a tensor of size 0x100 points to offset 0x20 of |
| # a storage at 0xA100 of size 0x100. (For simplicity, all of these |
| # sizes are given in bytes). HOWEVER, with the caching allocator, this storage |
| # might be part of a larger cudaMalloc allocation 0xA000 of size 0x4000. |
| # |
| # When we want to send this CUDA tensor over IPC, we must send the |
| # *entire* cudaMalloc allocation, i.e., the 0xA000 region, not just |
| # the storage 0xA100 (because that is what CUDA supports). So, on the |
| # other end, there simply isn't any way to say, "Wait, you gave me |
| # a bigger region (0xA000) than the one I wanted (0xA100)"; we have |
| # to just make a storage for the entire caching allocator block. |
| # |
| # This is fine, because all we need to do is just adjust the offset |
| # on the tensor itself: instead of: |
| # |
| # Tensor(size=0x100, offset=0x020, storage=Storage(data=0xA100, size=0x0100)) |
| # |
| # we have |
| # |
| # Tensor(size=0x100, offset=0x120, storage=Storage(data=0xA000, size=0x4000)) |
| # |
| # This strategy has a few implications: |
| # |
| # 1. When we serialize a CUDA tensor for IPC, we have to do it all in one |
| # go (non-compositionally), instead of first serializing storage, and |
| # then serializing tensor. This is because the base address of the |
| # storage allocation affects what offset we write into the tensor. |
| # |
| # 2. We MUST NOT let the new IPC tensor be resizable. Originally, a resize |
| # of the storage beyond 0x100 would merely have caused us to do a |
| # reallocation. You don't really want to do this, but if you did, |
| # all that would happen is that you would lose IPC sharing. But if |
| # you do this in the new world, we will happily let you write out of |
| # bounds of your "allocation", clobbering unrelated data in the cached |
| # allocator block. BAD! |
| # |
| # By the way, in old versions of PyTorch, we supported this situation |
| # natively using a "storage view", which permitted multiple storages to be |
| # views on each other. But this was the *only* use of storage views, so we |
| # eliminated it so that we could just use tensor views to implement the same |
| # thing. |
| # |
| if storage.is_cuda: |
| (device, handle, storage_size, storage_offset) = storage._share_cuda_() |
| tensor_offset = tensor.storage_offset() |
| |
| # WARNING! This call to _weak_ref could lead to O(n) deleter |
| # behavior, if you repeatedly call it on the same Storage (all |
| # other sites are guarded by shared_cache; maybe this site |
| # should be too?) |
| shared_cache[handle] = storage._weak_ref(StorageRef) |
| |
| return (rebuild_cuda_tensor, |
| (type(tensor), |
| tensor.size(), |
| tensor.stride(), |
| tensor_offset + storage_offset, |
| type(storage), |
| device, |
| handle, |
| storage_size)) |
| |
| metadata = (tensor.storage_offset(), tensor.size(), tensor.stride()) |
| return (rebuild_tensor, (type(tensor), storage, metadata)) |
| |
| |
| def fd_id(fd): |
| # Returns a tuple which uniquely identifies a file descriptor. In Mac OS, |
| # this doesn't work with shared memory handles, which is why we don't |
| # support the "file_descriptor" sharing method on that platform. |
| stat = os.fstat(fd) |
| return (stat.st_ino, stat.st_dev) |
| |
| |
| def storage_from_cache(cls, key): |
| storage_ref = shared_cache.get(key) |
| if storage_ref is None: |
| return None |
| return cls._new_with_weak_ptr(storage_ref) |
| |
| |
| def rebuild_storage_fd(cls, df, size): |
| if sys.version_info[0] == 2: |
| fd = multiprocessing.reduction.rebuild_handle(df) |
| else: |
| fd = df.detach() |
| try: |
| storage = storage_from_cache(cls, fd_id(fd)) |
| if storage is not None: |
| return storage |
| storage = cls._new_shared_fd(fd, size) |
| shared_cache[fd_id(fd)] = storage._weak_ref(StorageRef) |
| return storage |
| finally: |
| os.close(fd) |
| |
| |
| def rebuild_storage_filename(cls, manager, handle, size): |
| storage = storage_from_cache(cls, handle) |
| if storage is not None: |
| return storage._shared_decref() |
| storage = cls._new_shared_filename(manager, handle, size) |
| shared_cache[handle] = storage._weak_ref(StorageRef) |
| return storage._shared_decref() |
| |
| |
| def rebuild_storage_empty(cls): |
| return cls() |
| |
| |
| def reduce_storage(storage): |
| from . import get_sharing_strategy |
| if storage.is_cuda: |
| raise RuntimeError("Cannot pickle CUDA storage; try pickling a CUDA tensor instead") |
| elif get_sharing_strategy() == 'file_system': |
| metadata = storage._share_filename_() |
| cache_key = metadata[1] |
| rebuild = rebuild_storage_filename |
| storage._shared_incref() |
| elif storage.size() == 0: |
| # This is special cased because Empty tensors |
| # (with size 0) cannot be mmapped. |
| return (rebuild_storage_empty, (type(storage),)) |
| else: |
| fd, size = storage._share_fd_() |
| if sys.version_info[0] == 2: |
| df = multiprocessing.reduction.reduce_handle(fd) |
| else: |
| df = multiprocessing.reduction.DupFd(fd) |
| cache_key = fd_id(fd) |
| metadata = (df, size) |
| rebuild = rebuild_storage_fd |
| |
| # WARNING! This call to _weak_ref could lead to O(n) deleter |
| # behavior, if you repeatedly call it on the same Storage (all |
| # other sites are guarded by shared_cache; maybe this site |
| # should be too?) |
| shared_cache[cache_key] = storage._weak_ref(StorageRef) |
| return (rebuild, (type(storage),) + metadata) |
| |
| |
| def init_reductions(): |
| ForkingPickler.register(torch.cuda.Event, reduce_event) |
| |
| for t in torch._storage_classes: |
| ForkingPickler.register(t, reduce_storage) |
| |
| for t in torch._tensor_classes: |
| ForkingPickler.register(t, reduce_tensor) |
| |
| # TODO: Maybe this should be in tensor_classes? :) |
| ForkingPickler.register(torch.Tensor, reduce_tensor) |
