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
TensorImpl
StorageImpl
智能指针 Intrusive_ptr
Tensor Apply & Dynamic Dispatch
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:
// 指向实际数据存储位置，也就是指向StorageImpl
Storage 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)};
}
// 这四个参数用来构造一个DataPtr
return {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);
#else
CAFFE_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);
#else
free(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_ptr
c10::intrusive_ptr<at::StorageImpl>::reclaim(THStorage_(new)()), // Storage&& storage
at::CPUTensorId(), // TensorTypeId type_id
false // 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_ptr
data_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
#endif
FUNCTION_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)）赋给当前数组元素所对应的函数指针，最后判断一下当前架构是否支持该指令集，如果支持就退出循环；如果向量化指令集都不支持的话，最后还会指向默认实现的函数指针。
扫一扫
3292
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
