pytorch之empty源代码

api at::empty 初始化一个固定大小的Tensor

例如

torch::Tensor tensor = at::empty({3,3,4});

创建一个3*3*4的三维Tensor

在cpu上的实现代码在aten/src/ATen/native/TensorFactories.cpp中

Tensor empty_cpu(IntArrayRef size, const TensorOptions& options_, c10::optional<c10::MemoryFormat> optional_memory_format)

在cuda上的实现代码在aten/src/ATen/native/cuda/TensorFactories.cu中

Tensor empty_cuda(IntArrayRef size, const TensorOptions& options, c10::optional<MemoryFormat> optional_memory_format)

以empty_cpu为例函数主要部分可为三个步骤

1.获得分配器

c10::Allocator* allocator;

  if (options.pinned_memory()) {

    allocator = detail::getCUDAHooks().getPinnedMemoryAllocator();

  } else {

    allocator = at::getCPUAllocator();

  }

2.分配器分配空间后构建存储对象

  int64_t nelements = prod_intlist(size);

  auto dtype = options.dtype();

  int64_t size_bytes = nelements * dtype.itemsize();//类型占据字节数*单元个数

  auto storage_impl = c10::make_intrusive<StorageImpl>(

      c10::StorageImpl::use_byte_size_t(),

      size_bytes,

      allocator->allocate(size_bytes),

      allocator,

      /*resizeable=*/true);

再看allocator->allocate(size_bytes)内存分配，cpu分配器类DefaultCPUAllocator源代码在c10/core/CPUAllocator.cpp中

  at::DataPtr allocate(size_t nbytes) const override {

    void* data = alloc_cpu(nbytes);

    profiledCPUMemoryReporter().New(data, nbytes);

    return {data, data, &ReportAndDelete, at::Device(at::DeviceType::CPU)};

  }

其中alloc_cpu中调用的是_aligned_malloc分配函数

3.根据存储对象构建Tensor

auto tensor = detail::make_tensor<TensorImpl>(

std::move(storage_impl), at::DispatchKey::CPU, dtype);