PyTorch 源碼解讀之 cpp_extension：揭秘 C++/CUDA 算子實(shí)現(xiàn)和調(diào)用全流程

點(diǎn)藍(lán)色字關(guān)注“機(jī)器學(xué)習(xí)算法工程師”

設(shè)為星標(biāo)，干貨直達(dá)！

AI編輯：我是小將

本文作者：OpenMMLab @001024

https://zhuanlan.zhihu.com/p/348555597

本文已由原作者授權(quán)

“Python 用戶友好卻運(yùn)行效率低”，“C++ 運(yùn)行效率較高，但實(shí)現(xiàn)一個(gè)功能代碼量會(huì)遠(yuǎn)大于 Python”。平常學(xué)習(xí)工作中你是否常聽到類似的說法？在 Python 大行其道的今天，你是否經(jīng)常也會(huì)面臨代碼的瓶頸，而為運(yùn)行加速而煩惱呢？“我的代碼剛跑 10 步，隔壁同學(xué)的已經(jīng)跑完第一個(gè) epoch 了。”--這究竟是人性的扭曲還是科學(xué)的淪喪？荀子有言“君子性非異也，善假于物也”。本期《源碼解讀》帶你走進(jìn) “Pytorch 中 (神秘) 的 C++ / CUDA 擴(kuò)展"。

• 本期主題：結(jié)合 Python 與 C++ 各自的優(yōu)點(diǎn)，在 PyTorch 中加入 C++ / CUDA的擴(kuò)展，而讓我們自己更好地使用工具而不為工具所束縛。

• 代碼來源：MMCV, PyTorch。

• 注：C++ / CUDA 擴(kuò)展一般有”預(yù)編譯“ 與 ”實(shí)時(shí)編譯“ （just-in-time, JIT）模式。本期主要介紹”預(yù)編譯“模式。

1. 由擴(kuò)展的調(diào)用方式說起

當(dāng)你想為自己的代碼添加擴(kuò)展進(jìn)行加速時(shí)，我們可以先來看看經(jīng)典的例子中是怎么處理的。對(duì)檢測或分割稍有了解的同學(xué)應(yīng)該知道，nms 的計(jì)算是最常見的用到了 C++ / CUDA 擴(kuò)展的算子。

from mmcv import _ext as ext_module
from torch.autograd import Function

def nms(boxes, scores, iou_threshold, offset=0):
    inds = NMSop.apply(boxes, scores, iou_threshold, offset)
    dets = torch.cat((boxes[inds], scores[inds].reshape(-1, 1)), dim=1)
    return dets, inds

class NMSop(torch.autograd.Function):
    @staticmethod
    def forward(ctx, bboxes, scores, iou_threshold, offset):
        inds = ext_module.nms(
            bboxes, scores, iou_threshold=float(iou_threshold), offset=offset)
        return inds

    @staticmethod
    def symbolic(g, bboxes, scores, iou_threshold, offset):
        pass  # onnx 轉(zhuǎn)換相關(guān)

Function (見往期內(nèi)容torch.autograd)。NMSop的forward函數(shù)內(nèi)核調(diào)用的是mmcv._ext.nms模塊，但實(shí)際上我們?cè)?MMCV 源碼中是看不到 _ext module 的。只有在編譯好的mmcv 庫 (MMCV_WITH_OPS=True python setup.py build_ext --inplace) 會(huì)出現(xiàn) mmcv/_ext.cpython-xxx.so 文件，只有這時(shí)在 Python 中運(yùn)行 import mmcv._ext 才會(huì)成功。看來 C++ 擴(kuò)展是通過 setup.py 來執(zhí)行編譯的。

2. setup.py -- 擴(kuò)展的編譯文件

基于預(yù)編譯的擴(kuò)展由于需要編譯，而setup.py文件正是基于setuptools的編譯腳本。因此一個(gè) Python package 的擴(kuò)展是可以在setup.py文件中找到其蛛絲馬跡的。這里我們截取一段mmcv的 setup.py 文件，

