[原始碼解析] PyTorch 如何使用GPU
0x00 摘要
在 PyTorch DataParallel 訓練過程中，其會在多個GPU之上覆制模型副本，然後才開始訓練。筆者在分析過程中，發現如果不把一些GPU相關基礎知識整理出來，很難理解DataParallel的這個複製模型的過程，遂有此文。
本系列其他文章如下：
[原始碼解析]深度學習利器之自動微分(3) --- 示例解讀
[原始碼解析]PyTorch如何實現前向傳播(1) --- 基礎類(上)
[原始碼解析]PyTorch如何實現前向傳播(2) --- 基礎類(下)
[原始碼解析] PyTorch如何實現前向傳播(3) --- 具體實現
[原始碼解析] Pytorch 如何實現後向傳播 (1)---- 呼叫引擎
[原始碼解析] Pytorch 如何實現後向傳播 (2)---- 引擎靜態結構
[原始碼解析] Pytorch 如何實現後向傳播 (3)---- 引擎動態邏輯
[原始碼解析] PyTorch 如何實現後向傳播 (4)---- 具體演算法
[原始碼解析] PyTorch 分散式(1)------歷史和概述
0x01 問題
在 DataParallel 進行前向傳播之前，需要在GPU之上分散資料，複製模型，具體可見下圖。
由此我們有幾個問題：
- 移動模型到GPU這個動作的背後究竟做了哪些操作？
- 如何在 CPU 之上呼叫 GPU 操作？
- 如何在 CPU，GPU 操作之間無縫切換？
- 是否需要把損失函式移動到 GPU 之上？
我們接下來就一一分析。
注，關於CUDA和Dispatcher我們只是大致介紹，目的是可以讓讀者走通整個流程，有興趣的讀者可以自行深入研究。
0x02 移動模型到GPU
2.1 cuda 操作
CUDA 是NVIDIA公司開發的GPU程式設計模型，其提供了GPU程式設計介面，使用者可以基於CUDA程式設計來構建基於GPU計算的應用。
torch.cuda用於設定 cuda 和執行cuda操作。它跟蹤當前選定的GPU，預設情況下，使用者分配的所有CUDA張量都將在該裝置上建立。使用者可以使用 torch.cuda.device 來修改所選裝置。一旦分配了張量，您可以對其執行操作，而不考慮所選裝置，PyTorch 會把執行結果與原始張量放在同一裝置上。
預設情況下，除了~torch.Tensor.copy_和其他具有類似複製功能的方法（如~torch.Tensor.to和~torch.Tensor.cuda）之外，不允許跨GPU操作，除非啟用對等（peer-to-peer）記憶體訪問。
我們從原始碼之中找出一個具體示例如下，大家可以看到，張量可以在裝置上被建立，操作。
cuda = torch.device('cuda') # Default CUDA device
cuda0 = torch.device('cuda:0')
cuda2 = torch.device('cuda:2') # GPU 2 (these are 0-indexed)
x = torch.tensor([1., 2.], device=cuda0)
# x.device is device(type='cuda', index=0)
y = torch.tensor([1., 2.]).cuda()
# y.device is device(type='cuda', index=0)
with torch.cuda.device(1):
# allocates a tensor on GPU 1
a = torch.tensor([1., 2.], device=cuda)
# transfers a tensor from CPU to GPU 1
b = torch.tensor([1., 2.]).cuda()
# a.device and b.device are device(type='cuda', index=1)
# You can also use ``Tensor.to`` to transfer a tensor:
b2 = torch.tensor([1., 2.]).to(device=cuda)
# b.device and b2.device are device(type='cuda', index=1)
c = a + b
# c.device is device(type='cuda', index=1)
z = x + y
# z.device is device(type='cuda', index=0)
# even within a context, you can specify the device
# (or give a GPU index to the .cuda call)
d = torch.randn(2, device=cuda2)
e = torch.randn(2).to(cuda2)
f = torch.randn(2).cuda(cuda2)
# d.device, e.device, and f.device are all device(type='cuda', index=2)
2.2 Module
深度學習的模型可以看做是一種引數的容器，執行模型其實就是對輸入引數做了一些基本的矩陣運算。一般來說，使用者定義的模型都是派生自 nn.modules.module 類。而分散式訓練涉及到同步更新引數和把模型拷貝到多個worker之上，所以我們首先需要看看Module的狀況。從定義中可以看出來，Module的成員變數主要分為狀態引數和hooks函式。
class Module:
dump_patches: bool = False
_version: int = 1
training: bool
_is_full_backward_hook: Optional[bool]
def __init__(self):
"""
Initializes internal Module state, shared by both nn.Module and ScriptModule.
"""
torch._C._log_api_usage_once("python.nn_module")
self.training = True
self._parameters = OrderedDict() # 在訓練過程中會隨著 BP 而更新的引數
self._buffers = OrderedDict() # 在訓練過程中不會隨著 BP 而更新的引數
self._non_persistent_buffers_set = set()
self._backward_hooks = OrderedDict()
self._is_full_backward_hook = None
self._forward_hooks = OrderedDict()
self._forward_pre_hooks = OrderedDict()
self._state_dict_hooks = OrderedDict()
self._load_state_dict_pre_hooks = OrderedDict()
self._modules = OrderedDict()
我們主要對狀態引數進行說明。狀態引數之中，主要有四種：
-
self.training
- 本網路是否正在訓練。
-
self._modules
- 是本網路下屬的子模組，採取迭代的方式進行定義。
-
self._parameters
- 網路的引數。是在訓練過程中會隨著 BP 而更新的引數，就是梯度更新的物件。
-
self._buffers
- 在訓練過程中，不會隨著BP更新的引數，但需要被儲存，比如BatchNorm中的moving mean and variance，其優化不是通過梯度反向傳播而是通過其他途徑。
從本質上講，當一個模型的網路結構被定義之後，self._parameters 和 self._buffers的組合是一個模型的具體狀態。如果需要拷貝一個模型：
self._modules屬於網路結構的一部分，當我們拷貝模型到其他workers時，會一起拷貝過來。- 而
self._parameters和self._buffers都需要顯式拷貝到其他worker，這樣才能在不同的Python程式之中維持相同的狀態。
那麼，這是不是意味著我們只需要拷貝 self._modules，self._parameters 和 self._buffers 這些就可以了？讓我們繼續往下看。
2.3 移動
2.3.1 示例
前面看到了如何在 GPU 上操作張量，我們接下來看看如何把模型放置到 GPU 之上。
首先我們定義了一個模型。
class ToyModel(nn.Module):
def __init__(self):
super(ToyModel, self).__init__()
self.net1 = nn.Linear(10, 10)
self.relu = nn.ReLU()
self.net2 = nn.Linear(10, 5)
def forward(self, x):
return self.net2(self.relu(self.net1(x)))
然後通過如下方式使用模型。
model = ToyModel().cuda(device_ids[0]) # 這裡複製模型到 GPU 之上
ddp_model = DDP(model, device_ids)
loss_fn = nn.MSELoss() # 接著進行訓練
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
optimizer.zero_grad()
outputs = ddp_model(torch.randn(20, 10))
labels = torch.randn(20, 5).to(device_ids[0])
loss_fn(outputs, labels).backward()
optimizer.step()
2.3.2 操作
示例之中使用了 cuda 方法把模型複製到 GPU 之上，註釋中指出了是把模型的 parameters 和 buffers 移動到 GPU 之上。程式碼中實際就是使用 self._apply 來呼叫 cuda(device)。
def cuda(self: T, device: Optional[Union[int, device]] = None) -> T:
r"""Moves all model parameters and buffers to the GPU.
This also makes associated parameters and buffers different objects. So
it should be called before constructing optimizer if the module will
live on GPU while being optimized.
.. note::
This method modifies the module in-place.
Args:
device (int, optional): if specified, all parameters will be
copied to that device
Returns:
Module: self
"""
return self._apply(lambda t: t.cuda(device))
我們再看大家熟悉的另外一些函式。
首先，to 方法其實本質也是使用 self._apply 來呼叫 to(device)，我們省略了一些檢驗程式碼。
def to(self, *args, **kwargs):
r"""Moves and/or casts the parameters and buffers.
This can be called as
.. function:: to(device=None, dtype=None, non_blocking=False)
.. function:: to(dtype, non_blocking=False)
.. function:: to(tensor, non_blocking=False)
.. function:: to(memory_format=torch.channels_last)
Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
floating point or complex :attr:`dtype`s. In addition, this method will
only cast the floating point or complex parameters and buffers to :attr:`dtype`
(if given). The integral parameters and buffers will be moved
:attr:`device`, if that is given, but with dtypes unchanged. When
:attr:`non_blocking` is set, it tries to convert/move asynchronously
with respect to the host if possible, e.g., moving CPU Tensors with
pinned memory to CUDA devices.
See below for examples.
.. note::
This method modifies the module in-place.
Args:
device (:class:`torch.device`): the desired device of the parameters
and buffers in this module
dtype (:class:`torch.dtype`): the desired floating point or complex dtype of
the parameters and buffers in this module
tensor (torch.Tensor): Tensor whose dtype and device are the desired
dtype and device for all parameters and buffers in this module
memory_format (:class:`torch.memory_format`): the desired memory
format for 4D parameters and buffers in this module (keyword
only argument)
Returns:
Module: self
"""
device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
def convert(t):
if convert_to_format is not None and t.dim() in (4, 5):
return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
non_blocking, memory_format=convert_to_format)
return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
return self._apply(convert)
其次，cpu 方法也是使用 self._apply 來呼叫 cpu(device)。
def cpu(self: T) -> T:
r"""Moves all model parameters and buffers to the CPU.
.. note::
This method modifies the module in-place.
Returns:
Module: self
"""
return self._apply(lambda t: t.cpu())
因此，我們需要分析一下 _apply 方法。
2.3.3 _apply 方法
我們可以看到其主要邏輯是：
- 遍歷 _parameters：
- 對引數呼叫fn進行處理，得到param_applied。
- 用 param_applied 重新設定引數。
- 如果引數有梯度，則：
- 對引數的grad呼叫fn進行處理，得到grad_applied。
- 用 grad_applied 重新設定引數的梯度。
- 對引數呼叫fn進行處理，得到param_applied。
- 遍歷 _buffers：
- 對buf呼叫fn進行處理。
def _apply(self, fn):
for module in self.children():
module._apply(fn)
def compute_should_use_set_data(tensor, tensor_applied):
if torch._has_compatible_shallow_copy_type(tensor, tensor_applied):
# If the new tensor has compatible tensor type as the existing tensor,
# the current behavior is to change the tensor in-place using `.data =`,
# and the future behavior is to overwrite the existing tensor. However,
# changing the current behavior is a BC-breaking change, and we want it
# to happen in future releases. So for now we introduce the
# `torch.__future__.get_overwrite_module_params_on_conversion()`
# global flag to let the user control whether they want the future
# behavior of overwriting the existing tensor or not.
return not torch.__future__.get_overwrite_module_params_on_conversion()
else:
return False
# 遍歷 _parameters
for key, param in self._parameters.items():
if param is not None:
# Tensors stored in modules are graph leaves, and we don't want to
# track autograd history of `param_applied`, so we have to use
# `with torch.no_grad():`
with torch.no_grad():
param_applied = fn(param) # 對引數呼叫fn進行處理，得到param_applied
should_use_set_data = compute_should_use_set_data(param, param_applied)
if should_use_set_data:
param.data = param_applied # 用 param_applied 重新設定
else:
assert isinstance(param, Parameter)
assert param.is_leaf
# # 用 param_applied 重新設定
self._parameters[key] = Parameter(param_applied, param.requires_grad)
if param.grad is not None: # 如果引數有梯度
with torch.no_grad():
grad_applied = fn(param.grad) # 對引數的grad呼叫fn進行處理
should_use_set_data = compute_should_use_set_data(param.grad, grad_applied)
if should_use_set_data:
param.grad.data = grad_applied # 用 grad_applied 重新設定
else:
assert param.grad.is_leaf
self._parameters[key].grad = grad_applied.requires_grad_(param.grad.requires_grad) # 用 grad_applied 重新設定
# 遍歷 _buffers
for key, buf in self._buffers.items():
if buf is not None:
self._buffers[key] = fn(buf) # 對buf呼叫fn進行處理
return self
因此我們可以看到，移動模型到GPU，其實就是把模型的self._parameters 和 self._buffers 移動到 GPU，並沒有對 self._modules 進行移動。我們對模型進行 .cuda() 處理，是將模型的引數放到視訊記憶體上去（實際使用的時候也是通過這些引數做運算）。
比如原來模型在下圖左側，進行 Module.cuda() 操作之後，模型如右邊所示。
+
|
+---------------------------------+ | +----------------------------------+
| CPU | | | CPU |
| +--------------+ | | | +--------------------+ |
| |Module | | | | | Module | |
| | | | | | | | |
| | _parameters+----> Parameters | | | | _parameters ------+ |
| | | | | | | | | |
| | _buffers +------> Buffers | | | +-----+ _buffers | | |
| | | | | | | | | | |
| | _modules | | | | | | _modules | | |
| | | | | | | | | | |
| +--------------+ | | | | +--------------------+ | |
| | | | | | |
+---------------------------------+ | +----------------------------------+
| | |
+ | |
+-------------------------------> Module.cuda() +---------------------------------> Time
+ | |
| | |
+---------------------------------+ | +----------------------------------+
| GPU | | | GPU | | |
| | | | | | |
| | | | | Parameters <-----+ |
| | | | | |
| | | | | |
| | | | +----> Buffers |
| | | | |
| | | | |
+---------------------------------+ | +----------------------------------+
|
+
為什麼 self._modules 沒有被移動？這是因為沒有必要，因為_modules 可以認為是一個list，其主要起到了橋樑作用，對其遞迴遍歷可以被用來獲取網路所有的 parameters。而這個功能在後續操作之中不是必須的。
DP 就是在每次網路傳播開始前，會把master節點上的parameters和buffer廣播給其他節點，以此來維持狀態的統一。
2.4 小結
現在我們可以回答了第一個問題：移動模型到GPU這個動作的背後究竟做了哪些操作？
答案時：呼叫 cuda 或者 to 方法來移動模型到GPU，其實就是把模型的self._parameters 和 self._buffers 移動到 GPU，並沒有對 self._modules 進行移動。這個移動過程是遞迴呼叫的，是把模型每個葉子都移動到了 GPU 之上。
0x03 在GPU之上呼叫函式
3.1 CUDA程式設計模型基礎
我們首先介紹一下CUDA程式設計模型基礎。
3.1.1 異構模型
CUDA程式設計模型是一個異構模型。程式執行在一個異構系統之上，這個異構系統由CPU和GPU構成，它們之間由匯流排分開，程式執行時候是由CPU和GPU協同工作。
在CUDA之中，有兩個重要概念：host和device。
-
Host ：CPU及其記憶體。
-
Device ：GPU及其記憶體。
因此，CUDA 架構下的一個程式也對應分為兩個部份：Host 程式碼和Device程式碼，它們分別在CPU和GPU上執行。host與device之間可以通訊進行資料拷貝。
- 主機程式碼（Host Code）：在 CPU 上執行的部份，使用Linux（GNU gcc）和Windows（Microsoft Visual C）編譯器來編譯。大致可以認為認為C語言工作物件是CPU和記憶體條。
- 裝置程式碼（Device Code）：在GPU上執行的部份，使用 NVIDIA NVCC 編譯器來編譯。大致可以認為 CUDA C工作物件是GPU及GPU上記憶體（也叫裝置記憶體）。
+-------------------+ +--------------------+
| | | |
| +----------+ | | +----------+ |
| | | | | | | |
| | RAM | | | | RAM | |
| | | | | | | |
| +----+-----+ | | +----+-----+ |
| | +--------+ | |
| | | | | |
| +----+-----+ | | +----+-----+ |
| | | | | | | |
| | CPU | | | | GPU | |
| | | | | | | |
| +----------+ | | +----------+ |
| | | |
+-------------------+ +--------------------+
Host Device
3.1.2 並行思想
CUDA 程式設計的思路是並行思想，大致如下：
- 把一個很大的執行任務劃分成若干個簡單的可以重複的操作，然後使用若干個執行緒來分別執行這些操作，達到並行的目的。
- 執行任務處理的資料也要對應分組成多個小資料塊。比如一個大資料分成若干個GPU組，每個GPU組要再次分成多個執行緒組，執行緒組內的張量可能需要再細分為張量處理器能處理的小組。
因此，一個典型的CUDA程式包括序列程式碼和並行程式碼。
- 序列程式碼是標準C程式碼，由host執行。
- 並行程式碼是CUDA C程式碼，在device中執行。
CUDA 主程式由CPU開始，即程式由host執行序列程式碼開始，當遇到需要資料並行處理的部分，則由device執行並行程式碼來作為補足。device可以獨立於host進行大部分操作。當一個device程式碼啟動之後，控制權會立刻返還給CPU來執行其他任務，所以這是一個非同步過程。
圖來自 https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html。
3.1.3 處理流程
典型的CUDA程式的執行流程如下：
- 分配host記憶體空間並且初始化資料。
- 分配device視訊記憶體空間。
- 將要計算的資料從Host記憶體之上覆制到device視訊記憶體之上。
- 呼叫CUDA核函式在device上完成使用者指定的運算。
- 將計算後GPU記憶體上的結果複製到Host記憶體上。
- 釋放device和host上分配的記憶體。
具體可以參見下圖。
3.2 函式
3.2.1 核函式
核函式是在device執行緒中並行執行的函式。在 CUDA 程式中，主程式在呼叫GPU核心之前需要對核進行執行配置，以確定執行緒塊數，每個執行緒塊中執行緒數和共享記憶體大小。比如在呼叫時需要用<<引數1,引數2>>來指定核函式需要的執行緒數量以及執行緒是如何組織，這樣在GPU之中就會啟動若干個執行緒來並行執行這個核函式，每個執行緒被分配一個唯一的執行緒號。
CUDA通過函式型別限定詞來區別host和device上的函式，主要的三個函式型別限定詞為：
| 限定符 | 執行 | 呼叫 | 備註 |
|---|---|---|---|
__global__ |
裝置端執行 | 可以從主機呼叫也可以從某些特定裝置呼叫 | 非同步操作，host 將平行計算任務發射到GPU的任務呼叫單之後，不會等待kernel執行完就執行下一步 |
__device__ |
裝置端執行 | 裝置端呼叫 | 不可以和__global__同時用 |
__host__ |
主機端執行 | 主機呼叫 | 可省略，不可和__global__同時用，可和__device__同時用，此時函式在device和host都編譯。 |
具體如下：
具體如下：
+------------------------+ +------------------------+
| | | |
| | | |
| __host__ __global__ | | __device__ |
| + + | | |
| | | | | + |
| | | | | | |
| | v---------------> | |
| | | | | | |
| | | | | | |
| | | | | | |
| | | | | | |
| | | | | | |
| | | | | | |
| | | | | | |
| | | | | | |
| | | | | | |
| | +<--------------v | |
| | | | | | |
| | | | | | |
| | | | | | |
| v v | | v |
| | | |
+------------------------+ +------------------------+
Host Device
這三個限定詞其實也是 CUDA 中常見的三種執行場景。其中，device 函式和global函式因為需要在GPU上執行，因此不能呼叫常見的一些 C/C++ 函式（因為這些函式沒有對應的 GPU 實現）。
如下程式碼是 NVIDIA 的例子，使用內建的 threadIdx 變數，把 A 和 B 兩個張量進行相加，得到 C。因此，N 個執行緒之中每個都會執行 VecAdd() 。
// Kernel definition
__global__ void VecAdd(float* A, float* B, float* C)
{
int i = threadIdx.x;
C[i] = A[i] + B[i];
}
int main()
{
...
// Kernel invocation with N threads
VecAdd<<<1, N>>>(A, B, C);
...
}
3.2.2 PyTorch 樣例
我們從 third_party/cub/cub/device/dispatch/dispatch_reduce.cuh 找一個核函式例子來看看。
/**
* Reduce region kernel entry point (multi-block). Computes privatized reductions, one per thread block.
*/
template <
typename ChainedPolicyT, ///< Chained tuning policy
typename InputIteratorT, ///< Random-access input iterator type for reading input items \iterator
typename OutputIteratorT, ///< Output iterator type for recording the reduced aggregate \iterator
typename OffsetT, ///< Signed integer type for global offsets
typename ReductionOpT> ///< Binary reduction functor type having member <tt>T operator()(const T &a, const T &b)</tt>
__launch_bounds__ (int(ChainedPolicyT::ActivePolicy::ReducePolicy::BLOCK_THREADS))
__global__ void DeviceReduceKernel(
InputIteratorT d_in, ///< [in] Pointer to the input sequence of data items
OutputIteratorT d_out, ///< [out] Pointer to the output aggregate
OffsetT num_items, ///< [in] Total number of input data items
GridEvenShare<OffsetT> even_share, ///< [in] Even-share descriptor for mapping an equal number of tiles onto each thread block
ReductionOpT reduction_op) ///< [in] Binary reduction functor
{
// The output value type
typedef typename If<(Equals<typename std::iterator_traits<OutputIteratorT>::value_type, void>::VALUE), // OutputT = (if output iterator's value type is void) ?
typename std::iterator_traits<InputIteratorT>::value_type, // ... then the input iterator's value type,
typename std::iterator_traits<OutputIteratorT>::value_type>::Type OutputT; // ... else the output iterator's value type
// Thread block type for reducing input tiles
typedef AgentReduce<
typename ChainedPolicyT::ActivePolicy::ReducePolicy,
InputIteratorT,
OutputIteratorT,
OffsetT,
ReductionOpT>
AgentReduceT;
// Shared memory storage
__shared__ typename AgentReduceT::TempStorage temp_storage;
// Consume input tiles
OutputT block_aggregate = AgentReduceT(temp_storage, d_in, reduction_op).ConsumeTiles(even_share);
// Output result
if (threadIdx.x == 0)
d_out[blockIdx.x] = block_aggregate;
}
3.3 小結
目前我們知道了，PyTorch 其實可以通過呼叫 __global__ 方法來在GPU之上執行並行操作。這回答了我們的第二個問題：如何在 CPU 之上呼叫 GPU 操作？
0x04 在GPU/CPU之間切換
我們接下來分析如何在GPU/CPU之間切換。
由示例程式碼可以知道，只要呼叫了 cuda 函式把模型移動到 GPU 之上，我們就可以使用 CUDA global 核函式在GPU上進行並行運算。
model = ToyModel().cuda(device_ids[0]) # 這裡複製模型到 GPU 之上
ddp_model = DDP(model, device_ids)
loss_fn = nn.MSELoss()
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
optimizer.zero_grad()
outputs = ddp_model(torch.randn(20, 10))
但是我們忽略了一個問題，就是 PyTorch 怎麼知道此時應該呼叫GPU對應的 global 核函式？為什麼 PyTorch 就不呼叫 CPU 函式或者其他裝置的函式了？這就是我們接下來需要分析的。
4.1 Dispatcher 機制
此處我們主要借鑑 http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/。
4.1.1 問題
在PyTorch中，operator 所表現出預期行為是由很多機制共同作用導致的，比如：
- 做實際工作的kernel。
- 是否支援反向自動微分，例如，使 loss.backward() 正常工作的標記位。
- 是否啟用了torch.jit.trace。
- 如果你正在vmap呼叫中，所執行operator將會表現出不同的批處理行為。
因此，我們知道有太多不同的方式可以對PyTorch operator進行不同的解釋，如果我們試圖在一個名為add的單一函式裡面處理所有的行為，我們的實現程式碼會很快演變成一個不可維護的混亂局面。
所以我們需要有一個機制來解決這個問題，這個機制不僅僅是一個if語句這麼簡單，而是PyTorch內部一個非常重要的抽象，而且它必須在儘可能不降低PyTorch效能的情況下做到這一點。這個機制就是 Dispatcher。
4.1.2 什麼是 Dispatcher
什麼是dispatcher？dispatcher對於每個operator都會維護一個函式指標表，這些函式為每個dispatch key提供了對應的實現，這套機制大致對應於PyTorch中的一個橫切關注點。在上圖中，你可以看到在這個表中有針對不同後端（CPU、CUDA、XLA）以及更高階概念（例如 autograd 和跟蹤）的dispatch條目。dispatcher的工作是根據輸入的tensor和其他一些東西來計算出一個dispatch key，然後跳轉到函式指標表所指向的函式。
熟悉 C++ 的人可能會注意到，這個函式指標表與C++中的虛表非常相似。在C++中，物件的虛擬函式是通過將每個物件與一個虛表的指標相關聯來實現的，該虛表包含了有關物件上每個虛擬函式的實現。在PyTorch中，我們基本上重新實現了虛擬表，但有一些區別。
- dispatch表是按operator分配的，而虛表是按類分配的。這意味著我們可以通過分配一個新的dispatch表來擴充套件所支援的operator集。與其不同的是，對於一個C++物件，你可以通過繼承子類來擴充套件型別，但你不能輕易新增虛擬函式。與普通的物件導向系統不同，PyTorch大部分的可擴充套件性在於定義新的operator（而不是新的子類），所以這種權衡是合理的。此外，dispatch key的種類不是公開可擴充套件的，我們希望那些想新增新dispatch key的使用者通過向PyTorch核心團隊提交一個補丁來新增他們的dispatch key。
- 我們的dispatch key的計算考慮了operator的所有引數（multiple dispatch）以及執行緒本地狀態（TLS）。這與虛表不同，在虛表中只有第一個物件（this指標）很重要。
- 最後，dispatcher支援boxing和unboxing作為op的呼叫約定的一部分。在文章的最後部分會有更多關於這個的內容。
有趣的歷史筆記：我們曾經使用虛擬函式來實現動態dispatch，當我們意識到需要比虛表更多的能力時，我們重新實現了動態dispatch。
4.1.3 如何計算key
那麼，我們究竟是如何計算dispatch key的呢？我們是基於dispatch key set來完成的，dispatch key set是一個基本抽象，它是dispatch key的一個bitset。大致來講，我們綜合來自不同來源的dispatch key sets（在某些情況下遮蔽一些key）來得到一個最終的dispatch key set。然後我們在這個set中挑選優先順序最高的key（dispatch keys按某些優先順序隱式排序），這就是我們這次應該呼叫的結果。那麼，這些dispatch key sets的來源是什麼？
- 每個張量輸入都有一個由該張量上的所有dispatch key組成的dispatch key set（直觀地說，這些dispatch key的值會是類似 “CPU”字串這樣的東西，這告訴我們該張量是一個CPU張量，所以應該由dispatch表中的CPU handler來處理）。
- 我們還有一個local include set，用於 "模態（modal） "功能，例如tracing，它不與任何張量關聯，而是某種執行緒的本地模態，使用者可以在某些範圍內開啟或關閉。
- 最後，我們有一個global set，它包含了始終應該被考慮的dispatch key（自從寫下這張PPT以來，Autograd已經從global set轉移到了張量之上。然而系統的高階結構並沒有改變）。
除了這些，還有一個local exclude set，其用從dispatch排除某些dispatch key。一個常見的場景是一個handler負責處理一個key，然後通過local exclude set將自己遮蔽掉，這樣我們以後就不會嘗試重新處理這個key。
4.1.4 註冊
我們接下來看看如何註冊這個dispatch key 到 dispatch 表之中。這個過程通過operator registration API來實現。操作符註冊 API 有三種主要方式：
- 為operator定義模式。
- 然後在對應的key上註冊實現。
- 最後，有一個 fallback 方法，使用者可以使用它為某個key對應的所有運算子定義同一個處理程式。
為了視覺化 operator registration的工作，讓我們想象一下，所有op的dispatch表共同形成一個二維網格，像這樣：
- 縱軸上是PyTorch中支援的每個op。
- 橫軸上是系統支援的每個dispatch key。
operator registration 行為就是在這兩個軸定義出的單元格中填寫對應的實現。
在一個特定的dispatch key上為一個operator註冊kernel函式時，我們會填寫一個單元格（下面的藍色）的內容。
4.2 Dispatcher 程式碼
我們接下來通過原始碼來看看。
4.2.1 虛擬函式表
4.2.1.1 例子
我們可以從 aten/src/ATen/native/native_functions.yaml 之中找到一些虛擬函式的例子。
# zero 操作對應的虛擬函式表
- func: zero_(Tensor(a!) self) -> Tensor(a!)
device_check: NoCheck # TensorIterator
variants: method, function
dispatch:
CPU, CUDA: zero_
Meta: zero_meta_
SparseCPU, SparseCUDA: zero_sparse_
MkldnnCPU: mkldnn_zero_
# sub.out 對應的虛擬函式表
- func: sub.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)
device_check: NoCheck # TensorIterator
structured: True
structured_inherits: TensorIteratorBase
dispatch:
CPU, CUDA: sub_out
SparseCPU, SparseCUDA: sub_out_sparse
# sub.Tensor 對應的虛擬函式表
- func: sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
device_check: NoCheck # TensorIterator
variants: function, method
structured_delegate: sub.out
dispatch:
SparseCPU, SparseCUDA: sub_sparse
4.2.1.2 Operator的實現
我們可以看看 zero 的兩個實現，下面是MkldnnCPU的實現。
Tensor& mkldnn_zero_(Tensor& self) {
using Vec = vec::Vectorized<float>;
ideep::tensor& x = itensor_from_mkldnn(self);
auto n = x.get_nelems();
auto* x_ = static_cast<float*>(x.get_data_handle());
parallel_for(0, n, 2048, [x_](int64_t begin, int64_t end) {
vec::map(
[](Vec /* unused */) { return 0.0; },
x_ + begin,
x_ + begin,
end - begin);
});
return self;
}
又比如下面是SparseCPU, SparseCUDA 的對應實現：
// --------------------------------------------------------------------
// zero_(SparseTensor)
// --------------------------------------------------------------------
// hummu hummu
SparseTensor& zero_sparse_(SparseTensor& self) {
AT_ASSERT(self.is_sparse());
at::zeros_out(self, get_sparse_impl(self)->sizes());
return self._coalesced_(true);
}
4.2.2 Dispatcher 定義
我們接下來看看Dispatcher的定義，這裡只給出部分成員變數。
class TORCH_API Dispatcher final {
private:
// For direct access to backend fallback information
friend class impl::OperatorEntry;
struct OperatorDef final {
explicit OperatorDef(OperatorName&& op_name)
: op(std::move(op_name)) {}
impl::OperatorEntry op;
size_t def_count = 0;
size_t def_and_impl_count = 0;
};
friend class OperatorHandle;
template<class> friend class TypedOperatorHandle;
public:
static Dispatcher& realSingleton();
//儲存所有的運算元，並在其成員變數中儲存了每個運算元的不同版本，比如cpu，cuda，autograd....
std::list<OperatorDef> operators_;
//註冊運算元時會將運算元名稱和方法也儲存在這個裡面, 這樣就可以快速的通過名字查詢到運算元方法(其中包含了成員OperatorDef)
LeftRight<ska::flat_hash_map<OperatorName, OperatorHandle>> operatorLookupTable_;
// Map from namespace to debug string (saying, e.g., where the library was defined)
ska::flat_hash_map<std::string, std::string> libraries_;
std::array<impl::AnnotatedKernel, static_cast<uint8_t>(DispatchKey::NumDispatchKeys)> backendFallbackKernels_;
std::unique_ptr<detail::RegistrationListenerList> listeners_;
std::mutex mutex_;
};
4.2.3 註冊
我們接下來給出註冊虛擬函式表的方法。
RegistrationHandleRAII Dispatcher::registerImpl(
OperatorName op_name,
c10::optional<DispatchKey> dispatch_key,
KernelFunction kernel,
c10::optional<impl::CppSignature> cpp_signature,
std::unique_ptr<FunctionSchema> inferred_function_schema,
std::string debug
) {
std::lock_guard<std::mutex> lock(mutex_);
auto op = findOrRegisterName_(op_name);
auto handle = op.operatorDef_->op.registerKernel( // 進行註冊
*this,
dispatch_key,
std::move(kernel),
std::move(cpp_signature),
std::move(inferred_function_schema),
std::move(debug)
);
++op.operatorDef_->def_and_impl_count;
return RegistrationHandleRAII([this, op, op_name, dispatch_key, handle] {
deregisterImpl_(op, op_name, dispatch_key, handle);
});
}
4.2.3.1 登錄檔
OperatorEntry代表了一個運算元，以及該運算元的dispatch table，這裡只給出成員變數。
class TORCH_API OperatorEntry final { //代表了一個運算元，以及該運算元的dispatch table
public:
OperatorName name_;
c10::optional<AnnotatedSchema> schema_;
//儲存了不同key對應的運算元實現版本，比如cpu，cuda，autograd 等等，所有的運算元版本都會在這個table裡面
std::array<KernelFunction, static_cast<uint8_t>(DispatchKey::NumDispatchKeys)> dispatchTable_;
DispatchKeyExtractor dispatchKeyExtractor_;
//不同 DispatchKey對應了不同的版本的kernel運算元實現版本
ska::flat_hash_map<DispatchKey, std::list<AnnotatedKernel>> kernels_;
};
4.2.3.2 註冊行為
最終註冊行為就是往 dispatchTable_ 之中設定。
void OperatorEntry::updateDispatchTableEntry_(const c10::Dispatcher& dispatcher, DispatchKey dispatch_key) {
auto dispatch_ix = static_cast<uint8_t>(dispatch_key);
dispatchTable_[dispatch_ix] = computeDispatchTableEntry(dispatcher, dispatch_key);
dispatchKeyExtractor_.setOperatorHasFallthroughForKey(dispatch_key, dispatchTable_[dispatch_ix].isFallthrough());
}
4.2.4 如何dispatch
4.2.4.1 排程依據
PyTorch 之中會依據dtype、device和layout的不同來排程不同的operator。
- 大多數型別（比如int32）可以使用模版方式直接進行對映，但是某些operator 不支援模版功能，就需要dispatcher這樣的動態排程器。
- PyTorch的tensor不僅可以執行在CPU上，還可以跑在GPU，mkldnn和xla等裝置，這也需要動態排程。
- layout是指tensor中元素的排布，這就有strided layout和sparse layout的區別，所以也需要動態排程。
4.2.4.2 排程程式碼
我們這裡這是給出部分程式碼，有興趣的讀者繼續繼續深入。
template<class Return, class... Args>
C10_DISPATCHER_INLINE_UNLESS_MOBILE Return Dispatcher::call(const TypedOperatorHandle<Return(Args...)>& op, Args... args) const {
detail::unused_arg_(args...);
// 得到key set
auto dispatchKeySet = op.operatorDef_->op.dispatchKeyExtractor()
.template getDispatchKeySetUnboxed<Args...>(args...);
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(!c10::isAliasDispatchKey(dispatchKeySet.highestPriorityTypeId()));
// 得到運算元
const KernelFunction& kernel = op.operatorDef_->op.lookup(dispatchKeySet.highestPriorityTypeId());
// 進行排程
#ifndef PYTORCH_DISABLE_PER_OP_PROFILING
bool pre_sampled = false;
if (C10_UNLIKELY(at::shouldRunRecordFunction(&pre_sampled))) {
return callWithDispatchKeySlowPath<Return, Args...>(op, pre_sampled, dispatchKeySet, kernel, std::forward<Args>(args)...);
}
#endif // PYTORCH_DISABLE_PER_OP_PROFILING
return kernel.template call<Return, Args...>(op, dispatchKeySet, std::forward<Args>(args)...);
}
4.2.4.3 key
我們接下來看看key的定義，因為太多，所以我們只給出部分數值。
enum class DispatchKey : uint8_t {
CPU, // registered at build/aten/src/ATen/RegisterCPU.cpp
CUDA, // registered at build/aten/src/ATen/RegisterCUDA.cpp
HIP, // NB: I think this is not actually used, due to Note [Masquerading as
// CUDA]
FPGA, // Xilinx support lives out of tree at
// https://gitlab.com/pytorch-complex/vitis_kernels
MSNPU, // unused externally, but tested at
// test/cpp_extensions/msnpu_extension.cpp
XLA, // lives out of tree at https://github.com/pytorch/xla
MLC, // lives out of tree at https://github.com/pytorch/MLCompute
Vulkan,
Metal,
XPU, // For out of tree Intel's heterogeneous computing plug-in
HPU, // For out of tree & closed source integration of HPU / Habana
VE, // For out of tree & closed source integration of SX-Aurora / NEC
Lazy, // For lazy tensor backends
// A meta tensor is a tensor without any data associated with it. (They
// have also colloquially been referred to as tensors on the "null" device).
// A meta tensor can be used to dry run operators without actually doing any
// computation, e.g., add on two meta tensors would give you another meta
// tensor with the output shape and dtype, but wouldn't actually add anything.
Meta,
// Here are backends which specify more specialized operators
// based on the dtype of the tensor.
QuantizedCPU, // registered at build/aten/src/ATen/RegisterQuantizedCPU.cpp
QuantizedCUDA, // registered at build/aten/src/ATen/RegisterQuantizedCUDA.cpp
QuantizedXPU, // For out of tree Intel's heterogeneous computing plug-in
// This backend is to support custom RNGs; it lets you go
// to a different kernel if you pass in a generator that is not a
// traditional CPUGeneratorImpl/CUDAGeneratorImpl. To make use of this
// key:
// 1) set it as a second parameter of at::Generator constructor call in
// the user-defined PRNG class.
// 2) use it as a dispatch key while registering custom kernels
// (templatized kernels specialized for user-defined PRNG class)
// intended for out of tree use; tested by aten/src/ATen/test/rng_test.cpp
CustomRNGKeyId,
// Here are backends which specify more specialized operators
// based on the layout of the tensor. Note that the sparse backends
// are one case where ordering matters: sparse multi-dispatches with
// the corresponding dense tensors, and must be handled before them.
MkldnnCPU, // registered at build/aten/src/ATen/RegisterMkldnnCPU.cpp
// NB: not to be confused with MKLDNN, which is Caffe2 only
SparseCPU, // registered at build/aten/src/ATen/RegisterSparseCPU.cpp
SparseCUDA, // registered at build/aten/src/ATen/RegisterSparseCUDA.cpp
SparseHIP, // TODO: I think this is not actually used, due to Note
// [Masquerading as CUDA]
SparseXPU, // For out of tree Intel's heterogeneous computing plug-in
SparseVE, // For out of tree & closed source integration of SX-Aurora / NEC
SparseCsrCPU,
SparseCsrCUDA,
AutogradOther,
AutogradCPU,
AutogradCUDA,
AutogradXLA,
AutogradLazy,
AutogradXPU,
AutogradMLC,
AutogradHPU,
......
};
4.2.4.4 key的使用
因為篇幅所限，我們無法深入分析每一種情況，這裡只給出從 DeviceType 出發的情景。我們從下面函式可以看到，如何從 DeviceType 對映到 DispatchKey 型別。
template <typename Func>
inline CppFunction dispatch(c10::DeviceType type, Func&& raw_f) {
auto deviceTypeToDispatchKey = [](c10::DeviceType t){
switch (t) {
// This list is synchronized with the k-constants in c10/core/DeviceType.h
case c10::DeviceType::CPU:
return c10::DispatchKey::CPU;
case c10::DeviceType::CUDA:
return c10::DispatchKey::CUDA;
case c10::DeviceType::XLA:
return c10::DispatchKey::XLA;
case c10::DeviceType::Lazy:
return c10::DispatchKey::Lazy;
case c10::DeviceType::MLC:
return c10::DispatchKey::MLC;
case c10::DeviceType::Meta:
return c10::DispatchKey::Meta;
case c10::DeviceType::HIP:
return c10::DispatchKey::HIP;
case c10::DeviceType::MSNPU:
return c10::DispatchKey::MSNPU;
case c10::DeviceType::HPU:
return c10::DispatchKey::HPU;
default:
TORCH_CHECK(false,
"Device type ", t, " cannot be overloaded at dispatch time, "
"please file a bug report explaining what you were trying to do.");
}
};
return dispatch(deviceTypeToDispatchKey(type), std::forward<Func>(raw_f));
}
4.3 小結
至此，我們知道，通過 Dispatcher 機制，PyTorch 可以依據dtype、device和layout的不同來排程不同的operator。這就解答了我們第三個問題：如何在 CPU，GPU 操作之間無縫切換？
關於第四個問題：是否需要把損失函式移動到 GPU 之上？，我們也有了解答：
損失函式的引數是前向傳播的outputs和label，outputs已經在GPU之上（因為訓練資料已經在GPU之上），label 也被使用者手動設定到GPU之上。所以損失函式的引數都已經在GPU之上，這樣 Dispather 就依據device會呼叫到GPU對應的operator，所以不需要把損失函式移動到GPU之上。
我們整理一個總體邏輯如下，序列是：
- 把訓練資料 inputs 移動到GPU。
- 進行前向操作，假設只有一個operator，就是 op1，使用 device='GPU' 這個 dispatch key 去 Dispatcher 查詢。
- 找到了 op1-gpu 這個operator，進行計算，得出 outputs。
- outputs 就自動存在於 GPU 之上。
- 把 Labels 也放到 GPU 之上。
- 進行損失函式運算，假設只有一個 operator，就是 op2，此時損失函式的引數都在GPU之上，所以使用 device= 'GPU' 這個 dispatch key 去 Dispatcher 查詢。
- 找到了 op2-gpu 這個operator，進行計算，得出 loss。
+--------------------+
+-----------+ | Forward | +------------+ +------------------+
| GPU | | | | GPU | | Loss Function |
| +---> | op1 op1-gpu() +----> | +---> | | +--------+
| Inputs | 1 | | 4 | Outputs | | | | GPU |
| | | + ^ | | | | | | |
+-----------+ | | | | +------------+ | op2 op2-gpu() +-->+ loss |
| | | | | | | |
+--------------------+ +------------+ | + ^ | | |
| | | GPU | 5 | | | | +--------+
| | | +---> | | 6 | 7 |
2 | | 3 | Labels | | | | |
| | | | | | | |
| | +------------+ +------------------+
+----------------------------+ +--------------------------------+ | |
| | | |
+-----------------------------------------------------------------------------+ |
| | | |
| | +-------------------------------------------------------+ | |
| | | Dispather | | |
| | | + + + + | | |
| | | | XLA | CPU | Metal | GPU | | |
| | | +---------------------------------------------------+ | | |
| | | | | | | | | |
| +--------> | OP1 | op1-xla | op1-cpu | op1-metal | op1-gpu +---+ |
| 'device=GPU' | | | | | +------+ | |
| | +---------------------------------------------------+ | |
| | | | | | | |
+------------> | OP2 | op2-xla | op2-cpu | op2-metal | op2-gpu +---------------+
'device=GPU' | | | | | +------+ |
| +---------------------------------------------------+ |
| | | | | |
| OP3 | op3-xla | op3-cpu | op3-metal | op3-gpu |
| | | | | |
| +---------------------------------------------------+ |
+-------------------------------------------------------+
手機如下：
至此，GPU相關分析結束，下一篇我們開始分析DataParallel，敬請期待。
0xFF 參考
http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/
https://pytorch.org/tutorials/advanced/dispatcher.html
https://blog.csdn.net/qq_23858785/article/details/96476740
https://blog.csdn.net/weixin_42236014/article/details/116747358
https://blog.csdn.net/crazy_sunshine/article/details/97920534
深入淺出全連線層（fully connected layer）
Pytorch擴充進階(二)：Pytorch結合C++以及Cuda擴充
Pytorch擴充進階(一)：Pytorch結合C以及Cuda語言