一個簡單的CUDA實踐。用於實現前向傳播。
一、演算法設計
（一）問題背景描述和演算法設計？
問題描述：計算某個點 f 的特徵值的“插值”結果。
以二維為例，為了“插值”得到 f 的特徵值，需要用到：各頂點的特徵值 \(f_i\) 和 f 距離該頂點對面的兩條邊的距離的乘積之和。
如果擴充套件到三維，那麼需要求出 \(f\) 和 \(XYZ\) 三個面的距離（記為 uvw），然後再用一樣的思路即可寫出上面的式子。
（二）輸入輸出、以及它們的維度關係
輸入：
- feats 特徵向量：
(N,8,F)分別代表：N 個立方體、每個立方體有 8 個頂點、每個頂點的特徵向量的長度為 F； - points 待求解的點座標：
(N,3)分別代表：N 個待求解的f點（和 N 個立方體對應），每個點的座標是三維的。
輸出：feats_interp N 個點插值後的特徵向量結果。(N,F) 代表：N 個長度為 F 的特徵向量，對應 points 的每個點。
（三）如何並行運算呢？
我們這個問題一共有 N 個立方體，利用每個立方體八個格點的特徵值“插值”出最後的特徵值；並且每個特徵向量的長度是 F，而這 F 個計算也是相互獨立的。因此可以平行計算兩個量。
其實還有個小技巧：看最後的輸出的維度，能並行的數量不大於輸出的維度。比如這裡輸出是 (N,F) 是二維，所以最大能並行兩個量。
二、CUDA 程式設計
（一）GPU的基本結構
GPU 分為：Kernel - Grid - Block - Thread 四級架構。
（二）Block 層的作用是什麼？
為什麼使用 Block+thread 的二層結構，而不是讓 grid 直接管理 thread 的一層結構？
因為受限於 GPU 的硬體結構，Threads 數量是存在上限的。透過引入 Block 這個結構可以擴充套件執行緒數量。
（三）如何合理設定Block的大小？
假設 N=20,F=10，那麼最後的輸出維度是 20x10；我們一個Block的維度是 16*16，所以需要兩個Block才能覆蓋整個輸出。
具體到程式碼上就是下面兩行（也比較好理解）：
const dim3 threads(16, 16);
const dim3 blocks((N+threads.x-1)/threads.x, (F+threads.y-1)/threads.y);
上面這也提示我們一件事情：有的 thread 是不參與運算的，也就是說 thread 的行、列的 index 值不能超過 N 和 F 的範圍，否則會出錯。
（三）編寫核函式的注意事項？
以下面的核函式為例：
template <typename scalar_t>
__global__ void trilinear_fw_kernel(
const torch::PackedTensorAccessor<scalar_t, 3, torch::RestrictPtrTraits, size_t> feats,
const torch::PackedTensorAccessor<scalar_t, 2, torch::RestrictPtrTraits, size_t> points,
torch::PackedTensorAccessor<scalar_t, 2, torch::RestrictPtrTraits, size_t> feat_interp
){
const int n = blockIdx.x * blockDim.x + threadIdx.x;
const int f = blockIdx.y * blockDim.y + threadIdx.y;
// 如果超過輸出範圍，則跳過不能算
if (n>=feats.size(0) || f>=feats.size(2)) return;
// point -1~1. 需要正規化到 [0,1] 之間
const scalar_t u = (points[n][0]+1)/2;
const scalar_t v = (points[n][1]+1)/2;
const scalar_t w = (points[n][2]+1)/2;
const scalar_t a = (1-v)*(1-w);
const scalar_t b = (1-v)*w;
const scalar_t c = v*(1-w);
const scalar_t d = 1-a-b-c;
// 儲存計算結果
feat_interp[n][f] = (1-u)*(a*feats[n][0][f] +
b*feats[n][1][f] +
c*feats[n][2][f] +
d*feats[n][3][f]) +
u*(a*feats[n][4][f] +
b*feats[n][5][f] +
c*feats[n][6][f] +
d*feats[n][7][f]);
}
核函式沒有任何返回值，因此需要把輸入、輸出透過引數傳遞到函式內部。並且這裡是個泛型函式（scalar_t 代表所有可能的資料型別）。__global__ 是 CUDA 特有的關鍵字，代表該函式是被CPU呼叫、GPU執行；類似的還有 __host__ 代表 cpu 呼叫 + 執行、__device__ 代表 GPU 呼叫 + 執行。
torch::PackedTensorAccessor 的作用是把 Tensor 的型別具體化，才能給C++使用。也就是說：torch::Tensor 這種型別在 C++ 中指代不明確，需要進行具體化。我們以 feats 為例講一下它的幾個引數：
- 第二個引數 3 代表
[N,8,F]共三維結構，類比points是[N,F]是二維結構，所以它的引數寫的是 2 ； size_t是位置索引的型別，預設都是size_t型別；
最後還有個細節，如何確定每個 thread 的 index 呢？也就是程式碼的這兩行：
const int n = blockIdx.x * blockDim.x + threadIdx.x;
const int f = blockIdx.y * blockDim.y + threadIdx.y;
blockDim.x 代表 x 方向上（豎直向下）每個 block 由多少行 thread 構成，也就是 threads[0] 的值；blockIdx.x 則就是第幾個 block 了；threadIdx.x 代表當前 block 內的第 x 行。
（四）如何呼叫核函式？
呼叫核函式主要有兩個工作：輸入輸出、Block&Thread 的大小。
torch::Tensor trilinear_fw_cu(
const torch::Tensor feats,
const torch::Tensor points
){
// 1. 輸入輸出
const int N = feats.size(0), F = feats.size(2);
torch::Tensor feat_interp = torch::empty({N, F}, feats.options());
// 2. thread & block 大小
const dim3 threads(16, 16);// 這裡N和F是可並行的，共兩個次元（256不會出錯）
const dim3 blocks((N+threads.x-1)/threads.x, (F+threads.y-1)/threads.y);
// 建議呼叫該Kernel的第二個引數為當前函式名，方便報錯debug
AT_DISPATCH_FLOATING_TYPES(feats.type(), "trilinear_fw_cu",
([&] {
trilinear_fw_kernel<scalar_t><<<blocks, threads>>>(
feats.packed_accessor<scalar_t, 3, torch::RestrictPtrTraits, size_t>(),
points.packed_accessor<scalar_t, 2, torch::RestrictPtrTraits, size_t>(),
feat_interp.packed_accessor<scalar_t, 2, torch::RestrictPtrTraits, size_t>()
);
}));
return feat_interp;
}
我們以上面這個例子為例。feat_interp 是最後的輸出，它的維度是 (N,F) ，這裡直接透過 feats.options() 設定引數屬性；如果需要自定義的話，比如 torch::empty({N,F}, torch::dtype(torch::kInt32).device(feats.device)); 來自定義型別和位置。
threads 的總數一般是 256 或 512，最大不超過 1024。上面我們分析過，可並行量共有兩個，所以我們設定 256=16*16 作為 threads 的大小，blocks 的數量就按照上面的公式計算即可。
我們使用了 AT_DISPATCH_FLOATING_TYPES 這個函式來呼叫核函式，注意呼叫時透過 <<<, >>> 包裹 blocks 和 threads 的大小，這是 cuda 的特殊寫法。