setup(
    name='mmcv',
    install_requires=install_requires,
    # 需要編譯的c++/cuda擴(kuò)展
    ext_modules=get_extensions(),
    # cmdclass 為python setup.py --build_ext命令指定行為
    cmdclass={'build_ext':  torch.utils.cpp_extension.BuildExtension})

這里可以看到 setup函數(shù)中一個(gè)主要的參數(shù)ext_modules，該參數(shù)需要指定為一個(gè)Extension列表，代表實(shí)際需要編譯的擴(kuò)展。目前該參數(shù)由get_extensions函數(shù)獲得。其中 get_extensions函數(shù)定義如下（節(jié)選）

def get_extensions():
    extensions = []
    ext_name = 'mmcv._ext'
    from torch.utils.cpp_extension import (CUDAExtension, CppExtension)

    if torch.cuda.is_available():
        # CUDA編譯擴(kuò)展
        extra_compile_args['nvcc'] = [cuda_args] if cuda_args else []
        # 編譯./mmcv/ops/csrc/pytorch文件夾中的所有文件
        op_files = glob.glob('./mmcv/ops/csrc/pytorch/*')
        extension = CUDAExtension
    else:
        # C++ 編譯擴(kuò)展
        op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp')
        extension = CppExtension
     include_path = os.path.abspath('./mmcv/ops/csrc')
     ext_ops = extension(
         name=ext_name,  # 擴(kuò)展模塊名
         sources=op_files,  # 擴(kuò)展文件
         include_dirs=[include_path],  # 頭文件地址
         define_macros=define_macros,  # 預(yù)定義宏
         extra_compile_args=extra_compile_args)  # 其他編譯選項(xiàng)
      extensions.append(ext_ops)
    return extensions

在上述代碼中我們終于看到了mmcv._ext，該名字正是新定義的擴(kuò)展的名字。由此我們便知道上文中提到的mmcv._ext模塊實(shí)際上是在 setup.py 文件中指定其模塊名字的。另外我們發(fā)現(xiàn)用于生成擴(kuò)展的函數(shù)會(huì)隨系統(tǒng)環(huán)境不同而有所區(qū)別，當(dāng)系統(tǒng)中沒有 CUDA 時(shí)會(huì)調(diào)用 CppExtension，且只編譯所有 .cpp文件，反之則調(diào)用 CUDAExtension。其實(shí) CppExtension 與 CUDAExtension 都是基于setuptools.Extension的擴(kuò)展，這兩個(gè)函數(shù)都額外將系統(tǒng)目錄中的 torch/include 加入到 C++ 編譯時(shí)的include_dirs中，另外 CUDAExtension 會(huì)額外將CUDA相關(guān)的庫以及頭文件加到默認(rèn)編譯搜索路徑中。由 setup.py 文件我們還了解到送給編譯的其他信息，如擴(kuò)展文件的源文件地址，在 MMCV中則是存放于 ./mmcv/ops/csrc/pytorch/中。其他信息如 include_dirs, define_macros, extra_compile_args 則會(huì)在 torch/utils/cpp_extension.py:BuildExtension一并形成最終的 gcc /nvcc 的命令。

class BuildExtension(build_ext, object):
    # 只顯示核心代碼
    def build_extensions(self):
        # 檢查二進(jìn)制接口兼容性
        self._check_abi()

        # 注冊(cè) cuda 代碼 (.cu, .cuh)
        self.compiler.src_extensions += ['.cu', '.cuh']
        def unix_wrap_compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
            try:
                original_compiler = self.compiler.compiler_so
                if _is_cuda_file(src):
                    # 對(duì) cuda 文件調(diào)用 nvcc 命令
                    nvcc = _join_cuda_home('bin', 'nvcc')
                    self.compiler.set_executable('compiler_so', nvcc)
                    if isinstance(cflags, dict):
                        cflags = cflags['nvcc']
                    cflags = COMMON_NVCC_FLAGS + ['--compiler-options',
                                                  "'-fPIC'"] + cflags + _get_cuda_arch_flags(cflags)
                elif isinstance(cflags, dict):
                    # 默認(rèn) c++ 程序的 flags
                    cflags = cflags['cxx']
                # 強(qiáng)制性使用 --std=c++11
                if not any(flag.startswith('-std=') for flag in cflags):
                    cflags.append('-std=c++11')
                # c++ / cuda 程序編譯入口
                original_compile(obj, src, ext, cc_args, cflags, pp_opts)
            finally:
                # 將之前覆蓋的默認(rèn)編譯器還原
                self.compiler.set_executable('compiler_so', original_compiler)

