/usr/local/lib64/python3.6/site-packages/torch/include/ATen/native/cuda
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/ATen/native/cuda/Resize.cuh (1919B)
#pragma once
#include
#include
#include
#include
namespace at { namespace native {
// These functions are called by native::resize_ as well as (legacy) THC resize.
// They are not in THC/THCTensor.cpp because the at namespace is easier
// to benchmark than THC; I can't get gbenchmark to call fns from THTensor.cpp
static inline void maybe_resize_storage_cuda(TensorImpl* self, uint64_t new_size) {
// It does not make sense to try to resize a storage
// to hold 0 elements, and this can break
// if storage_offset is positive but
// new_size is 0, so just bail in that case
// (same comment is in Resize.h)
if (new_size == 0) {
return;
}
if (!THTensor_getStoragePtr(self)) {
TORCH_CHECK(false, "Tensor: invalid null storage");
}
uint64_t new_size_bytes = (new_size + self->storage_offset()) * self->dtype().itemsize();
if (new_size_bytes > self->storage().nbytes()) {
THCStorage_resizeBytes(
globalContext().getTHCState(),
THTensor_getStoragePtr(self),
new_size_bytes
);
}
}
inline TensorImpl* resize_impl_cuda_(
TensorImpl* self,
IntArrayRef size,
c10::optional stride,
bool device_guard = true) {
if (self->sizes() == size && (!stride || self->strides() == stride)) {
return self;
}
// NB: We don't need to hold the device guard when calling from TH
cuda::OptionalCUDAGuard guard;
if (device_guard) {
guard.set_index(self->storage().device().index());
}
int64_t storage_size = 1;
if (stride) {
self->set_sizes_and_strides(size, *stride);
// NB: storage size can be different from numel.
storage_size = storage_size_for(size, *stride);
} else {
self->set_sizes_contiguous(size);
storage_size = self->numel();
}
maybe_resize_storage_cuda(self, storage_size);
return self;
}
}}