PyTorch源码浅析(1)：THTensor

PyTorch中Tensor的存储和表示分开，多个THTensor可能共享一个THStorage，每个THTensor可能拥有不同的view（e.g. size, stride）。这样设计的好处是，有时看起来不一样的数据底层是共享的，比如矩阵与矩阵的转置、二维矩阵与二维矩阵变成一维时的矩阵。这部分的主要实现在pytorch/aten文件夹中，这里面既实现了底层的Tensor操作库，也封装了名为 ATen 的 C++11接口。

aten/src里有几个重要的文件夹，它们实现的内容如下：

src
├── ATen        # Tensor操作的C++11接口
├── TH          # CPU Tensor 的底层实现（本篇内容）
├── THC         # GPU Tensor (CUDA) 的底层实现（在下一篇讲）
├── THCUNN      # CUDA版底层神经网络实现(下下篇讲)
└── THNN        # CPU版底层神经网络实现(下下篇讲)
这篇讲的主要代码也都在TH文件夹内。

THTensor & THStorage

TH里面的核心类型就是THTensor和THStorage了，前者是Tensor的view，后者是Tensor数据的存储地址。由于Tensor的数据类型可以是多种多样的，而每种类型的API都一致，所以需要用到范型来减少重复代码，TH中是使用宏实现范型功能（因为Torch一开始是用C实现的），不了解的读者可以先参考一下这篇知乎专栏:https://zhuanlan.zhihu.com/p/34496542

THTensor的定义在aten/src/TH/generic/THTensor.h中：

#define THTensor at::TensorImpl#define THFloatTensor THTensor
#define THDoubleTensor THTensor
#define THHalfTensor THTensor
#define THByteTensor THTensor
#define THCharTensor THTensor
#define THShortTensor THTensor
#define THIntTensor THTensor
#define THLongTensor THTensor
#define THBoolTensor THTensor

同样的，THStorage的定义在Aten/src/TH/generic/THStorage.h：

#define THStorage at::StorageImpl#define THFloatStorage THStorage
#define THDoubleStorage THStorage
#define THHalfStorage THStorage
#define THByteStorage THStorage
#define THCharStorage THStorage
#define THShortStorage THStorage
#define THIntStorage THStorage
#define THLongStorage THStorage
#define THBoolStorage THStorage

在THTensor.h中有很多类似这样的API声明：

TH_API THStorage* THTensor_(storage)(const THTensor *self);

其中，THTensor_的宏定义为

#define THTensor_(NAME)   TH_CONCAT_4(TH,Real,Tensor_,NAME)

上面的TH_CONCAT_4宏就是把它的四个参数连接起来，其中Real会被定义为实际类型（Float, Bool等），所以上面的API会被展开成：

at::StorageImpl THFloatTensor_storage(const at::TensorImpl *self);
at::StorageImpl THBoolTensor_storage(const at::TensorImpl *self);
at::StorageImpl THLongTensor_storage(const at::TensorImpl *self);
...

在这些API中还会用到scalar_t类型，这个也是一个宏，特化的时候会传入具体的类型（如用来确保THFloatTensor里的StorageImpl存储的float类型）。

不难发现，所有的THTensor类型最终都会替换成at::TensorImpl，所有的THStorage类型也都会替换成at::StorageImpl，前者的实现在c10/core/TensorImpl.h中，后者的实现在c10/core/StorageImpl.h中。这两个类型的实现在c10中，也就是说Tensor类型的实现转移到了c10中，但API的实现依然在TH中。

TensorImpl

接着具体来看一下TensorImpl的实现，首先来看一下它的声明和属性：