以上過程了解清楚之后我們運(yùn)行 MMCV_WITH_OPS=True python setup.py build_ext --inplace 指令。

/usr/local/cuda-10.0/bin/nvcc -DMMCV_WITH_CUDA -I/home-to-/mmcv/mmcv/ops/csrc -I/home-to-//torch/include -I/home-to/torch/include/torch/csrc/api/include -I/home-to/torch/include/TH -I/home-to/torch/include/THC -I/usr/local/cuda-10.0/include -I/home-to/python3.7m -c ./mmcv/ops/csrc/pytorch/nms_cuda.cu -o build/temp.linux-x86_64-3.7/./mmcv/ops/csrc/pytorch/nms_cuda.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=_ext -D_GLIBCXX_USE_CXX11_ABI=0 -gencode=arch=compute_61,code=sm_61 -std=c++11

gcc -pthread -B compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -DMMCV_WITH_CUDA -I/home-to/mmcv/ops/csrc -I/home-to/torch/include -I/home-to/torch/include/torch/csrc/api/include -I/home-to/torch/include/TH -I/home-to/torch/include/THC -I/usr/local/cuda-10.0/include -I/home-to/python3.7m -c ./mmcv/ops/csrc/pytorch/nms.cpp -o build/temp.linux-x86_64-3.7/./mmcv/ops/csrc/pytorch/nms.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=_ext -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11
...

在上述運(yùn)行結(jié)果中我們可以看到

• 編譯器對(duì) CUDA 文件自動(dòng)調(diào)用 nvcc 而對(duì) .cpp 文件則調(diào)用 gcc

• 被 CUDAExtension包裝過后系統(tǒng)自動(dòng)加入了 Python，PyTorch， CUDA 等庫中的頭文件以及庫地址，系統(tǒng)架構(gòu)信息(-gencode)與編譯優(yōu)化信息(-O3 等)

• 通過 -DTORCH_EXTENSION_NAME=_ext 將TORCH_EXTENSION_NAME宏賦值為_ext。這看來也絕非是閑來之筆，欲知后事如何，我們看下一節(jié)分解

3. PYBIND11_MODULE -- Python 與 C++ 的橋梁

上文說到通過 setup.py 我們編譯了擴(kuò)展文件。可是目前仍然有個(gè)疑問，為什么編譯出來的 C++ / CUDA 二進(jìn)制文件可以在 Python 中直接被調(diào)用呢？再次檢測編譯的所有文件，發(fā)現(xiàn)其中有個(gè)文件 pybind.cpp 十分可疑，其打開后大致形式如下。

#include <torch/extension.h>// 函數(shù)聲明，具體實(shí)現(xiàn)在其他文件
Tensor nms(Tensor boxes, Tensor scores, float iou_threshold, int offset);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("nms", &nms, "nms (CPU/CUDA) ", py::arg("boxes"), py::arg("scores"),
        py::arg("iou_threshold"), py::arg("offset"));
}

這里PYBIND11_MODULE是一個(gè)宏，定義在 pybind11 庫中(見pybind11/include/pybind11/pybind11.h)。而 pybind11 是一個(gè)用來在 C++ 代碼中創(chuàng)建 Python的連接的庫。找到了源頭，我們進(jìn)一步分析。

