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}$$

参考资料

3DGS CUDA 代码解读

3DGS CUDA Kernel