struct C10_API TensorImpl : public c10::intrusive_ptr_target {protected:// 指向实际数据存储位置，也就是指向StorageImplStorage storage_;// 用于自动微分，只对Variable适用std::unique_ptr<c10::AutogradMetaInterface> autograd_meta_ = nullptr;SmallVector<int64_t,5> sizes_;      // Tensor的实际大小SmallVector<int64_t,5> strides_;    // 各个维度直接的间隔int64_t storage_offset_ = 0;int64_t numel_ = 1; // Tensor中元素个数，也就是sizes_数组中元素的乘积caffe2::TypeMeta data_type_;        // 数据类型TensorTypeId type_id_;bool is_contiguous_ = true;bool is_variable_ = false;bool is_wrapped_number_ = false;bool allow_tensor_metadata_change_ = true;bool reserved_ = false;
}

TensorImpl继承自intrusive_ptr_target，后者是为了通过使用intrusive_ptr实现引用计数而设计的基类，需要实现引用计数的类只需继承它即可。TensorImpl中的每个成员是干什么的基本都有注释，其中strides_是用来实现内存寻址的，即某个ijk脚标对应的内存地址为：

storage_offset_ + i * strides_[0] + j * strides_[1] + k * strides_[2]

正常情况下用sizes_代替strides_可以实现同样的功能，但是有的Tensor是由较大的Tensor分割而来，维度之间的间隔不是sizes_，所以需要用另一个数组strides_存储维度间隔。

TensorImpl中的方法较多，就不一一列出了，实现了对Tensor（和Variable）的基本操作，Aten中的API也是基于这些基本操作实现的。

Variable与Tensor的合并

在早期的PyTorch版本中，Variable与Tensor是不同的类，Variable用来保存需要计算梯度的Tensor，但Variable的实现并不美好：一方面Variable::Impl是Tensor的子类，而它的成员里又拥有一个Tensor（存储数据），这违反了里氏替换原则，而且让Tensor的实现变得很复杂。而在现版本中，已经把Variable变成Tensor了，把一些Variable特有的方法（e.g.requires_grad）移到了TensorImpl里。

StorageImpl

StorageImpl的声明和属性如下：

struct C10_API StorageImpl final : public c10::intrusive_ptr_target {private:caffe2::TypeMeta data_type_;  // 数据类型DataPtr data_ptr_;            // 指向存储数据的内存块int64_t numel_;               // 数据总数bool resizable_;Allocator* allocator_;        // 内存分配器
}

其中，caffe2::TypeMeta data_type_存储了数据类型信息，包括：类型id、大小、类型名字等；DataPtr其实就是 unique_ptr，但指针类型固定 void*。

分配内存

在Allocator.cpp中定义了全局变量allocator_array来存储所有的 allocator，每个设备类型对应一个：

C10_API at::Allocator* allocator_array[at::COMPILE_TIME_MAX_DEVICE_TYPES];void SetAllocator(at::DeviceType t, at::Allocator* alloc) {allocator_array[static_cast<int>(t)] = alloc;
}at::Allocator* GetAllocator(const at::DeviceType& t) {auto* alloc = allocator_array[static_cast<int>(t)];AT_ASSERTM(alloc, "Allocator for ", t, " is not set.");return alloc;
}

同时配备了SetAllocator和GetAllocator来设置和获取相应的分配器。

Allocator是一个虚基类，它的声明如下：

struct C10_API Allocator {virtual ~Allocator() = default;virtual DataPtr allocate(size_t n) const = 0;// 重载这个函数很关键，用来释放申请到的内存virtual DeleterFnPtr raw_deleter() const {return nullptr;}void* raw_allocate(size_t n) {auto dptr = allocate(n);AT_ASSERT(dptr.get() == dptr.get_context());return dptr.release_context();}void raw_deallocate(void* ptr) {auto d = raw_deleter();AT_ASSERT(d);d(ptr);}
};