這里PYBIND11_MODULE 的作用是為 C++ 代碼接入 Python 解釋器提供入口。以上述代碼為例， TORCH_EXTENSION_NAME 正是在上文 gcc編譯過程中出現(xiàn)的宏，對(duì)應(yīng)為extension的 name 變量。因此在這里會(huì)被解釋成 _ext（注意沒有雙引號(hào)） 。m 則代表 TORCH_EXTENSION_NAME 所對(duì)應(yīng)的模塊實(shí)例（實(shí)際上可以指定為任何名字）。{}中的每個(gè) m.def都定義了一個(gè) _ext 的成員函數(shù)，其一般形式為 m.def("函數(shù)名"，具體 C++ 實(shí)現(xiàn)的函數(shù)指針, "文檔", 參數(shù)列表)。通過這種形式，nms也就順利地成為了mmcv._ext的成員函數(shù)，其具體實(shí)現(xiàn)為已經(jīng)定義好的 nms 函數(shù)（對(duì)這個(gè)函數(shù)的分析我們會(huì)在下節(jié)講到）。在 Python 中也就可以運(yùn)行from mmcv._ext import nms了。如果對(duì)這里的定義仍然不清楚，我們可以把該宏用 C++ 編譯器展開一下：

Tensor nms(Tensor boxes, Tensor scores, float iou_threshold, int offset);
static void pybind11_init__ext(pybind11::module &); 
extern "C" __attribute__ ((visibility("default"))) PyObject *PyInit__ext() 
  { 
    // 省略部分代碼  
    auto m = pybind11::module("_ext");  // m 變量的初始化是在宏內(nèi)部
    try { pybind11_init__ext(m); return m.ptr(); } 
} 
void pybind11_init__ext(pybind11::module &m) {
   // 添加成員函數(shù)
   m.def("nms", &nms, "nms (CPU/CUDA) ", py::arg("boxes"), py::arg("scores"),
       py::arg("iou_threshold"), py::arg("offset"));
}

其中 PyObject *PyInit_ 定義在 Python.h中，正是 C++ 中聲明 Python module 的官方方法（可見官方 Python 文檔）。這里 PyInit_后接的_ext其實(shí)就是 TORCH_EXTENSION_NAME宏解釋得到。意味著新聲明了一個(gè) 名為_ext 的 Python module。

4. cpp/cu文件 -- 算子的具體實(shí)現(xiàn)

通過對(duì) PYBIND11_MODULE 的分析后，我們了解了 mmcv._ext.nms 具體的實(shí)現(xiàn)是一個(gè)聲明為 Tensor nms(Tensor boxes, Tensor scores, float iou_threshold, int offset); 的函數(shù)。該函數(shù)定義在 mmcv/ops/csrc/pytorch/nms.cpp 中

#include <torch/extension.h>
Tensor nms(Tensor boxes, Tensor scores, float iou_threshold, int offset) {
  if (boxes.device().is_cuda()) {
    // cuda 實(shí)現(xiàn)
    return nms_cuda(boxes, scores, iou_threshold, offset);
  } else {
    // c++ 實(shí)現(xiàn)
    return nms_cpu(boxes, scores, iou_threshold, offset);
  }
}

可以看到這時(shí)實(shí)際的實(shí)現(xiàn)方式針對(duì)設(shè)備的不同分為了 nms_cuda 與 nms_cpu 兩種。這里我們先來看 cpp 的實(shí)現(xiàn)。

4.1 CPP 算子實(shí)現(xiàn)

#include <torch/extension.h>using namespace at; // 適當(dāng)改寫
Tensor nms_cpu(Tensor boxes, Tensor scores, float iou_threshold, int offset) {
  // 僅顯示核心代碼
  for (int64_t _i = 0; _i < nboxes; _i++) {
    // 遍歷所有檢測框，稱為主檢測框
    if (select[_i] == false) continue;
    for (int64_t _j = _i + 1; _j < nboxes; _j++) {
      // 對(duì)每個(gè)主檢測框，遍歷其他檢測框，稱為次檢測框
      // 這里只用遍歷上三角元素即可，節(jié)省計(jì)算
      if (select[_j] == false) continue;
      auto ovr = inter / (iarea + areas[j] - inter);
      // 如果次檢測框與主檢測框 iou 過大，則去除該次檢測框
      if (ovr >= iou_threshold) select[_j] = false;
    }
  }
  return order_t.masked_select(select_t);
}

