[原始碼解析] NVIDIA HugeCTR，GPU版本引數伺服器---(3)
0x00 摘要
在本系列中，我們介紹了 HugeCTR，這是一個面向行業的推薦系統訓練框架，針對具有模型並行嵌入和資料並行密集網路的大規模 CTR 模型進行了優化。
本文主要介紹HugeCTR所依賴的輸入資料和一些基礎資料結構。其中借鑑了HugeCTR原始碼閱讀 這篇大作，特此感謝。因為 HugeCTR 實際上是一個具體而微的深度學習系統，所以它也實現了眾多基礎功能，值得想研究深度學習框架的朋友仔細研讀。
本系列其他文章如下：
[原始碼解析] NVIDIA HugeCTR，GPU 版本引數伺服器 --(1)
[原始碼解析] NVIDIA HugeCTR，GPU版本引數伺服器--- (2)
0x01 回顧
我們首先回歸一下前文內容，流水線邏輯關係如下：
訓練流程如下：
基於前文知識，我們接下來看看如何處理資料。
0x02 資料集
HugeCTR 目前支援三種資料集格式，即Norm、Raw和Parquet，具體格式參見如下：
Fig. 1: (a) Norm (b) Raw (c) Parquet Dataset Formats
2.1 Norm
為了最大化資料載入效能並最小化儲存，Norm 資料集格式由一組二進位制資料檔案和一個 ASCII 格式的檔案列表組成。模型檔案應指定訓練和測試（評估）集的檔名，樣本中的元素（鍵）最大數目和標籤維度，具體如圖 1（a）所示。
2.1.1 資料檔案
一個資料檔案是一個讀取執行緒的最小讀取粒度，因此每個檔案列表中至少需要10個檔案才能達到最佳效能。資料檔案由header和實際表格（tabular ）資料組成。
Header定義：
typedef struct DataSetHeader_ {
long long error_check; //0: 沒有錯誤檢查；1：check_num
long long number_of_records； //此資料檔案中的樣本數
long long label_dim; //標籤的維度
long long density_dim; //密集特徵的維度
long long slot_num; //每個嵌入的 slot_num
long long reserved[ 3 ]; //保留以備將來使用
資料集頭；
資料定義（每個樣本）：
typedef struct Data_ {
int length; //此示例中的位元組數（可選：僅在 check_sum 模式下）
float label[label_dim];
float dense[dense_dim];
Slot slots[slot_num];
char checkbits; //此樣本的校驗位（可選：僅在checksum模式下）
} Data;
typedef struct Slot_ {
int nnz;
unsigned int* keys; //可在配置檔案的 `solver` 物件中使用 `"input_key_type"` 更改為 `long long`
} Slot;
資料欄位（field）通常有很多樣本。每個樣本以格式化為整數的標籤開始，然後是nnz（非零數）和使用 long long（或無符號整數）格式的輸入key，如圖 1（a）所示。
categorical 的輸入key分佈到插槽（slot）中，不允許重疊。例如：slot[0] = {0,10,32,45}, slot[1] = {1,2,5,67}。如果有任何重疊，它將導致未定義的行為。例如，給定slot[0] = {0,10,32,45}, slot[1] = {1,10,5,67}，查詢10鍵的表將產生不同的結果，結果根據插槽分配給 GPU 的方式。
2.1.2 檔案列表
檔案列表的第一行應該是資料集中資料檔案的數量，然後是這些檔案的路徑，具體如下所示：
$ cat simple_sparse_embedding_file_list.txt
10
./simple_sparse_embedding/simple_sparse_embedding0.data
./simple_sparse_embedding/simple_sparse_embedding1.data
./simple_sparse_embedding/simple_sparse_embedding2.data
./simple_sparse_embedding/simple_sparse_embedding3.data
./simple_sparse_embedding/simple_sparse_embedding4.data
./simple_sparse_embedding/simple_sparse_embedding5.data
./simple_sparse_embedding/simple_sparse_embedding6.data
./simple_sparse_embedding/simple_sparse_embedding7.data
./simple_sparse_embedding/simple_sparse_embedding8.data
./simple_sparse_embedding/simple_sparse_embedding9.data
使用例子如下：
reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Norm,
source = ["./wdl_norm/file_list.txt"],
eval_source = "./wdl_norm/file_list_test.txt",
check_type = hugectr.Check_t.Sum)
2.2 Raw
Raw 資料集格式與 Norm 資料集格式的不同之處在於訓練資料出現在一個二進位制檔案中，並且使用 int32。圖 1 (b) 顯示了原始資料集樣本的結構。
注意：此格式僅接受獨熱資料。
Raw資料集格式只能與嵌入型別 LocalizedSlotSparseEmbeddingOneHot 一起使用。
例子：
reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Raw,
source = ["./wdl_raw/train_data.bin"],
eval_source = "./wdl_raw/validation_data.bin",
check_type = hugectr.Check_t.Sum)
2.3 Parquet
Parquet 是一種面向列的、開源的資料格式。它可用於 Apache Hadoop 生態系統中的任何專案。為了減小檔案大小，它支援壓縮和編碼。圖 1 (c) 顯示了一個示例 Parquet 資料集。有關其他資訊，請參閱parquet 文件。
請注意以下事項：
- Parquet 資料載入器當前不支援巢狀列型別。
- 不允許列中有任何缺失值。
- 與 Norm 資料集格式一樣，標籤和密集特徵列應使用浮點格式。
- Slot 特徵列應使用 Int64 格式。
- Parquet 檔案中的資料列可以按任何順序排列。
- 要從每個 parquet 檔案中的所有行和每個標籤、密集（數字）和槽（分類）特徵的列索引對映中獲取所需資訊，需要一個單獨的
_metadata.json檔案。
例子 _metadata.json：
{
"file_stats": [{"file_name": "file1.parquet", "num_rows": 6528076}, {"file_name": "file2.parquet", "num_rows": 6528076}],
"cats": [{"col_name": "C11", "index": 24}, {"col_name": "C24", "index": 37}, {"col_name": "C17", "index": 30}, {"col_name": "C7", "index": 20}, {"col_name": "C6", "index": 19}],
"conts": [{"col_name": "I5", "index": 5}, {"col_name": "I13", "index": 13}, {"col_name": "I2", "index": 2}, {"col_name": "I10", "index": 10}],
"labels": [{"col_name": "label", "index": 0}]
}
使用如下：
reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Parquet,
source = ["./criteo_data/train/_file_list.txt"],
eval_source = "./criteo_data/val/_file_list.txt",
check_type = hugectr.Check_t.Non,
slot_size_array = [278899, 355877, 203750, 18573, 14082, 7020, 18966, 4, 6382, 1246, 49, 185920, 71354, 67346, 11, 2166, 7340, 60, 4, 934, 15, 204208, 141572, 199066, 60940, 9115, 72, 34])
我們提供了通過一個選項slot_size_array，可以為每個插槽新增偏移量。slot_size_array是一個長度等於槽數的陣列。為了避免新增offset後出現key重複，我們需要保證第i個slot的key範圍在0到slot_size_array[i]之間。我們將以這種方式進行偏移：對於第 i 個槽鍵，我們將其新增偏移量 slot_size_array[0] + slot_size_array[1] + ... + slot_size_array[i - 1]。在上面提到的配置片段中，對於第 0 個插槽，將新增偏移量 0。對於第一個插槽，將新增偏移量 278899。對於第三個插槽，將新增偏移量 634776。
0x03 CSR 格式
嵌入層是基於CSR格式基礎之上搭建的，所以我們首先看看CSR格式。
3.1 什麼是CSR
稀疏矩陣指的是矩陣中的元素大部分是0的矩陣，實際上現實問題中大多數的大規模矩陣都是稀疏矩陣，因此就出現了很多專門針對稀疏矩陣的高效儲存格式，Compressed Sparse Row（CSR）就是其中之一。
這是最簡單的一種格式，每一個元素需要用一個三元組來表示，分別是（行號，列號，數值），對應上圖右邊的一列。這種方式簡單，但是記錄單資訊多（行列），每個三元組自己可以定位，因此空間不是最優。
CSR需要三類資料來表達：數值，列號，行偏移。它不是用三元組來表示一個元素，而是一個整體編碼方式。
- 數值：一個元素。
- 列號 ：元素的列號，
- 行偏移：某一行的第一個元素在values裡面的起始偏移位置。
上圖中，第一行元素1是0偏移，第二行元素2是2偏移，第三行元素5是4偏移，第4行元素6是7偏移。最後會在行偏移之後加上矩陣總的元素個數，本例子中是9。
3.2 HugeCTR 之中的CSR
我們從中找出一個例子看看。因為只是用來儲存slot裡的sparse key，所以沒有列號，因為一個slot裡的sparse key可以直接順序儲存。
* For example data:
* 4,5,1,2
* 3,5,1
* 3,2
* Will be convert to the form of:
* row offset: 0,4,7,9
* value: 4,5,1,2,3,5,1,3,2
我們再從原始碼之中找一些資訊 samples/ncf/preprocess-20m.py。
def write_hugeCTR_data(huge_ctr_data, filename='huge_ctr_data.dat'):
print("Writing %d samples"%huge_ctr_data.shape[0])
with open(filename, 'wb') as f:
#write header
f.write(ll(0)) # 0: no error check; 1: check_num
f.write(ll(huge_ctr_data.shape[0])) # the number of samples in this data file
f.write(ll(1)) # dimension of label
f.write(ll(1)) # dimension of dense feature
f.write(ll(2)) # long long slot_num
for _ in range(3): f.write(ll(0)) # reserved for future use
for i in tqdm.tqdm(range(huge_ctr_data.shape[0])):
f.write(c_float(huge_ctr_data[i,2])) # float label[label_dim];
f.write(c_float(0)) # dummy dense feature
f.write(c_int(1)) # slot 1 nnz: user ID
f.write(c_uint(huge_ctr_data[i,0]))
f.write(c_int(1)) # slot 2 nnz: item ID
f.write(c_uint(huge_ctr_data[i,1]))
3.3 操作類
3.3.1 定義
這裡只給出成員變數，具體可以和上面csr格式進行印證。
class CSR {
private:
const size_t num_rows_; /**< num rows. */
const size_t max_value_size_; /**< number of element of value the CSR matrix will have for
num_rows rows. */
Tensor2<T> row_offset_tensor_;
Tensor2<T> value_tensor_; /**< a unified buffer for row offset and value. */
T* row_offset_ptr_; /**< just offset on the buffer, note that the length of it is
* slot*batchsize+1.
*/
T* value_ptr_; /**< pointer of value buffer. */
size_t size_of_row_offset_; /**< num of rows in this CSR buffer */
size_t size_of_value_; /**< num of values in this CSR buffer */
size_t check_point_row_; /**< check point of size_of_row_offset_. */
size_t check_point_value_; /**< check point of size_of_value__. */
}
3.3.2 建構函式
建構函式之中，會在GPU之上進行分配記憶體。
/**
* Ctor
* @param num_rows num of rows is expected
* @param max_value_size max size of value buffer.
*/
CSR(size_t num_rows, size_t max_value_size)
: num_rows_(num_rows),
max_value_size_(max_value_size),
size_of_row_offset_(0),
size_of_value_(0) {
std::shared_ptr<GeneralBuffer2<CudaHostAllocator>> buff =
GeneralBuffer2<CudaHostAllocator>::create();
buff->reserve({num_rows + 1}, &row_offset_tensor_);
buff->reserve({max_value_size}, &value_tensor_);
buff->allocate();
row_offset_ptr_ = row_offset_tensor_.get_ptr();
value_ptr_ = value_tensor_.get_ptr();
}
3.3.3 生成新行
new_row 之中會生成新行，並且把目前value總數設定到row_offset之中。
/**
* Insert a new row to CSR
* Whenever you want to add a new row, you need to call this.
* When you have pushed back all the values, you need to call this method
* again.
*/
inline void new_row() { // call before push_back values in this line
if (size_of_row_offset_ > num_rows_) CK_THROW_(Error_t::OutOfBound, "CSR out of bound");
row_offset_ptr_[size_of_row_offset_] = static_cast<T>(size_of_value_);
size_of_row_offset_++;
}
3.3.4 插入資料
這裡會插入資料，並且增加value總數。
/**
* Push back a value to this object.
* @param value the value to be pushed back.
*/
inline void push_back(const T& value) {
if (size_of_value_ >= max_value_size_)
CK_THROW_(Error_t::OutOfBound, "CSR out of bound " + std::to_string(max_value_size_) +
"offset" + std::to_string(size_of_value_));
value_ptr_[size_of_value_] = value;
size_of_value_++;
}
0x04 基礎資料結構
因為 HugeCTR 實際上是一個具體而微的深度學習系統，所以它也實現了眾多基礎功能，為了更好的進行分析，我們需要首先介紹一些基礎資料結構。以下只給出各個類的成員變數和必要函式。
4.1 張量
首先就是最基礎的張量概念。
4.1.1 TensorBuffer2
TensorBuffer2 是張量底層的資料，也許聯絡到 PyTorch 的 data 或者 storage 可以更好的理解。
class TensorBuffer2 {
public:
virtual ~TensorBuffer2() {}
virtual bool allocated() const = 0;
virtual void *get_ptr() = 0;
};
4.1.2 Tensor2
這就對應了TF或者PyTorch的張量。
template <typename T>
class Tensor2 {
std::vector<size_t> dimensions_;
size_t num_elements_;
std::shared_ptr<TensorBuffer2> buffer_;
}
成員函式我們選介紹兩個如下：
static Tensor2 stretch_from(const TensorBag2 &bag) {
return Tensor2(bag.dimensions_, bag.buffer_);
}
TensorBag2 shrink() const {
return TensorBag2(dimensions_, buffer_, TensorScalarTypeFunc<T>::get_type());
}
具體如下：
4.1.3 Tensors2
Tensors2 就是 Tensor2 的一個vector。
template <typename T> using Tensors2 = std::vector<Tensor2<T>>;
4.1.4 TensorBag2
PyTorch 之中也有一些Bag字尾名字的類，比如 nn.Embedding和nn.EmbeddingBag。當構建袋子模型時，做一個Embedding跟隨Sum或是Mean常見的。對於可變長度序列，nn.EmbeddingBag 來提供了更加高效和更快速的處理方式，特別是對於可變長度序列。
在 HugeCTR，TensorBag2 可以認為是把 Tensor 放在袋子裡統一處理的類。
class TensorBag2 {
template <typename T>
friend class Tensor2;
std::vector<size_t> dimensions_;
std::shared_ptr<TensorBuffer2> buffer_;
TensorScalarType scalar_type_;
};
using TensorBags2 = std::vector<TensorBag2>;
關於 Tensor 和 Bag 的聯絡，可以參見下面的函式。
template <typename T>
Tensors2<T> bags_to_tensors(const std::vector<TensorBag2> &bags) {
Tensors2<T> tensors;
for (const auto &bag : bags) {
tensors.push_back(Tensor2<T>::stretch_from(bag));
}
return tensors;
}
template <typename T>
std::vector<TensorBag2> tensors_to_bags(const Tensors2<T> &tensors) {
std::vector<TensorBag2> bags;
for (const auto &tensor : tensors) {
bags.push_back(tensor.shrink());
}
return bags;
}
4.1.5 SparseTensor
SparseTensor 是 Sparse 型別的張量，這是3.2 版本加入的，目的是為了統一處理CSR格式，或者說是統一處理稀疏矩陣，可以有效儲存和處理大多數元素為零的張量。後續在讀取資料到GPU時候會有分析。我們對比一下 CSR 格式，就可以看出來其內部機制就對應了CSR 的 rowoffset 和 value。其具體定義如下：
template <typename T>
class SparseTensor {
std::vector<size_t> dimensions_;
std::shared_ptr<TensorBuffer2> value_buffer_;
std::shared_ptr<TensorBuffer2> rowoffset_buffer_;
std::shared_ptr<size_t> nnz_; // maybe size_t for FixedLengthSparseTensor
size_t rowoffset_count_;
};
示意圖如下：
我們從中找出一個例子看看。因為只是用來儲存slot裡的sparse key，所以沒有列號，因為一個slot裡的sparse key可以直接順序儲存。
* For example data:
* 4,5,1,2
* 3,5,1
* 3,2
* Will be convert to the form of:
* row offset: 0,4,7,9
* value: 4,5,1,2,3,5,1,3,2
對應下圖：
成員函式介紹如下：
static SparseTensor stretch_from(const SparseTensorBag &bag) {
return SparseTensor(bag.dimensions_, bag.value_buffer_, bag.rowoffset_buffer_, bag.nnz_,
bag.rowoffset_count_);
}
SparseTensorBag shrink() const {
return SparseTensorBag(dimensions_, value_buffer_, rowoffset_buffer_, nnz_, rowoffset_count_,
TensorScalarTypeFunc<T>::get_type());
}
PyTorch
PyTorch 有 sparse_coo_tensor 可以實現類似的功能。PyTorch 支援不同layout的張量，大家可以從 torch/csrc/utils/tensor_layouts.cpp 找到，比如 at::Layout::Strided，at::Layout::Sparse，at::Layout::SparseCsr，at::Layout::Mkldnn 等等，這些對應了不同的記憶體佈局模式。
使用稀疏張量時候，提供一對 dense tensors：一個value張量，一個二維indice張量，也有其他輔助引數。
>>> i = [[1, 1]]
>>> v = [3, 4]
>>> s=torch.sparse_coo_tensor(i, v, (3,))
>>> s
tensor(indices=tensor([[1, 1]]),
values=tensor( [3, 4]),
size=(3,), nnz=2, layout=torch.sparse_coo)
TensorFlow
TensorFlow 也有 SparseTensor 型別來表示多維稀疏資料。一個 SparseTensor 使用三個稠密張量來表示：
- indices 表示稀疏張量的非零元素座標。
- values 則對應每個非零元素的值。
- shape 表示本稀疏張量轉換為稠密形式後的形狀。
比如下面程式碼：
indices = tf.constant([[0, 0], [1, 1], [2,2]], dtype=tf.int64)
values = tf.constant([1, 2, 3], dtype=tf.float32)
shape = tf.constant([3, 3], dtype=tf.int64)
sparse = tf.SparseTensor(indices=indices,
values=values,
dense_shape=shape)
dense = tf.sparse_tensor_to_dense(sparse, default_value=0)
with tf.Session() as session:
sparse, dense = session.run([sparse, dense])
print('Sparse is :\n', sparse)
print('Dense is :\n', dense)
列印出來如下：
Sparse is :
SparseTensorValue(indices=array([[0, 0],
[1, 1],
[2, 2]]), values=array([1., 2., 3.], dtype=float32), dense_shape=array([3, 3]))
Dense is :
[[1. 0. 0.]
[0. 2. 0.]
[0. 0. 3.]]
4.1.6 SparseTensorBag
這個類似 TensorBag 的功能，具體如下：
class SparseTensorBag {
template <typename T>
friend class SparseTensor;
std::vector<size_t> dimensions_;
std::shared_ptr<TensorBuffer2> value_buffer_;
std::shared_ptr<TensorBuffer2> rowoffset_buffer_;
std::shared_ptr<size_t> nnz_;
size_t rowoffset_count_;
TensorScalarType scalar_type_;
SparseTensorBag(const std::vector<size_t> &dimensions,
const std::shared_ptr<TensorBuffer2> &value_buffer,
const std::shared_ptr<TensorBuffer2> &rowoffset_buffer,
const std::shared_ptr<size_t> &nnz, const size_t rowoffset_count,
TensorScalarType scalar_type)
: dimensions_(dimensions),
value_buffer_(value_buffer),
rowoffset_buffer_(rowoffset_buffer),
nnz_(nnz),
rowoffset_count_(rowoffset_count),
scalar_type_(scalar_type) {}
public:
SparseTensorBag() : scalar_type_(TensorScalarType::None) {}
const std::vector<size_t> &get_dimensions() const { return dimensions_; }
};
4.1.7 向量類
以下是兩個向量類，用來方便使用者使用。
using TensorBags2 = std::vector<TensorBag2>;
template <typename T>
using SparseTensors = std::vector<SparseTensor<T>>;
4.2 記憶體
我們接下來看看一些記憶體相關類。
4.2.1 Allocator
首先看看如何為tensor等變數分配記憶體。
4.2.1.1 HostAllocator
HostAllocator 作用是在host之上管理記憶體。
class HostAllocator {
public:
void *allocate(size_t size) const { return malloc(size); }
void deallocate(void *ptr) const { free(ptr); }
};
後面幾個實現都是呼叫了CUDA函式來進行記憶體分配，比如 cudaHostAlloc，有興趣讀者可以深入學習。
4.2.1.2 CudaHostAllocator
呼叫CUDA方法在主機上分配記憶體
class CudaHostAllocator {
public:
void *allocate(size_t size) const {
void *ptr;
CK_CUDA_THROW_(cudaHostAlloc(&ptr, size, cudaHostAllocDefault));
return ptr;
}
void deallocate(void *ptr) const { CK_CUDA_THROW_(cudaFreeHost(ptr)); }
};
4.2.1.3 CudaManagedAllocator
cudaMallocManaged 分配旨在供主機或裝置程式碼使用的記憶體，算是一種統一分配記憶體的方法。
class CudaManagedAllocator {
public:
void *allocate(size_t size) const {
void *ptr;
CK_CUDA_THROW_(cudaMallocManaged(&ptr, size));
return ptr;
}
void deallocate(void *ptr) const { CK_CUDA_THROW_(cudaFree(ptr)); }
};
4.2.1.4 CudaAllocator
該類是在裝置上分配記憶體。
class CudaAllocator {
public:
void *allocate(size_t size) const {
void *ptr;
CK_CUDA_THROW_(cudaMalloc(&ptr, size));
return ptr;
}
void deallocate(void *ptr) const { CK_CUDA_THROW_(cudaFree(ptr)); }
};
4.2.2 GeneralBuffer2
分析完如何分配記憶體，我們接下來看看如何封裝記憶體，具體通過 GeneralBuffer2 完成的。GeneralBuffer2 可以認為是一個對大段記憶體的統一封裝，具體在其上可以有若干Tensor。
4.2.2.1 定義
這裡都忽略了成員函式，內部類也忽略了成員函式。
- allocator ：具體記憶體分配器，也區分在GPU分配還是CPU分配。
- ptr_ ：指向分配的記憶體；
- total_size_in_bytes_ ：記憶體大小；
- reserved_buffers_ ：前期預留buffer，後續會統一分配;
具體內部類為：
- BufferInternal 是介面。
- TensorBufferImpl 是 Tensor2 對應的buffer實現。
- BufferBlockImpl 則是在構建網路時候會用到。
具體程式碼如下：
template <typename Allocator>
class GeneralBuffer2 : public std::enable_shared_from_this<GeneralBuffer2<Allocator>> {
class BufferInternal {
public:
virtual ~BufferInternal() {}
virtual size_t get_size_in_bytes() const = 0;
virtual void initialize(const std::shared_ptr<GeneralBuffer2> &buffer, size_t offset) = 0;
};
class TensorBufferImpl : public TensorBuffer2, public BufferInternal {
size_t size_in_bytes_;
std::shared_ptr<GeneralBuffer2> buffer_;
size_t offset_;
};
template <typename T>
class BufferBlockImpl : public BufferBlock2<T>, public BufferInternal {
size_t total_num_elements_;
std::shared_ptr<TensorBufferImpl> buffer_impl_;
Tensor2<T> tensor_;
bool finalized_;
std::vector<std::shared_ptr<BufferInternal>> reserved_buffers_;
};
Allocator allocator_;
void *ptr_;
size_t total_size_in_bytes_;
std::vector<std::shared_ptr<BufferInternal>> reserved_buffers_;
}
4.2.2.2 TensorBufferImpl
就是指向了一個 GeneralBuffer2，然後設定了自己的offset和大小。
void initialize(const std::shared_ptr<GeneralBuffer2> &buffer, size_t offset) {
buffer_ = buffer;
offset_ = offset;
}
4.2.2.2 BufferBlockImpl 關鍵函式
BufferBlockImpl 和 TensorBufferImpl 可以來比較一下。
其中，BufferBlock2 是 BufferBlockImpl 的介面類。
template <typename T>
class BufferBlock2 {
public:
virtual ~BufferBlock2() {}
virtual void reserve(const std::vector<size_t> &dimensions, Tensor2<T> *tensor) = 0;
virtual Tensor2<T> &as_tensor() = 0;
};
BufferBlockImpl 是一組連續的 Tensor，某些特定的實現需要連續的記憶體，比如權重。
std::shared_ptr<BufferBlock2<float>> train_weight_buff = blobs_buff->create_block<float>();
// 省略其他程式碼......
network->train_weight_tensor_ = train_weight_buff->as_tensor();
BufferBlockImpl 多了一個reserve方法，用來預留記憶體空間，在此空間之上生成內部tensor。
void reserve(const std::vector<size_t> &dimensions, Tensor2<T> *tensor) override {
if (finalized_) {
throw std::runtime_error(ErrorBase + "Buffer block is finalized.");
}
size_t num_elements = get_num_elements_from_dimensions(dimensions);
size_t size_in_bytes = num_elements * TensorScalarSizeFunc<T>::get_element_size();
std::shared_ptr<TensorBufferImpl> buffer_impl =
std::make_shared<TensorBufferImpl>(size_in_bytes);
reserved_buffers_.push_back(buffer_impl);
*tensor = Tensor2<T>(dimensions, buffer_impl);
total_num_elements_ += num_elements;
}
initialize 會對內部進行配置
void initialize(const std::shared_ptr<GeneralBuffer2> &buffer, size_t offset) {
size_t local_offset = 0;
for (const std::shared_ptr<BufferInternal> &buffer_impl : reserved_buffers_) {
buffer_impl->initialize(buffer, offset + local_offset);
local_offset += buffer_impl->get_size_in_bytes();
}
reserved_buffers_.clear();
if (!finalized_) {
buffer_impl_ = std::make_shared<TensorBufferImpl>(
total_num_elements_ * TensorScalarSizeFunc<T>::get_element_size());
tensor_ = Tensor2<T>({total_num_elements_}, buffer_impl_);
finalized_ = true;
}
buffer_impl_->initialize(buffer, offset);
}
4.2.2.3 GeneralBuffer2 關鍵函式
reserve 方法會把某一個張量對應的記憶體需求用 TensorBufferImpl 的形式記錄在reserved_buffers_之中，然後生成這個張量，而且就是用TensorBufferImpl 生成。
template <typename T>
void reserve(const std::vector<size_t> &dimensions, Tensor2<T> *tensor) {
if (allocated()) {
throw std::runtime_error(ErrorBase + "General buffer is finalized.");
}
size_t size_in_bytes =
get_num_elements_from_dimensions(dimensions) * TensorScalarSizeFunc<T>::get_element_size();
std::shared_ptr<TensorBufferImpl> buffer_impl =
std::make_shared<TensorBufferImpl>(size_in_bytes);
reserved_buffers_.push_back(buffer_impl);
*tensor = Tensor2<T>(dimensions, buffer_impl);
}
create_block 會針對BufferBlock2進行建立。
template <typename T>
std::shared_ptr<BufferBlock2<T>> create_block() {
if (allocated()) {
throw std::runtime_error(ErrorBase + "General buffer is finalized.");
}
std::shared_ptr<BufferBlockImpl<T>> block_impl = std::make_shared<BufferBlockImpl<T>>();
reserved_buffers_.push_back(block_impl);
return block_impl;
}
allocate 會遍歷註冊的 BufferInternal，累積其總大小，最後呼叫 allocator_ 進行分配記憶體。
void allocate() {
if (ptr_ != nullptr) {
throw std::runtime_error(ErrorBase + "Memory has already been allocated.");
}
size_t offset = 0;
for (const std::shared_ptr<BufferInternal> &buffer : reserved_buffers_) {
// 對 BufferInternal（比如TensorBufferImpl）內部進行配置
buffer->initialize(this->shared_from_this(), offset);
size_t size_in_bytes = buffer->get_size_in_bytes();
if (size_in_bytes % 32 != 0) {
size_in_bytes += (32 - size_in_bytes % 32);
}
offset += size_in_bytes;
}
reserved_buffers_.clear();
total_size_in_bytes_ = offset;
if (total_size_in_bytes_ != 0) {
ptr_ = allocator_.allocate(total_size_in_bytes_);
}
}
4.2.4 小結
至此，Tensor的邏輯擴充一下：
- TensorBufferImpl 的 buffer 是GeneralBuffer2；
- GeneralBuffer2 的 ptr 是由CudaAllocator在GPU之中分配的；GeneralBuffer2 可以認為是一個對大段記憶體的統一封裝，在其上可以有若干Tensor。這些Tensor先reserve記憶體，然後統一分配。
- TensorBufferImpl 的 offset_ 就指向了 GeneralBuffer2 的 ptr 之中具體的某一個記憶體偏移；
- BufferBlockImpl 用來實現一個連續的Tensor記憶體。
如果還有另外一個 Tensor2，則其 TensorBufferImpl.offset 會指向 GPU記憶體的另外一個offset，比如下面有兩個張量，Tensor 1 和 Tensor 2。
0xFF 參考
https://web.eecs.umich.edu/~justincj/teaching/eecs442/notes/linear-backprop.html