这个分配器有两种使用方法，第一种就是直接调用raw_allocate和raw_deallocate分配和释放内存。第二种方法是调用Allocator::allocate，这个方法将返回一个DataPtr类型的指针，也就是 unique_ptr，由于deleter存储在指针中，在释放指针的时候会释放相应的内存。不过两种方法的正确性都依赖于Allocator::raw_deleter能正确返回相应的释放器（deleter），否则两种方法都不能正确释放内存。这里需要注意的是，释放器（deleter）未必就是C库函数free：根据操作系统的不同，也可能是_aligned_free；根据设备的不同也可能是其他函数。

下面是在CPU上内存分配的具体实现：

struct C10_API DefaultCPUAllocator final : at::Allocator {...at::DataPtr allocate(size_t nbytes) const override {void* data = alloc_cpu(nbytes);if (FLAGS_caffe2_report_cpu_memory_usage && nbytes > 0) {getMemoryAllocationReporter().New(data, nbytes);return {data, data, &ReportAndDelete, at::Device(at::DeviceType::CPU)};}// 这四个参数用来构造一个DataPtrreturn {data, data, &free_cpu, at::Device(at::DeviceType::CPU)};}at::DeleterFnPtr raw_deleter() const override {if (FLAGS_caffe2_report_cpu_memory_usage) {return &ReportAndDelete;}return &free_cpu;}
};void* alloc_cpu(size_t nbytes) {if (nbytes == 0) {return nullptr;}void* data;// 分配64字节对齐的连续内存块
#ifdef __ANDROID__data = memalign(gAlignment, nbytes); // gAlignment = 64
#elif defined(_MSC_VER)data = _aligned_malloc(nbytes, gAlignment);
#elseCAFFE_ENFORCE_EQ(posix_memalign(&data, gAlignment, nbytes), 0);
#endif// 移动数据到线程的NUMA节点中NUMAMove(data, nbytes, GetCurrentNUMANode());// 填充内存if (FLAGS_caffe2_cpu_allocator_do_zero_fill) {memset(data, 0, nbytes);} else if (FLAGS_caffe2_cpu_allocator_do_junk_fill) {memset_junk(data, nbytes);}return data;
}void free_cpu(void* data) {#ifdef _MSC_VER_aligned_free(data);
#elsefree(data);
#endif
}

分配时使用memalign/_aligned_malloc/posix_memalign函数确保内存是64字节对齐的。

智能指针 Intrusive_ptr

PyTorch中使用intrusive_ptr来管理THTensor和THStorage的引用计数，其中引用分为强引用和弱引用（弱引用为了解决循环引用问题），对应的类名 intrusive_ptr和weak_intrusive_ptr。由于弱引用和强引用的实现类似，为了简单起见，我把弱引用的代码都去掉了，简化intrusive_ptr的实现如下：

class intrusive_ptr_target {mutable std::atomic<size_t> refcount_;// 声明友元类使得只能指针可以访问 refcount_template <typename T>friend class intrusive_ptr;protected:// 隐藏析构函数，防止直接析构对象virtual ~intrusive_ptr_target() {}constexpr intrusive_ptr_target() noexcept : refcount_(0) {}private:// 在摧毁对象时会调用次函数释放资源virtual void release_resources() {}
};template <class TTarget>
class intrusive_ptr final {private:TTarget* target_;// 增加引用计数void retain_() {if (target_ != nullptr) {size_t new_refcount = ++target_->refcount_;AT_ASSERTM(new_refcount != 1,"intrusive_ptr:Cannot increase refcount after it""reached zero.");}}// 销毁智能指针，减少引用计数，当引用计数为0时摧毁对象void reset_() noexcept {if (target_ != nullptr && --target_->refcount_ == 0) {target_->release_resources();delete target_;}target_ = nullptr;}// 隐藏通过普通指针初始化的构造函数（只能通过make_intrusive调用）explicit intrusive_ptr(TTarget* target) noexcept : target_(target) {}public:intrusive_ptr() noexcept : intrusive_ptr(nullptr) {}// 通过其他智能指针初始化，需增加引用计数intrusive_ptr(const intrusive_ptr& rhs) : target_(rhs.target_) {retain_(); }~intrusive_ptr() noexcept { reset_(); }TTarget* get() const noexcept { return target_; }// 重载运算符使其能当作正常指针使用const TTarget& operator*() const noexcept { return *target_; }TTarget& operator*() noexcept { return *target_; }const TTarget* operator->() const noexcept { return target_; }TTarget* operator->() noexcept { return target_; }operator bool() const noexcept { return target_ != nullptr; }// 获取当前的引用计数size_t use_count() const noexcept {if (target_ == nullptr) {return 0;}return target_->refcount_.load();}// 把普通指针转换为智能指针（增加引用计数）template <class... Args>static intrusive_ptr make(Args&&... args) {auto result = intrusive_ptr(new TTarget(std::forward<Args>(args)...));++result.target_->refcount_;return result;}// 把智能指针转换为普通指针TTarget* release() noexcept {TTarget* result = target_;target_ = nullptr;return result;}// 把普通指针转换为智能指针（不增加引用计数）static intrusive_ptr reclaim(TTarget* owning_ptr) {AT_ASSERTM(owning_ptr == nullptr || owning_ptr->refcount_.load() > 0,"intrusive_ptr: Can only intrusive_ptr::reclaim() owning pointers that were created using intrusive_ptr::release().");return intrusive_ptr(owning_ptr);}
};// 用于构造智能指针
template <class TTarget, class... Args>
inline intrusive_ptr<TTarget> make_intrusive(Args&&... args) {return intrusive_ptr<TTarget>::make(std::forward<Args>(args)...);
}

intrusive_ptr_target是被引用对象的父类，TensorImpl和StorageImpl都继承自它，主要工作是声明了一个引用计数refcount_，并且把智能指针类声明为友元类，这样在intrusive_ptr里面就可以直接操作refcount_了。

再来看intrusive_ptr的实现，它有一个私有成员TTarget* target_，是被引用对象的普通指针；还有两个私有方法retain_和reset_，分别用来增加和减少引用计数。需要注意的是：通过普通指针初始化intrusive_ptr的构造函数也是私有的，不能被外部调用，只能通过静态函数make来调用（详见下文）；在通过其他智能指针初始化的时候需要增加引用计数。

最后还有两个方法release和reclaim，它们实现了普通指针和智能指针间的相互转换，用于C风格的API中（在Aten中经常用到）。除此之外，intrusive_ptr里还重载了许多运算符，让它可以像普通指针一样使用，还有许多其他方法就不一一介绍了。

创建Tensor

THTensor *THTensor_(new)(void)
{return c10::make_intrusive<at::TensorImpl>(// 下面三个参数将通过intrusive_ptr::make传给TensorImpl的构造函数，然后// 通过TensorImpl的指针构造intrusive_ptrc10::intrusive_ptr<at::StorageImpl>::reclaim(THStorage_(new)()),                                     // Storage&& storageat::CPUTensorId(),  // TensorTypeId type_idfalse                                // bool is_variable).release(); // release返回普通指针，TensorImpl*
}THStorage* THStorage_(new)(void)
{return THStorage_new(caffe2::TypeMeta::Make<scalar_t>());
}THStorage* THStorage_new(caffe2::TypeMeta data_type) {THStorage* storage = c10::make_intrusive<at::StorageImpl>(// 同理下面四个参数将通过intrusive_ptr::make传给StorageImpl的构造函// 数，然后通过StorageImpl的指针构造intrusive_ptrdata_type,0,getTHDefaultAllocator(),true).release();return storage;
}

上面是新建一个tensor的过程，通过c10::make_instrusive构造智能指针，然后调用release返回普通指针。传入的三个参数：storage智能指针，tensortype，is_variable会被转发到intrusive_ptr::make函数中，然后用这三个参数构造一个TensorImpl对象：

TensorImpl(Storage&& storage, TensorTypeId type_id, bool is_variable);

再用该对象的指针初始化智能指针：

intrusive_ptr(TTarget* target);

同时增加引用计数，最后make_intrusive返回智能指针。THStorage的构造过程同理。

Tensor Apply & Dynamic Dispatch
TH中的Tensor Apply就相当于map函数，把一个函数应用到tensor的每个数据中。举个例子来看一下THTensor_(cadd)的具体实现（简化过的）：

void THTensor_(cadd)(THTensor *r_, THTensor *t, scalar_t value, THTensor *src)
{THTensor_(resizeAs)(r_, t);int64_t r_Size = THTensor_(nElement)(r_);int64_t srcSize = THTensor_(nElement)(src);int r_Contig = THTensor_(isContiguous)(r_);int tContig = THTensor_(isContiguous)(t);int srcContig = THTensor_(isContiguous)(src);int serial_path = 0;if (srcSize == r_Size){if (r_Contig && tContig && srcContig) {TH_TENSOR_APPLY3_CONTIG(scalar_t, r_, scalar_t, t, scalar_t, src, THVector_(cadd)(r__data, t_data, src_data, value, r__len););} else {serial_path = 1;}} else {serial_path = 1;}if (serial_path) {TH_TENSOR_APPLY3(scalar_t, r_, scalar_t, t, scalar_t, src, *r__data = *t_data + value * *src_data;);}
}

这个函数实现的功能很简单，就是*r_ = *t + value * *src，把src的值乘以value然后和t相加，最后赋值给r_，其中r_, t, src都是THTensor。函数首先获取元素个数和是否连续等信息，如果连续的话调用TH_TENSOR_APPLY3_CONTIG来处理，否则调用TH_TENSOR_APPLY3处理。后者太复杂了，我们主要看前者的实现：

#ifdef _OPENMP
// 多线程加速
#define TH_TENSOR_APPLY3_CONTIG(TYPE1, TENSOR1, TYPE2, TENSOR2, TYPE3, TENSOR3, CODE) \
{ \ int inOmp = omp_in_parallel(); \ptrdiff_t TH_TENSOR_size = THTensor_(nElement)(TENSOR1); \ # 启动 OpenMP 多线程PRAGMA(omp parallel if ((TH_TENSOR_size > TH_OMP_OVERHEAD_THRESHOLD) && (!inOmp))) \{ \size_t num_threads = omp_get_num_threads(); \// 获取线程数size_t tid = omp_get_thread_num(); \// 计算开始和结尾ptrdiff_t TH_TENSOR_offset = tid*(TH_TENSOR_size/num_threads);\ptrdiff_t TH_TENSOR_end =(tid==num_threads-1)?TH_TENSOR_size: \TH_TENSOR_offset + TH_TENSOR_size / num_threads; \// 每个线程负责 TH_TENSOR_offset 到 TH_TENSOR_end 之间的数据ptrdiff_t TENSOR1##_len = TH_TENSOR_end - TH_TENSOR_offset; \// 获取Tensor的数据TYPE1 *TENSOR1##_data = TENSOR1->data<scalar_t>() + TH_TENSOR_offset; \TYPE2 *TENSOR2##_data = TENSOR2->data<scalar_t>() + TH_TENSOR_offset; \TYPE3 *TENSOR3##_data = TENSOR3->data<scalar_t>() + TH_TENSOR_offset; \CODE \} \
}
#else
// 普通实现
#define TH_TENSOR_APPLY3_CONTIG(TYPE1, TENSOR1, TYPE2, TENSOR2, TYPE3, TENSOR3, CODE) \
{ \// 获取Tensor的数据TYPE1 *TENSOR1##_data = TENSOR1->data<scalar_t>(); \TYPE2 *TENSOR2##_data = TENSOR2->data<scalar_t>(); \TYPE3 *TENSOR3##_data = TENSOR3->data<scalar_t>(); \ptrdiff_t TENSOR1##_len = THTensor_(nElement)(TENSOR1); \// 执行传入的代码CODE \
}
#endif

上半部分的实现为OPENMP的多线程加速版，下面是普通版。这个宏接收7个参数：三个Tensor和对应类型，还有要执行的操作。THTensor_(cadd)中传入的操作是THVector_(cadd)(r__data, t_data, src_data, value, r__len);，它的定义为：

// THVectorDispatch.cpp
void THVector_(cadd)(scalar_t *z, const scalar_t *x, const scalar_t *y, const scalar_t c, const ptrdiff_t n) {THVector_(cadd_DISPATCHPTR)(z, x, y, c, n);
}
// 动态派发函数指针
static void (*THVector_(cadd_DISPATCHPTR))(scalar_t *, const scalar_t *, const scalar_t *, const scalar_t, const ptrdiff_t) = &THVector_(cadd_DEFAULT);
// 对各种向量化指令的支持
static FunctionDescription THVector_(cadd_DISPATCHTABLE)[] = {#if defined(__NEON__)#if defined(TH_REAL_IS_FLOAT)FUNCTION_IMPL(THVector_(cadd_NEON), SIMDExtension_NEON),#endif#endif#if defined(USE_AVX2)#if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT)FUNCTION_IMPL(THVector_(cadd_AVX2), SIMDExtension_AVX2),#endif#endif#if defined(USE_AVX)#if defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_FLOAT)FUNCTION_IMPL(THVector_(cadd_AVX), SIMDExtension_AVX),#endif#endifFUNCTION_IMPL(THVector_(cadd_DEFAULT), SIMDExtension_DEFAULT)
};// THVectorDefault.cpp
void THVector_(cadd_DEFAULT)(scalar_t *z, const scalar_t *x, const scalar_t *y, const scalar_t c, const ptrdiff_t n)
{ptrdiff_t i = 0;for(; i<n-4; i+=4){z[i] = x[i] + c * y[i];z[i+1] = x[i+1] + c * y[i+1];z[i+2] = x[i+2] + c * y[i+2];z[i+3] = x[i+3] + c * y[i+3];}for(; i<n; i++)z[i] = x[i] + c * y[i];
}

这里涉及到对SIMD向量化指令的支持和动态派发，可以看到THVector_(cadd)函数里调用的是THVector_(cadd_DISPATCHPTR)，而它是一个函数指针，默认指向THVector_(cadd_DEFAULT)，这个是不支持向量化指令的实现。代码中间部分的数组FunctionDescription THVector_(cadd_DISPATCHTABLE)[]记录各种向量化指令的实现，最后是默认的实现。

动态派发的实现在TH/vector/simd.h：

#define INIT_DISPATCH_PTR(OP)                                     \ do {                                                            \ size_t i;                                                     \ for (i = 0; i < sizeof(THVector_(OP ## _DISPATCHTABLE)) /     \ sizeof(FunctionDescription); ++i) {      \THVector_(OP ## _DISPATCHPTR) =                             \reinterpret_cast<decltype(THVector_(OP ## _DISPATCHPTR))> \(THVector_(OP ## _DISPATCHTABLE)[i].function);            \if (THVector_(OP ## _DISPATCHTABLE)[i].supportedSimdExt     \& hostSimdExts) {                                       \break;                                                    \}                                                           \}                                                             \} while(0)typedef struct FunctionDescription
{void *function;uint32_t supportedSimdExt;
} FunctionDescription;

INIT_DISPATCH_PTR这个宏做的事就是遍历THVector_(cadd_DISPATCHTABLE)数组，循环内把动态派发指针（THVector_(cadd_DISPATCHPTR)）赋给当前数组元素所对应的函数指针，最后判断一下当前架构是否支持该指令集，如果支持就退出循环；如果向量化指令集都不支持的话，最后还会指向默认实现的函数指针。