以上即為 nms_cpu 的核心代碼，對(duì)該算法想要有進(jìn)一步了解的同學(xué)可以參看源碼 。這里出現(xiàn)了兩個(gè)for 循環(huán)，實(shí)現(xiàn)上這正是我們希望實(shí)現(xiàn) nms 的 C++ / CUDA 擴(kuò)展的原因。對(duì)于有一定 C++ 基礎(chǔ)的同學(xué)來說代碼應(yīng)該較好理解 （注意這里 int64_t 可理解為C99規(guī)約的為支持不同平臺(tái)的int64類型的 typedef 定義，可直接理解為 int64），但這里同時(shí)也出現(xiàn)了一些新的變量類型，最典型的是 Tensor 數(shù)據(jù)類型。

其實(shí)這里 Tensor 數(shù)據(jù)類型由 torch/extension.h 支持，來源于 pytorch 中 C++ API 中三大 namespace（at, torch 與 c10)中的 at。

小知識(shí)點(diǎn)：at, torch 與 c10這三個(gè)命名空間中 at 代表 ATen (A Tensor Library)，負(fù)責(zé)聲明和定義Tensor運(yùn)算相關(guān)的邏輯，是pytorch擴(kuò)展c++接口中最常用到的命名空間，c10 (Caffe Tensor Library)其實(shí)是 ATen 的基礎(chǔ)，包含了PyTorch的核心抽象、Tensor和Storage數(shù)據(jù)結(jié)構(gòu)的實(shí)際實(shí)現(xiàn)。torch 命名空間下定義的 Tensor 相比于ATen 增加自動(dòng)求導(dǎo)功能，但 c++ 擴(kuò)展中一般不常見）

該類型功能十分強(qiáng)大，基本支持 PyTorch 中 Tensor 的所有運(yùn)算方式（如 +, -, *, /, >, < 等運(yùn)算符，.view, .reshape, .unsqueeze等維度變化功能等）。Tensor 的 API 接口可見官方鏈接。當(dāng)然除了 Tensor 類型外 at 命名空間也支持幾乎所有和 Tensor 有關(guān)的函數(shù) （如 at::ones, at::zeros， at::where等）， ATen 的 API 接口可見官方鏈接。基本上只要在程序中加入了 #include <torch/extension.h> 就可以在 C++ 中調(diào)用所有 PyTorch 支持的功能。

4.2 CUDA 算子實(shí)現(xiàn)

4.2.1 (番外篇) CUDA 編程基礎(chǔ)

該部分內(nèi)容部分參考 CUDA編程入門極簡教程（https://blog.csdn.net/xiaohu2022/article/details/79599947），感興趣的同學(xué)可以看原文。

基本概念

CUDA 是建立在 NVIDIA GPU上的一個(gè)通用并行計(jì)算平臺(tái)和編程模型，CUDA編程可以利用GPUs 的并行計(jì)算引擎來更加高效地解決比較復(fù)雜的計(jì)算難題。CUDA 的語法和 C++ 大多部分情況下是一致的，其默認(rèn)文件名后綴是 .cu，默認(rèn)頭文件名后綴是 .cuh。CUDA 編程是異構(gòu)的，即CPU負(fù)責(zé)處理邏輯復(fù)雜的串行程序，而 GPU 重點(diǎn)處理數(shù)據(jù)密集型的并行計(jì)算程序，從而發(fā)揮最大功效。其中 CPU 所在位置稱為為主機(jī)端（host），而 GPU 所在位置稱為設(shè)備端（device）。

CUDA程序的設(shè)計(jì)流程

一般而言，CUDA 程序執(zhí)行會(huì)依照如下流程：

1. 分配 host 內(nèi)存，并進(jìn)行數(shù)據(jù)初始化

2. 分配 device 內(nèi)存，并從 host 將數(shù)據(jù)拷貝到 device 上

3. 調(diào)用 CUDA 的核函數(shù)在 device 上完成指定的運(yùn)算

4. 將 device 上的運(yùn)算結(jié)果拷貝到 host 上

5. 釋放 device 和 host 上分配的內(nèi)存

而對(duì) PyTorch 的 CUDA 擴(kuò)展來說， CUDA 擴(kuò)展傳入和傳出的 Tensor 都已經(jīng)在 GPU 上，因此這里的 5 個(gè)步驟只有第 3 步了，這會(huì)為我們省下比較寶貴的時(shí)間而將更多注意力放到具體的程序?qū)崿F(xiàn)上。

