大家好~本文使用WebGPU和光线追踪算法，从0开始实现和逐步优化Demo，展示了从渲染500个2D物体都吃力到流畅渲染4百万个2D物体的优化过程和思路

• Optimizing GPU occupancy and resource usage with large thread groups

  我们知道了一个计算单位中只有64KB的内存用于存储（我们只用到了64KB的VGPRs (Vector General-Purpose Registers))，并且由于一个计算单位会并行执行两个线程组，所以每个线程组只有32KB的内存大小。

  然而在_intersectScene的for循环中遍历了所有的圆环。这意味着会将所有圆环的数据都加载到计算单位的内存中，从而超过了32KB的大小！
  这导致了计算单位只能使用1个线程组，并且该线程组只能使用1个而不是所有的线程！
  所以不仅相当于没启用局部单位，反而可能因为试图启用局部单位而造成的同步开销，导致FPS不升反降！

  结论

  综上分析，我们需要引入BVH，大幅减少for循环中需要遍历的数据，使其小于32KB，从而能够将其载入到计算单位的内存中；
  然后我们再启用8*8的局部单位！

  6、实现BVH

  BVH是一个用于空间划分的树，相关介绍可参考：
  场景管理方法之BVH介绍
  GAMES101-现代计算机图形学入门-闫令琪

  我们需要在CPU端构造BVH树，将其传入到CS；然后在CS->_intersectScene函数中遍历BVH

  6.1、实现构造BVH

  构造BVH树

  我们在CPU端使用最简单暴力的Middle方法构造BVH树，步骤如下：

  1、计算所有圆环的包围盒AABB
  2、构造根节点，以所有AABB形成的整体AABB为该节点的包围盒
  3、从x轴方向，按照AABB的中心点位置排序所有的AABB
  4、以中间的AABB为分界线分割，形成两个子节点，每个子节点以其包含的所有AABB形成的整体AABB为该子节点的包围盒
  5、递归地构造这两个子节点，并且交替地从y轴方向开始（即x轴->y轴->x轴。。。。。。交替）排序该节点包含的所有AABB和分割。。。。。。直到节点包含的AABB个数<=5或达到最大深度时结束递归

  构造加速结构
  因为CS中只能用数组，所以需要将BVH树拍平成数组，作为加速结构传送到CS

  我们将加速结构设计为两层：TopLevel和BottomLevel

  BottomLevel的数据类型如下：

  type worldMinX = number
  type worldMinY = number
  type worldMaxX = number
  type worldMaxY = number
  
  type instanceIndex = number
  
  type bottomLevel = Array<[worldMinX, worldMinY, worldMaxX, worldMaxY, instanceIndex]>
  

  BottomLevel用来保存所有圆环对应的包围盒、instanceIndex

  TopLevel的数据类型如下：

  type wholeWorldMinX = number
  type wholeWorldMinY = number
  type wholeWorldMaxX = number
  type wholeWorldMaxY = number
  
  type leafInstanceOffset = number
  type leafInstanceCount = number
  type child1Index = number
  type child2Index = number
  
  type topLevel = Array<[
  	wholeWorldMinX, wholeWorldMinY, wholeWorldMaxX, wholeWorldMaxY,
  	leafInstanceOffset,
  	leafInstanceCount,
  	child1Index,
  	child2Index
  ]>
  
  

  TopLevel用来保存BVH节点的包围盒、节点包含的AABB在BottomLevel数组中的索引（如果该节点不是叶节点，则leafInstanceCount=0）、子节点在topLevel数组中的索引

  然后将加速结构传到CS中，CS中对应数据结构如下：

  struct TopLevel {
    worldMin : vec2,
    worldMax : vec2,
  
    leafInstanceOffset: f32,
    leafInstanceCount: f32,
    child1Index: f32,
    child2Index: f32
  }
  
  struct BottomLevel {
    worldMin : vec2,
    worldMax : vec2,
  
    instanceIndex: f32,
    pad_0: f32,
    pad_1: f32,
    pad_2: f32,
  }
  
   struct TopLevels {
    topLevels : array,
  }
  
   struct BottomLevels {
    bottomLevels : array,
  }
  
  @binding(0) @group(0) var topLevel : TopLevels;
  @binding(1) @group(0) var bottomLevel : BottomLevels;
  

  6.2、CPU端实现遍历BVH

  为了方便测试，我们先在CPU端实现遍历BVH；解决所有bug后，再移植到CS中

  可以分别用递归和迭代来实现

  考虑到WGSL不支持递归函数，所以我们只用迭代来实现
  相关代码如下：

  type traverseResult = {
  	isHit: boolean,
  	instanceIndex: instanceIndex | null
  }
  
  let _isPointIntersectWithAABB = (
  	point,
  	wholeWorldMinX, wholeWorldMinY, wholeWorldMaxX, wholeWorldMaxY,
  ) => {
  	return point[0] > wholeWorldMinX && point[0] < wholeWorldMaxX && point[1] > wholeWorldMinY && point[1] < wholeWorldMaxY
  }
  
  let _isPointIntersectWithTopLevelNode = (point, node: topLevelNodeData) => {
  	let [
  		wholeWorldMinX, wholeWorldMinY, wholeWorldMaxX, wholeWorldMaxY,
  		leafInstanceOffset,
  		leafInstanceCount,
  		child1Index,
  		child2Index
  	] = node
  
  	return _isPointIntersectWithAABB(
  		point,
  		wholeWorldMinX, wholeWorldMinY, wholeWorldMaxX, wholeWorldMaxY,
  	)
  }
  
  let _isLeafNode = (node: topLevelNodeData) => {
  	let leafInstanceCountOffset = 5
  
  	return node[leafInstanceCountOffset] !== 0
  }
  
  let _handleIntersectWithLeafNode = (intersectResult, isIntersectWithInstance, point, node: topLevelNodeData, bottomLevelArr: bottomLevelArr) => {
  	let [
  		wholeWorldMinX, wholeWorldMinY, wholeWorldMaxX, wholeWorldMaxY,
  		leafInstanceOffset,
  		leafInstanceCount,
  		child1Index,
  		child2Index
  	] = node
  
  	while (leafInstanceCount > 0) {
  		let [worldMinX, worldMinY, worldMaxX, worldMaxY, instanceIndex] = bottomLevelArr[leafInstanceOffset]
  
  		if (_isPointIntersectWithAABB(
  			point,
  			worldMinX, worldMinY, worldMaxX, worldMaxY
  		)) {
  			if (isIntersectWithInstance(point, instanceIndex)) {
  				intersectResult.isHit = true
  				intersectResult.instanceIndex = instanceIndex
  
  				break;
  			}
  		}
  
  		leafInstanceCount -= 1
  		leafInstanceOffset += 1
  	}
  }
  
  let _hasChild = (node, childIndexOffset) => {
  	return node[childIndexOffset] !== 0
  }
  
  export let traverse = (isIntersectWithInstance, point, topLevelArr: topLevelArr, bottomLevelArr: bottomLevelArr): traverseResult => {
  	let rootNode = topLevelArr[0]
  
  	let child1IndexOffset = 6
  	let child2IndexOffset = 7
  
  	let stackContainer = [rootNode]
  	let stackSize = 1
  
  	let intersectResult = {
  		isHit: false,
  		instanceIndex: null
  	}
  
  	while (stackSize > 0) {
  		let currentNode = stackContainer[stackSize - 1]
  
  		stackSize -= 1
  
  		if (_isPointIntersectWithTopLevelNode(point, currentNode)) {
  			if (_isLeafNode(currentNode)) {
  				_handleIntersectWithLeafNode(intersectResult, isIntersectWithInstance, point, currentNode, bottomLevelArr)
  
  				if (intersectResult.isHit) {
  					break
  				}
  			}
  			else {
  				if (_hasChild(currentNode, child1IndexOffset)) {
  					stackContainer[stackSize] = topLevelArr[currentNode[child1IndexOffset]]
  					stackSize += 1
  				}
  				if (_hasChild(currentNode, child2IndexOffset)) {
  					stackContainer[stackSize] = topLevelArr[currentNode[child2IndexOffset]]
  					stackSize += 1
  				}
  			}
  		}
  	}
  
  	return intersectResult
  }
  

  这里用了栈(stackContainer)来保存需要遍历的节点

  本来我们可以直接通过stackContainer.push方法将节点push到栈中，但考虑到WGSL的数组没有push操作，所以这里我们就增加了stackSize这个数据，从而能够通过"stackContainer[stackSize] = 节点"来代替"stackContainer.push(节点)"

  6.3、GPU端实现遍历BVH

  CPU端测试通过后，我们将其移植到CS中

  这里值得注意的是因为WGSL创建数组时必须定义大小，所以栈的大小必须预先确定且为常数

  栈的大小其实就是BVH树的最大深度，我们可以先暂时指定为20

  相关代码如下：

  fn _intersectScene(ray: Ray)->RingIntersect {
  const MAX_DEPTH = 20;
  
  var stackContainer:array;
  
  ...
  }
  

  7、测试渲染极限

  现在我们将圆环数量增加200倍，渲染500*200=10万个圆环，测试下FPS

  我们将圆环的半径和圆环宽度缩小为1/10，这样方便显示
  渲染结果如下图所示：

  FPS没有变化

  结论

  通过引入BVH，渲染性能提高了200倍

  8、设置workgroup_size

  因为引入了BVH，需要遍历的节点数大大减少了，所以减少了显存占用

  现在再次启用局部单位：
  我们将work group减少为原来的1/64，将work group的size设为(8,8,1)：

  运行Demo后，发现FPS变为60了

  理论上可以提高64倍的渲染速度
  所以我们将圆环数量增加40倍，渲染10万*40=4百万个圆环，测试下FPS
  结果FPS跟之前10万个时一样

  结论

  通过启用局部单位，渲染性能提高了40倍

  9、测试内存占用

  通过Chrome dev tool->Memory->Take heap snapshot，可以看到包含4百万个圆环的场景在CPU端只占用了211MB左右的内存，说明内存占用确实小

  10、使用LBVH算法来构造BVH

  现在当圆环数量为4百万个时，CPU端构造BVH需要花费100秒以上的时间！

  因此，我们保持“构造加速结构”的代码不变，修改“构造BVH树”的算法为LBVH算法

  它的步骤如下：

  1、计算所有圆环的包围盒AABB
  2、构造根节点，以所有AABB形成的AABB为该节点的包围盒
  3、根据根节点的包围盒，在x、y轴方向上将其1024等分，根据AABB的中心点在哪个区域而计算出AABB在x、y轴方向上的格子坐标
  4、将格子坐标转换为Morton Code
  5、根据Morton Code将所有的AABB排序
  6、对于该有序Morton Code数组，我们利用二分查找出第一个不同的bit位（也即是从0变为1的index），此时我们即可将最高位为0的所有BVs归入此node（此时是root节点）的左子树，最高位为1的所有BVs归入右子树；同理，我们对左右子树按照下一个bit位来递归的处理，直到递归的处理完全部bit位，LBVH即可建立完毕

  最后一步如下图所示：
  

  结论

  通过改为LBVH算法，构造时间降低为10秒左右，性能提高了10倍

  之所以LBVH算法更快，是因为只排序了一次！

  11、实现剔除

  之前分析过剔除的实现思路：

  首先圆环加上“层”的数据；
  然后遍历所有圆环，判断像素在哪些圆环上；
  最后取出最大“层”的圆环，将它的颜色作为像素的颜色

  因此，可以在圆环的transform组件数据中增加layer数据（从1开始的正整数），用来表示“层”；
  然后在构造加速结构BottomLevel时，读取transform组件中的layer数据，将其保存到BottomLevel数据中；
  最后修改CS代码，在遍历BVH->检测到像素在圆环上时比较layer，并且不再停止遍历，而是继续遍历栈的其它节点；在遍历栈的其它节点时，如果找到了像素在其上的圆环，也不再停止

  实现剔除后，运行Demo渲染1百万个圆环时FPS都才15（之前是4百万个 45 FPS），渲染性能估计降低了10倍

  12、改进遍历BVH

  我们从下面几个方面对剔除进行优化

  traverse order优化

  现在如果找到了像素在其上的圆环，会继续遍历其它节点。
  我们需要减少遍历的节点数量

  我们可以在构造BVH树时，为每个节点增加maxLayer数据，它为该节点包含的所有的圆环中最大的层；
  然后在遍历BVH时：
  检测到像素在圆环上时，将圆环的layer记录到相交结果intersectResult中；
  在遍历栈中的节点的while循环中，如果该节点的maxLayer <= intersectResult.layer，则说明该节点包含的所有圆环都被遮挡了，直接continue，跳过；
  另外，在遍历叶节点的所有圆环时，如果像素所在的圆环的layer==该节点的maxLayer，则说明已经找到了叶节点包含的所有圆环中的最大层的圆环，则break，停止搜索该叶节点包含的其它圆环。

  通过上面的优化，可以大幅降低遍历的节点数量

  合并数据

  因为BVH树的节点需要保存maxLayer数据，而这个数据实际上是保存在TopLevel中的，对应到CS中的数据结构就是：

  struct TopLevel {
  	worldMin : vec2,
  	worldMax : vec2,
  
  	leafInstanceOffset: f32,
  	leafInstanceCount: f32,
  	child1Index: f32,
  	child2Index: f32
  
  	maxLayer: f32,
  	pad_0: f32,
  	pad_1: f32,
  	pad_2: f32,
  }
  
  

  我们可以看到，因为增加了maxLayer，需要增加3个pad数据来对齐，这样浪费了显存

  而占用尽可能少的显存是非常重要的，因为经过之前的分析，我们知道一个计算单位只有32KB的显存可用，超过的话就会导致启用局部单位失效！

  仔细观察后，我们发现leafInstanceCount和maxLayer只需要占用<32位的字节数
  我们可以让这两个数据各占16位，其中leafInstanceCount在高位，maxLayer在低位；然后将其合成一个32位f32
  从而TopLevel修改为：

  struct TopLevel {
  	worldMin : vec2,
  	worldMax : vec2,
  
  	leafInstanceOffset: f32,
  	leafInstanceCountAndMaxLayer: f32,
  	child1Index: f32,
  	child2Index: f32
  }
  

  这样就消除了pad数据，减少了显存占用

  但是当渲染的圆环数量超过1百万个时，会出现“从leafInstanceCountAndMaxLayer中取出的maxLayer为0（应该>=1）”的bug！
  这是因为leafInstanceCount（叶节点的圆环个数）过大，占用了超出了16位的字节数，从而影响到maxLayer的值！

  所以我们重新分配，让leafInstanceCount占24位，maxLayer占8位，则解决了bug

  注：理论上16位的leafInstanceCount可以最大为65535，但实际上我发现当leafInstanceCount>1000时，就出现了超过16位的情况！我估计是WGLSL可能占用了f32类型的数据的最高几位，导致leafInstanceCount实际可用位数<16位

  13、测试渲染极限

  现在我们将圆环数量恢复为4百万个圆环，FPS又恢复为45左右

  当我们尝试渲染5百万个圆环时，遇到了“我们BottomLevel Buffer数据超出了Storage Buffer的最大限制：128MB”的问题

  关于这个限制，WebGPU官方有相关的讨论issue：
  Limit for the maximum buffer size

  绕过该限制的可能方案是将其拆成多个Storage Buffer

  结论

  通过traverse order优化，渲染性能提高了10倍左右

  当我们尝试渲染5百万个圆环时，遇到了超出Storage Buffer最大大小的限制

  总结

  感谢大家的学习~

  在本文中，我们先提出了需求；然后按照需求来设计和选择算法；然后实现最简单的版本；接着不断优化，直到达到Storage Buffer的最大大小限制为止

  我们的优化的成果为：

  • 通过引入BVH，渲染性能提高了200倍
  • 通过启用局部单位，渲染性能提高了40倍
  • 通过改为LBVH算法，性能提高了10倍
  • 通过traverse order优化，使得剔除的性能提高了10倍左右

  目前我们最多渲染4百万个圆环，因为再多就会超出Storage Buffer最大大小的限制

  后续的改进方向

  后面我们希望能够渲染千万级2D物体，可以从下面的方向改进：

  • 将加速结构拆成多个Storage Buffer
  • 优化构造BVH，使叶节点包含的圆环数量尽量少，且节点的包围盒尽量不重叠，这样才能提高遍历BVH的性能
    可考虑使用HLBVH算法
  • 优化遍历BVH：考虑并行遍历BVH、无栈的遍历
    参考Ray Tracing学习之Traversal

  参考资料

  OpenGL4.3 新特性: 计算着色器 Compute Shader
  Bad preformance of simple ray trace compute shader
  Optimizing GPU occupancy and resource usage with large thread groups
  我所理解的DirectX Ray Tracing
  并行构建BVH
  Build LBVH on GPUs
  Ray Tracing学习之Traversal
  光线求交加速算法：边界体积层次结构(Bounding Volume Hierarchies)3-LBVH(Linear Bounding Volume Hierarchies)
  WebGPU 计算管线、计算着色器（通用计算）入门案例：2D 物理模拟