1 #pragma once 2 3 #include <c10/core/Storage.h> 4 #include <c10/core/StorageImpl.h> 5 #include <c10/util/intrusive_ptr.h> 6 7 namespace at { 8 9 class TensorBase; 10 11 // Here we define a series of utils to create/manipulate ATen backed 12 // c10 storage implementations. 13 14 /** 15 * Create a new shared memory storage impl managed by file descriptor 16 * 17 * @param size size in bytes 18 */ 19 C10_EXPORT c10::intrusive_ptr<c10::StorageImpl> new_shm_fd_storage(size_t size); 20 21 /** 22 * Copy src to dst 23 * Caller must guarantee the validness of the storage objects 24 * during the entire copy process, esp. when it's async. 25 * 26 * This can probably live in c10 namespace later if needed, 27 * but for now keep it in at to keep implementation simple. 28 * 29 * @param dst dst tensor 30 * @param src src tensor 31 * @param non_blocking (default false) whether this operation blocks caller 32 */ 33 C10_EXPORT void storage_copy( 34 c10::Storage& dst, 35 const c10::Storage& src, 36 bool non_blocking = false); 37 38 /** 39 * In place change the storage to shm based. 40 * 41 * This is only applicable to CPU tensors not already shared. 42 * Otherwise, it's a no op to mirror the THP tensor behavior: 43 * https://pytorch.org/docs/stable/generated/torch.Tensor.share_memory_.html 44 * 45 * @param t a tensor 46 */ 47 C10_EXPORT void share_memory_(TensorBase& t); 48 49 } // namespace at 50