CUDA 中指定函數(shù)設(shè)備關(guān)鍵字

由于 CUDA 編程為異步，因此函數(shù)的定義與調(diào)用很可能不在同一個(gè) device 上面，因此 CUDA 中騎過增加額外函數(shù)類型來規(guī)約函數(shù)的定義與調(diào)用設(shè)備。- __global__：在 device 上執(zhí)行，從 host 中調(diào)用（一些特定的 GPU 也可以從 device 上調(diào)用），返回類型必須是 void，不支持可變參數(shù)參數(shù)，不能成為類成員函數(shù)。注意用__global__定義的 kernel 是異步的，這意味著 host 不會(huì)等待 kernel 執(zhí)行完就執(zhí)行下一步。- __device__：在 device 上執(zhí)行，單僅可以從 device 中調(diào)用，不可以和__global__同時(shí)用。- __host__：在 host 上執(zhí)行，僅可以從 host 上調(diào)用，一般省略不寫，不可以和__global__同時(shí)用，但可和__device__，此時(shí)函數(shù)會(huì)在 device 和 host 都編譯。

CUDA 中線程邏輯架構(gòu)形式

一旦一個(gè) kernel 在 device 上執(zhí)行，device 上很多經(jīng)量級(jí)的線程會(huì)被啟動(dòng)，一個(gè) kernel 所啟動(dòng)的所有線程分成兩級(jí)架構(gòu)。所有線程歸為一個(gè)網(wǎng)格（grid），同一個(gè)網(wǎng)格上的線程共享相同的全局內(nèi)存空間，而網(wǎng)格又可以分為很多線程塊（block），一個(gè)線程塊里面包含很多線程。線程兩層組織結(jié)構(gòu)如下圖所示，這是一個(gè) gird 和 block 均為 2-dim 的線程組織。grid 和 block 都是定義為 dim3 類型的變量，dim3 可以看成是包含三個(gè)無符號(hào)整數(shù)（x，y，z）成員的結(jié)構(gòu)體變量，在定義時(shí)，可定義為一維或二維，剩下維度缺少值為 1。當(dāng)然這里的 grid, block 層次劃分實(shí)際上只是邏輯層次，線程在 GPU 中的流處理器 (SM) 中是用“線程束”管理的，一個(gè)線程束包含 32 個(gè)線程。因此一般在設(shè)計(jì) block 時(shí)要保證其線程個(gè)數(shù)為 32 的整數(shù)倍。

為了更好地理解這里的線程架構(gòu)，我們可以直接將一個(gè)kernel 開辟的所有線程理解為一個(gè)小區(qū)，這個(gè)小區(qū)就被稱為 grid，而該小區(qū) (grid) 是由不同樓棟 (block)組成的，每個(gè)樓棟 (block)有其在小區(qū)內(nèi)的三維坐標(biāo) (x, y, z)。在每個(gè)樓棟中的所有線程按其在該 block 的三維坐標(biāo) (x, y, z)來進(jìn)行索引。

CUDA 中核函數(shù)調(diào)用

核函數(shù) (kernel) 是在 device上線程中并行執(zhí)行的函數(shù)，核函數(shù)用__global__符號(hào)聲明，在調(diào)用時(shí)需要用<<<grid, block>>>來指定 kernel 要執(zhí)行的線程數(shù)量。這里 grid 與 block 都需要提前定義好，在 CUDA 中，每一個(gè)線程都要執(zhí)行核函數(shù)，并且每個(gè)線程會(huì)分配一個(gè)唯一的線程號(hào) thread ID，這個(gè) ID 值可以通過核函數(shù)的內(nèi)置變量 threadIdx 來獲得。下面代碼即為在上圖線程邏輯架構(gòu)下的核函數(shù)調(diào)用方式。

dim3 grid(3, 2);
dim3 block(5, 3);
kernel_fun<<< grid, block >>>(prams...);

CUDA 中核函數(shù)編寫

