TH库学习(一): THTensor, THStorage, THAllocator介绍

pytorch中的底层很多代码都是来源于的torch的低层Tensor库
TH = TorcH
THC = TorcH Cuda
THCS = TorcH Cuda Sparse
THCUNN = TorcH CUda Neural Network (see cunn)
THD = TorcH Distributed
THNN = TorcH Neural Network
THS = TorcH Sparse
特别推荐查看Torch7的官方介绍Tensor

前面我们介绍strided indexing scheme时就已经说过，很多矩阵(ndarray)库都是存储和表示分开，以便很多变量共享内存，TH库也是这样

THTensor, THStorage, THAllocator的关系

先整体概览一下他们3者的关系:

数据存储: THStorage

所有在CPU上的张量实际上都是内存中的一个一维C数组（C指针）data来存储，并且使用引用计数（reference count）来管理内存。

typedef struct THStorage
{
    real *data;
    ptrdiff_t size;
    int refcount; // 引用计数
    char flag;
    THAllocator *allocator;
    void *allocatorContext; // GPU或CPU的上下文，在GPU时有用，在CPU暂无用到
    struct THStorage *view;
} THStorage;
// flag是4位标识符号,例如flag=0b0101 代表可以释放内存且采用了引用计数 ，
// 所以在storage释放内存时首先判断是否可以释放，然后如果采用了引用计数，只有计数为0时才可以释放内存。
#define TH_STORAGE_REFCOUNTED 1
#define TH_STORAGE_RESIZABLE  2
#define TH_STORAGE_FREEMEM    4
#define TH_STORAGE_VIEW       8

所有构造新THStorage的函数都以new开头，后面跟具有相关含义的后缀名。

// 空的THStorage -> return THStorage_(newWithSize)(0);
TH_API THStorage* THStorage_(new)(void); 
// 指定大小的THStorage -> return THStorage_(newWithAllocator)(size, &THDefaultAllocator, NULL);
// 未赋值，直接malloc了内存，未初始化这块内存
TH_API THStorage* THStorage_(newWithSize)(ptrdiff_t size);
TH_API THStorage* THStorage_(newWithSize1)(real);
TH_API THStorage* THStorage_(newWithSize2)(real, real);
TH_API THStorage* THStorage_(newWithSize3)(real, real, real);
TH_API THStorage* THStorage_(newWithSize4)(real, real, real, real);
TH_API THStorage* THStorage_(newWithMapping)(const char *filename, ptrdiff_t size, int flags);

/* takes ownership of data */
// 生成一个THStorage，其data直接指向传入的data，不再重新开辟内存
// return THStorage_(newWithDataAndAllocator)(data, size, &THDefaultAllocator, NULL);
TH_API THStorage* THStorage_(newWithData)(real *data, ptrdiff_t size);

TH_API THStorage* THStorage_(newWithAllocator)(ptrdiff_t size,
                                               THAllocator* allocator,
                                               void *allocatorContext);
TH_API THStorage* THStorage_(newWithDataAndAllocator)(
    real* data, ptrdiff_t size, THAllocator* allocator, void *allocatorContext);

内存分配: THAllocator

// THAllocator.h
/*
 * THAllocator 定义了3个函数指针 malloc realloc free
 * 即指向函数的指针,在THAllocator.c中会把具体的函数地址赋给它们
 */
typedef struct THAllocator {
  void* (*malloc)(void*, ptrdiff_t);
  void* (*realloc)(void*, void*, ptrdiff_t);
  void (*free)(void*, void*);
} THAllocator;

// THAllocator.c
// 它最终实际调用的就是#include 中的malloc(size)
// 但是它这么套了几层，是为了考虑各个平台上malloc时数据对齐的问题，这里不再深究
static void *THDefaultAllocator_alloc(void* ctx, ptrdiff_t size) {
  return THAlloc(size);
}
// realloc(ptr, size)
static void *THDefaultAllocator_realloc(void* ctx, void* ptr, ptrdiff_t size) {
  return THRealloc(ptr, size);
}
// free(ptr)
static void THDefaultAllocator_free(void* ctx, void* ptr) {
  THFree(ptr);
}
THAllocator THDefaultAllocator = {
  &THDefaultAllocator_alloc,
  &THDefaultAllocator_realloc,
  &THDefaultAllocator_free
};

数据查看(表示): THTensor

typedef struct THTensor
{
    // 一个用来存储张量各个维度大小的一维数组
    int64_t *size;
    // 一个用来存储张量各个角标偏移量的数组，之前strided indexing scheme介绍过
    int64_t *stride;
    // 维度
    int nDimension; 
    // Note: storage->size may be greater than the recorded size
    // of a tensor
    // 持有的实际存储的指针，存储大小可能实际上是大于于等于张量大小
    THStorage *storage; 
    // 存储偏移
    ptrdiff_t storageOffset;
    // 引用计数
    int refcount;
    // 同THStorage
    char flag;

} THTensor;

如果通过THTensor去查看数据实际存储地THStorage呢？我们在strided indexing scheme这章里讲得很清楚了，主要就通过storageOffset和stride。

这个图片还缺了一个storageOffset未加上去
我们以torch.narrow这个函数为例，解读如何通过storageOffset和stride实现THStorage的不同视图THTensor

# narrow(dimension, start, length) → Tensor
# 抽取Tensor某个维度上的一部分
# stride: 3,1 size: 3,3
x = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
 1  2  3
 4  5  6
 7  8  9
# 抽取x第0个维度，从0开始，长度为2，所有就是第0行和第1行
# dimension=0 start=0 length=2
# self->storageOffset += firstIndex*self->stride[dimension];
# self->size[dimension] = size;
# 可以看到storageOffset=0*3=0，不变
# size: length,3 = 2,3
# stride不变
x.narrow(0, 0, 2)
 1  2  3
 4  5  6
[torch.FloatTensor of size (2,3)]
# 抽取x第1个维度，从1开始，长度为2，所有就是第1行和第2行
# dimension=1 start=1 length=2
# self->storageOffset += firstIndex*self->stride[dimension];
# self->size[dimension] = size;
# 可以看到storageOffset=1*1=1,其起始偏移变为了1，刚好把第一个元素1跳过去
# size: 3,length = 3,2
# stride不变
x.narrow(1, 1, 2)
 2  3
 5  6
 8  9
[torch.FloatTensor of size (3,2)]

低层实现如下:

void THTensor_(narrow)(THTensor *self, THTensor *src, int dimension, int64_t firstIndex, int64_t size)
{
  if(!src)
    src = self;

  THArgCheck( (dimension >= 0) && (dimension < src->nDimension), 2, "out of range");
  THArgCheck( (firstIndex >= 0) && (firstIndex < src->size[dimension]), 3, "out of range");
  THArgCheck( (size > 0) && (firstIndex <= src->size[dimension] - size), 4, "out of range");

  THTensor_(set)(self, src);

  if(firstIndex > 0)
    self->storageOffset += firstIndex*self->stride[dimension];

  self->size[dimension] = size;
}

---

PyTorch源码浅析（一）
PyTorch – Internal Architecture Tour
机器之心翻译: PyTorch – Internal Architecture Tour
torch.Tensor