NIUHE

日々私たちが过ごしている日常というのは、実は奇迹の连続なのかもしれんな

PyTorch源码浅析(1)：THTensor

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

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

1
2
3
4
5
6
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实现的），不了解的读者可以先参考一下这篇知乎专栏。

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

1
2
3
4
5
6
7
8
9
10
11
#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：

1
2
3
4
5
6
7
8
9
10
11
#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声明：

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

其中，THTensor_的宏定义为

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

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

1
2
3
4
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的实现，首先来看一下它的声明和属性：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
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<T>实现引用计数而设计的基类，需要实现引用计数的类只需继承它即可。TensorImpl中的每个成员是干什么的基本都有注释，其中strides_是用来实现内存寻址的，即某个ijk脚标对应的内存地址为：

1
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的声明和属性如下：

1
2
3
4
5
6
7
8
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，每个设备类型对应一个：

1
2
3
4
5
6
7
8
9
10
11
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是一个虚基类，它的声明如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
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上内存分配的具体实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
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的实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
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对象：

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

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

1
intrusive_ptr(TTarget* target);

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

Tensor Apply & Dynamic Dispatch

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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处理。后者太复杂了，我们主要看前者的实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
#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);，它的定义为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
// 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：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
#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)）赋给当前数组元素所对应的函数指针，最后判断一下当前架构是否支持该指令集，如果支持就退出循环；如果向量化指令集都不支持的话，最后还会指向默认实现的函数指针。

つづく

上一篇：简介 下一篇：THC
博客 PyTorchTensorTHTensorTHStorage

Powered by Hexo and Theme by Hacker
© 2019 NIUHE