核函數(shù)調(diào)用過程中將需要并行執(zhí)行的部分用不同的的線程進(jìn)行完成。因此在實(shí)際 CUDA 的核函數(shù)中，系統(tǒng)定義了兩個(gè)兩個(gè)內(nèi)置的坐標(biāo)變量blockIdx 與 threadIdx 來唯一標(biāo)識(shí)一個(gè)線程，它們都是 dim3 類型變量（包括x, y, z成員），其中blockIdx指明線程所在grid中的位置，而threaIdx指明線程所在block中的位置，這里 grid 與 block 正是在核函數(shù)調(diào)用過程中定義好的，在核函數(shù)的定義中也有 dim3 類型變量 gridDim與 blockDim 來分別指定 grid 與 block 的維度。如下核函數(shù)為矩陣相加的 CUDA 代碼。程序執(zhí)行過程中會(huì)按 blockIdx 與 threadIdx 的坐標(biāo)信息將該核函數(shù)分配給不同的線程來完成，因此實(shí)現(xiàn)高效并行化計(jì)算。以下為一個(gè)較為典型的矩陣相加的核函數(shù)設(shè)計(jì)。

__global__ void MatAdd(float A[N][N], float B[N][N], float C[N][N]) 
{ 
    int i = blockIdx.x * blockDim.x + threadIdx.x; 
    int j = blockIdx.y * blockDim.y + threadIdx.y; 
    if (i < N && j < N) 
        C[i][j] = A[i][j] + B[i][j]; 
}

至此我們完成了一般 CUDA 算子實(shí)現(xiàn)的基礎(chǔ)，在下一小節(jié)中我們?cè)賮矸治?nms CUDA 算子的實(shí)例。

4.2.2 CUDA 算子實(shí)例

// 以下程序適當(dāng)改寫，只顯示核心代碼
#include <torch/extension.h>using namespace at; 


#define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0))
int const threadsPerBlock = sizeof(unsigned long long int) * 8; // 64
__device__ inline bool devIoU(float const *const a, float const *const b,
                              const int offset, const float threshold) {
  // 定義在 device 上的函數(shù)，用于返回iou
  // 定義省略
}

__global__ void nms_cuda(const int n_boxes, const float iou_threshold,
                         const int offset, const float *dev_boxes,
                         unsigned long long *dev_mask) {
  // 核函數(shù)
  const int row_start = blockIdx.y;  // block 在 grid 中的位置
  const int col_start = blockIdx.x;  // block 在 grid 中的位置
  const int tid = threadIdx.x;  // thread ID (0-63)

  if (row_start > col_start) return;

  const int row_size =
      fminf(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
  const int col_size =
      fminf(n_boxes - col_start * threadsPerBlock, threadsPerBlock);

  //  Shared memory 只會(huì)被一個(gè) block 中的所有線程共享
  __shared__ float block_boxes[threadsPerBlock * 4];
  if (tid < col_size) {
    block_boxes[tid * 4 + 0] =
        dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 0];
    block_boxes[tid * 4 + 1] =
        dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 1];
    block_boxes[tid * 4 + 2] =
        dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 2];
    block_boxes[tid * 4 + 3] =
        dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 3];
  }
  __syncthreads(); 

  if (tid < row_size) {
    const int cur_box_idx = threadsPerBlock * row_start + tid;
    const float *cur_box = dev_boxes + cur_box_idx * 4;
    int i = 0;
    unsigned long long int t = 0;
    int start = 0;
    if (row_start == col_start) {
      start = tid + 1;  // 每個(gè) bbox 只需要和上三角元素計(jì)算（節(jié)省計(jì)算）
    }
    for (i = start; i < col_size; i++) {
      if (devIoU(cur_box, block_boxes + i * 4, offset, iou_threshold)) {
        t |= 1ULL << i;
      }
    }
    dev_mask[cur_box_idx * gridDim.y + col_start] = t;
  }
}

