How is tensor storage size determined

ashari4 · April 15, 2022, 11:13pm

In what cases does the underlying storage size not equal t.numel() * t.element_size()?

A known case is a tensor with all strides equal to 0.

>>> t = torch.empty_strided(size=(3,3),stride=(0,0),dtype=torch.float32)
>>> t.storage().nbytes() 
4

ashari4 · April 16, 2022, 1:43am

looks like the storage size of tensors created via empty_strided is given by “1 bigger than the offset of the last element according to stride”

github.com

pytorch/pytorch/blob/051802e4f4fc6c03b653dd4552998354730fdef2/aten/src/ATen/EmptyTensor.cpp#L60

      
        
            }
            
            
size_t computeStorageNbytes(
                IntArrayRef sizes,
                IntArrayRef strides,
                size_t itemsize_bytes,
                size_t storage_offset
              ) {
              // Ignore overflow checks on mobile
            #ifndef C10_MOBILE
              // size of the underlying storage is 1 bigger than the offset
              // of the last element according to stride
              uint64_t size = storage_offset + 1;
              bool overflowed = false;
              for (const auto i : c10::irange(sizes.size())) {
                if (sizes[i] == 0) {
                  return 0;
                }
            
            
    uint64_t strided_size;
                overflowed |= c10::mul_overflows(strides[i], sizes[i] - 1, &strided_size);

And the size of a contiguous tensor and views of such tensors is given by the product along the dimensions

github.com

pytorch/pytorch/blob/051802e4f4fc6c03b653dd4552998354730fdef2/aten/src/ATen/EmptyTensor.cpp#L39

      
        
            }  // namespace (anonymous)
            
            
size_t computeStorageNbytesContiguous(
                IntArrayRef sizes,
                size_t itemsize_bytes,
                size_t storage_offset
              ) {
              // Ignore overflow checks on mobile
            #ifndef C10_MOBILE
              uint64_t size = 1;
              bool overflowed = c10::safe_multiplies_u64(sizes, &size);
              overflowed |= c10::add_overflows(size, storage_offset, &size);
              overflowed |= c10::mul_overflows(size, itemsize_bytes, &size);
              overflowed |= size > storage_max();
              TORCH_CHECK(!overflowed,
                          "Storage size calculation overflowed with sizes=", sizes);
              return static_cast<size_t>(size);
            #else
              const auto numel = c10::multiply_integers(sizes);
              return itemsize_bytes * (storage_offset + numel);
            #endif