3DGS前向渲染(2)CUDA光栅化
<摘要>
光栅化是计算机图形学中，将连续的三维几何场景转换为离散的屏幕像素的图形处理管线；高斯泼溅中为了加速光栅化的渲染过程，基于 torch 的 CUDA 接口提供了 GPU 并行的编程实现。
GPU 并行编程基础
为了理清这部分的逻辑，有必要先介绍最基本的 GPU 并行编程概念。GPU 上并行执行的指令称为核函数；CUDA 在逻辑访问上的层级结构分为 grid、block 和 thread 三级。grid 下分为多个 block，一个 block 中包含多个 thread，一个 block 内可以使用共享内存(共享内存是 GPU 中比全局内存访问速度快得多的一种存储结构，但大小十分有限)。block size 表示一个 block 内含有的 thread 数，block dim 表示 block 内线 thread 的组织结构。组织结构可以这样理解，把一个 thread 想象成一个边长为 1 的小立方体，block dim = (4, 3, 2)表示立方体排列成长宽高分别为 4、3、2 的长方体，如下图所示。组织结构支持一维到三维，未指定的维度默认设置为 1。类似地，gird size 表示一个 grid 内含有的 block 数，grid dim 表示 grid 内 block 的组织结构。核函数以整个 gird 为范围、以 thread 为单位执行，即每个线程都会运行一次核函数，多个线程同时运行核函数即可实现并行数据处理。
并行渲染策略
高斯泼溅中将图像按分辨率向上取整划分为多个 tile，每个 tile 包含 16*16=256 个像素。整个并行逻辑分为 preprocess 和 render 两部分，preprocess 部分每个 thread 负责一个高斯椭球，block size 设置为(256, 1, 1)，grid size 根据高斯椭球总数进行设置；render 部分每个 thread 负责一个像素，且每个 block 负责一个 tile，block size 设为(16, 16, 1)，与 tile 中的像素对应，grid size 根据划分出的 tile 数进行设置。
例如，图像分辨率为 1920*1080，高斯总数为 30000，\(\left\lceil x \right\rceil\)为向上取整的符号，上述过程计算如下：
\[\begin{aligned} &\text{preprocess}部分的\text{block size}为: (256, 1, 1)\\ &\text{preprocess}部分的\text{grid size}为: \left\lceil\frac{30000}{256}\right\rceil = 118 \Rightarrow (118, 1, 1)\\ &\text{render}部分的\text{block size}为: (16, 16, 1)\\ &划分的\text{tile}数: \left\lceil \frac{1920}{16} \right\rceil \times \left\lceil \frac{1080}{16} \right\rceil = 120 \times 68 = 8160\\ &\text{render}部分的\text{grid size}为: (120, 68, 1)\\ \end{aligned}\]在 preprocess 部分，以某一个 thread 为例，按照前面介绍的高斯椭球转化过程，计算得到像素坐标系下的高斯椭圆，即 xy 位置和距离相机的深度、2D 协方差矩阵、颜色和不透明度。接下来需要得到这个高斯椭圆所涉及的 tile 索引范围和数量，为此对二维协方差进行特征值分解，取最大特征值的三倍（即高斯函数中的\(3\sigma\)）这个范围的 AABB（Axis Aligned Bounding Box，轴对齐包围盒，是边与坐标轴平行的外切矩形，可由一组点中的最小坐标和最大坐标确定）作为该高斯椭圆的影响区域，记录该区域涉及的 tile 索引范围且可以据此得出高斯椭圆影响的 tile 数量。
将所有高斯椭圆涉及的 tile 数求和，就是实际要处理的渲染任务总数。在高斯泼溅中，这些椭圆按照到相机的距离从近到远的顺序渲染到图像上，完成全部高斯椭球的处理后的结构类似于下图。高斯椭圆和 tile 之间是多对多的关系，即一个高斯椭圆可能影响多个 tile，一个 tile 也可能受多个高斯椭圆影响。
为了处理这种多对多的关系并实现这种按深度顺序进行渲染的逻辑，构造tile id | depth <--> gaussian id的键值对，将 tile 索引置于键的高位，深度置于键的低位，高斯椭圆的索引作为值构成一条 entry。这样的 entry 假定一个高斯椭圆影响一个 tile，根据高斯椭圆影响 tile 的数量对其进行复制，如某个高斯椭圆影响 5 个 tile，则将其复制 5 份并分别构造 entry。注：这里的复制只是一种理解方式，指的是重复使用这个高斯索引来与它影响的不同 tile 分别构成 entry，而不是在内存上将高斯参数进行复制。对这个 entry 表进行键值对的排序，就可以将高斯索引按 tile 索引排列好，同一个 tile 内按深度排列好，举例如下图。
注：高斯泼溅在构建这样的 entry 时其实还会根据高斯椭圆对 tile 的贡献值进行排序，小于一定阈值则认为影响可以忽略，并不会构造这个 entry。具体来说是，对某一个椭圆影响的某一个 tile，按照椭圆梯度的方法计算这个 tile 内高斯函数的最大值，大于阈值才认为这个椭圆对这个 tile 有贡献，才会构建这个 entry。
排序完成之后，就获得了每个 tile 上需要处理的高斯椭圆的索引，接下来进入每个 thread 负责一个像素、每个 block 负责一个 tile 的 render 部分。在一个 tile 内，为了利用 GPU 的共享内存机制，对所有的高斯椭圆进行分轮次处理，每一轮处理前都先将数据读入共享内存，高斯泼溅中将每一轮处理的高斯数量设为 256。在一个轮次内，遍历这 256 个高斯椭圆，计算它们对这个 tile 内所有像素的影响。对某一个高斯椭圆来说，一个线程处理一个像素，整个 block 内的 thread 并行即可处理完某一个高斯椭圆对整个 tile 的影响。由此经过轮次的循环和高斯椭圆的循环，即可处理完整个 tile 上所有高斯椭圆对这个 tile 中所有像素的影响。
\(\alpha\)-blending 图像合成
一个轮次中的 256 个高斯椭圆对这个 tile 中某一个像素的颜色累积过程称为\(\alpha-\text{blending}\)。每个椭圆参与合成的实际\(\alpha\)，由像素中心和高斯函数中心之间的马氏距离进行椭圆平均加权。
椭圆加权平均(EWA)是一种高质量的采样和抗锯齿算法，适合处理椭圆形状分布的元素，核心思想是基于二维协方差的马氏距离。马氏距离是一种距离指标，可以应对高维线性分布的数据中各维度间非独立同分布的问题，修正了欧式距离中各维度尺度不一致且相关的问题，本质上是对 PCA 中的主成分进行标准化，将变量按主成分旋转(主成分就是特征向量的方向)，使各维度之间相互独立，然后再标准化(缩放特征值)使它们同分布。马氏距离其实就是标准化之后的欧式距离。
\[\begin{aligned} D_M(\mathbf{x}) &= \sqrt{(\mathbf{x} - \mathbf{\mu})^\top\Sigma^{-1}(\mathbf{x} - \mathbf{\mu})}\\ D_M(\mathbf{x}, \mathbf{y}) &= \sqrt{(\mathbf{x} - \mathbf{y})^\top\Sigma^{-1}(\mathbf{x} - \mathbf{y})}\\ 二维高斯函数: G(\mathbf{x}) &= \frac{1}{2\pi \sqrt{|\mathbf{\Sigma}|}} \exp\left( -\frac{1}{2} (\mathbf{x} - \mathbf{\mu})^\top \mathbf{\Sigma}^{-1} (\mathbf{x} - \mathbf{\mu}) \right),\\ \text{其中 } \mathbf{x} &= \begin{bmatrix} x_1 \\ x_2 \end{bmatrix},\quad \mathbf{\mu} = \begin{bmatrix} \mu_1 \\ \mu_2 \end{bmatrix},\quad \mathbf{\Sigma} = \begin{bmatrix} \sigma_{11} & \sigma_{12} \\ \sigma_{21} & \sigma_{22} \end{bmatrix}\\ \end{aligned}\]设高斯椭圆中心坐标为\((\text{center}_x, \text{center}_y)\)，像素中心坐标为\((\text{pixel}_x, \text{pixel}_y)\)，EWA 中权重的计算公式如下，其中\(\mathbf{\Sigma}\)为高斯椭圆的协方差矩阵。
\[\begin{aligned} \begin{bmatrix} d_x\\ d_y\\ \end{bmatrix} &= \begin{bmatrix} \text{center}_x\\ \text{center}_y\\ \end{bmatrix} - \begin{bmatrix} \text{pixel}_x\\ \text{pixel}_y\\ \end{bmatrix} \\ \text{power} &= -\frac{1}{2}(D_M(\text{center}, \text{pixel}))^2 = -\frac{1}{2}\left( \begin{bmatrix} d_x & d_y \end{bmatrix} \mathbf{\Sigma}^{-1} \begin{bmatrix} d_x\\ d_y\\ \end{bmatrix} \right)\\ \end{aligned}\]像素颜色累积公式如下，其中\(C\)为像素颜色，\(c\)为高斯椭圆颜色；\(\alpha\)为高斯椭圆对像素的透明度贡献，即参与合成的实际\(\alpha\)；\(T\)为像素在接受当前高斯椭圆渲染之前的透射率，即像素的剩余透明度。事实上，颜色值的更新，就是在当前处理的高斯椭圆颜色值，和截止当前的颜色累积值之间进行插值的过程。
\[\begin{aligned} 颜色更新:\quad &C_\text{final}= c_i\cdot\alpha_i+c_\text{accum}\cdot(1 - \alpha_i) \quad &即C=\sum_{i = 1}^{N} T_i\alpha_i c_i\\ 剩余透明度更新:\quad &T_\text{final}=T_\text{accum}\cdot(1 - \alpha_i) \quad &即T_i= \prod_{j = 1}^{i - 1} (1 - \alpha_j)\\ \end{aligned}\]颜色积累的迭代更新实现见如下伪代码，值得注意的是：第一，透明度贡献\(\alpha\)大于阈值的高斯椭圆才会参与对当前像素的渲染，这就是所谓的像素离散，将连续的高斯椭圆离散渲染到屏幕像素上。第二，剩余透明度\(T\)小于阈值时，表示对于当前像素而言，前面的高斯椭圆已经将光吸收完了，后面的高斯椭圆对该像素的颜色不再有贡献。
完成所有轮次后，tile 中像素的剩余透明度\(T\)参差不齐，再用统一的背景色 background color 填充每个像素的剩余透明度，就得到了所有像素最终的颜色，也就得到了完整的渲染图像。
\[C_{\text{result}} = C_{\text{final}} + T \times \text{background color}\]