Tensor NMSCUDAKernelLauncher(Tensor boxes, Tensor scores, float iou_threshold,
                             int offset) {
  // cuda 程序入口
  at::cuda::CUDAGuard device_guard(boxes.device());  // 指定默認(rèn)的顯卡

  auto order_t = std::get<1>(scores.sort(0, /*descending=*/true));
  auto boxes_sorted = boxes.index_select(0, order_t);

  int boxes_num = boxes.size(0);
  // 這里將
  const int col_blocks = DIVUP(boxes_num, threadsPerBlock);
  // mask 是用來存儲(chǔ) bboxes 之間兩兩 iou 是否大于閾值的一個(gè) mask
  // 本來長度應(yīng)該是 (boxes_num, boxes_num)，但這里采用位存儲(chǔ)的方式，一個(gè) LongLong 類型可以存取 64 
  // 個(gè) bool 值，因此存儲(chǔ)空間可以縮小64倍，只用開辟 (boxes_num, boxes_num/64)長度即可。

  Tensor mask =
      at::empty({boxes_num, col_blocks}, boxes.options().dtype(at::kLong));
  dim3 blocks(col_blocks, col_blocks); 
  dim3 threads(threadsPerBlock);  // 每 64 個(gè)線程放到一個(gè) block 中，遍歷一個(gè)LongLong數(shù)的所有位
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  // 更完整的核函數(shù)調(diào)用 <<< blocks, threads, shared_memory, stream>>>
  nms_cuda<<<blocks, threads, 0, stream>>>(
      boxes_num, iou_threshold, offset, boxes_sorted.data_ptr<float>(),
      (unsigned long long*)mask.data_ptr<int64_t>());

  at::Tensor mask_cpu = mask.to(at::kCPU);
  unsigned long long* mask_host =
      (unsigned long long*)mask_cpu.data_ptr<int64_t>();

  std::vector<unsigned long long> remv(col_blocks);
  memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks);

  at::Tensor keep_t =
      at::zeros({boxes_num}, boxes.options().dtype(at::kBool).device(at::kCPU));
  bool* keep = keep_t.data_ptr<bool>();

  for (int i = 0; i < boxes_num; i++) {
    int nblock = i / threadsPerBlock;
    int inblock = i % threadsPerBlock;

    if (!(remv[nblock] & (1ULL << inblock))) {
      keep[i] = true;
      // set every overlap box with bit 1 in remv
      unsigned long long* p = mask_host + i * col_blocks;
      for (int j = nblock; j < col_blocks; j++) {
        remv[j] |= p[j];
      }
    }
  }

  AT_CUDA_CHECK(cudaGetLastError());
  return order_t.masked_select(keep_t.to(at::kCUDA));
}

推薦閱讀

谷歌提出Meta Pseudo Labels，刷新ImageNet上的SOTA！

"未來"的經(jīng)典之作ViT：transformer is all you need！

大道至簡！深度解讀CVPR2021論文RepVGG！

PyTorch 源碼解讀之 torch.autograd

漲點(diǎn)神器FixRes：兩次超越ImageNet數(shù)據(jù)集上的SOTA

Transformer為何能闖入CV界秒殺CNN？

SWA：讓你的目標(biāo)檢測模型無痛漲點(diǎn)1% AP

CondInst：性能和速度均超越Mask RCNN的實(shí)例分割模型

centerX: 用新的視角的方式打開CenterNet

mmdetection最小復(fù)刻版(十一)：概率Anchor分配機(jī)制PAA深入分析

MMDetection新版本V2.7發(fā)布，支持DETR，還有YOLOV4在路上!

CNN：我不是你想的那樣

TF Object Detection 終于支持TF2了！

無需tricks，知識(shí)蒸餾提升ResNet50在ImageNet上準(zhǔn)確度至80%+

不妨試試MoCo，來替換ImageNet上pretrain模型！

重磅！一文深入深度學(xué)習(xí)模型壓縮和加速

從源碼學(xué)習(xí)Transformer！

mmdetection最小復(fù)刻版(七)：anchor-base和anchor-free差異分析

mmdetection最小復(fù)刻版(四)：獨(dú)家yolo轉(zhuǎn)化內(nèi)幕

機(jī)器學(xué)習(xí)算法工程師

---

                                    一個(gè)用心的公眾號(hào)

---