WebGPU實現Ray Packet

大家好~本文在如何用WebGPU流暢渲染百萬級2D物體？基礎上進行優化，使用WebGPU實現了Ray Packet，也就是將8*8=64條射線作為一個Packet一起去訪問BVH的節點。這樣做的好處是整個Packet只需要一個維護BVH節點的Stack，節省了GPU Shared Memory；壞處是一個Packet的64條射線是并行計算的，需要實現同步邏輯，并且針對GPU的架構進行并行優化

相關文章如下：
如何用WebGPU流暢渲染百萬級2D物體？

成果
算法實現
并行計算的優化點
為什么沒有提升FPS？
為什么不使用改進的Ranged traversal算法？
具體案例分析
參考資料

成果

我們渲染1百萬個圓環，對比優化前的Demo：
性能指標：

FPS基本上沒有變化
GPU Shared Memory的占用減少為1/5
在FPS不變的情況下，Stack的大小可以從之前的最大值20提升到最大值100

硬件：

Mac OS Big Sur 11.4操作系統
Canary瀏覽器
Intel Iris Pro 1536 MB集成顯卡

算法實現

每條射線使用一個局部單位，所以一個Packet共使用8*8=64個局部單位

這里將射線與相機近平面的交點變換到了屏幕坐標系下，變量名為pointInScreen（之前的Demo中的交點是在世界坐標系下）

整個算法的相關實現代碼如下所示：

var<workgroup>stackContainer: array<TopLevel, 20>;
var<workgroup>rayPacketAABBData: array<f32, 4>;
var<workgroup>isRayPacketAABBIntersectWithTopLevelNode: bool;
var<workgroup>rayPacketRingIntersectLayer: array<u32, 64>;
var<workgroup>stackSize: u32;
var<workgroup>isAddChild1: bool;
var<workgroup>isAddChild2: bool;

...

fn _intersectScene(ray: Ray, LocalInvocationIndex : u32) -> RingIntersect {
  var intersectResult: RingIntersect;

  intersectResult.isClosestHit = false;
  intersectResult.layer = 0;

  var rootNode = topLevel.topLevels[0];

  var pointInScreen = ray.rayTarget;

  // 用兩個局部單位并行創建Packet的AABB
  if (LocalInvocationIndex == 0) {
    rayPacketAABBData[0] = pointInScreen.x;
    rayPacketAABBData[1] = pointInScreen.y;
  }
  if (LocalInvocationIndex == 63) {
    rayPacketAABBData[2] = pointInScreen.x;
    rayPacketAABBData[3] = pointInScreen.y;
  }

  //用一個局部單位并行初始化共享變量
  if (LocalInvocationIndex == 1) {
    stackSize = 1;

    stackContainer[0] = rootNode;
  }

  workgroupBarrier();

  //遍歷BVH節點
  while (stackSize > 0) {
    //在循環開始的時候要同步下（暫時不清楚原因，反正是因為stackSize是共享變量造成的）
    workgroupBarrier();

    if (LocalInvocationIndex == 0) {
      stackSize -= 1;
    }

    workgroupBarrier();


    var currentNode = stackContainer[stackSize];
    var leafInstanceCountAndMaxLayer = u32(currentNode.leafInstanceCountAndMaxLayer);

    var leafInstanceCount = _getLeafInstanceCount(leafInstanceCountAndMaxLayer);

    //如果是葉節點，則Packet中所有射線都與該節點進行相交檢測
    if (_isLeafNode(leafInstanceCount)) {
      //判斷射線是否與節點的AABB相交
      if (_isPointIntersectWithTopLevelNode(pointInScreen, currentNode)) {
        var leafInstanceOffset = u32(currentNode.leafInstanceOffset);

        var maxLayer = _getMaxLayer(u32(currentNode.leafInstanceCountAndMaxLayer));

        //判斷射線是否與節點包含的所有圓環相交
        while (leafInstanceCount > 0) {
          var bottomLevel = bottomLevel.bottomLevels[leafInstanceOffset];
          //判斷射線是否與圓環的AABB相交
          if (_isPointIntersectWithAABB(pointInScreen, bottomLevel.screenMin, bottomLevel.screenMax)) {
            var instance: Instance = sceneInstanceData.instances[u32(bottomLevel.instanceIndex)];
            var geometry: Geometry = sceneGeometryData.geometrys[u32(instance.geometryIndex)];

            //判斷射線是否與圓環相交
            if (_isIntersectWithRing(pointInScreen, instance, geometry)) {
              var layer = u32(bottomLevel.layer);

              if (!intersectResult.isClosestHit || layer > intersectResult.layer) {
                intersectResult.isClosestHit = true;
                intersectResult.layer = layer;
                intersectResult.instanceIndex = bottomLevel.instanceIndex;
              }
            }
          }

          leafInstanceCount = leafInstanceCount - 1;
          leafInstanceOffset = leafInstanceOffset + 1;
        }
      }
    }
    //如果不是葉節點，則通過剔除檢測后加入兩個子節點
    else {
      // 一個非葉節點必有兩個子節點
      var child1Node = topLevel.topLevels[u32(currentNode.child1Index)];
      var child2Node = topLevel.topLevels[u32(currentNode.child2Index)];

      var child1NodeMaxLayer = _getMaxLayer(u32(child1Node.leafInstanceCountAndMaxLayer));
      var child2NodeMaxLayer = _getMaxLayer(u32(child2Node.leafInstanceCountAndMaxLayer));

      //并行計算Packet中所有射線的最小layer

      rayPacketRingIntersectLayer[LocalInvocationIndex] = intersectResult.layer;

      workgroupBarrier();

      if (LocalInvocationIndex < 32) {
        _minForRayPacketRingIntersectLayer(LocalInvocationIndex, LocalInvocationIndex + 32);
      }
      workgroupBarrier();
      if (LocalInvocationIndex < 16) {
        _minForRayPacketRingIntersectLayer(LocalInvocationIndex, LocalInvocationIndex + 16);
      }
      workgroupBarrier();
      if (LocalInvocationIndex < 8) {
        _minForRayPacketRingIntersectLayer(LocalInvocationIndex, LocalInvocationIndex + 8);
      }
      workgroupBarrier();
      if (LocalInvocationIndex < 4) {
        _minForRayPacketRingIntersectLayer(LocalInvocationIndex, LocalInvocationIndex + 4);
      }
      workgroupBarrier();
      if (LocalInvocationIndex < 2) {
        _minForRayPacketRingIntersectLayer(LocalInvocationIndex, LocalInvocationIndex + 2);
      }
      workgroupBarrier();

      //用兩個局部單位并行判斷子節點是否在Packet的前面以及Packet的AABB是否與子節點相交

      if (LocalInvocationIndex == 0) {
        isAddChild1 = child1NodeMaxLayer > rayPacketRingIntersectLayer[0] && _isRayPacketAABBIntersectWithTopLevelNode(rayPacketAABBData, child1Node);
      }
      if (LocalInvocationIndex == 1) {
        isAddChild2 = child2NodeMaxLayer > rayPacketRingIntersectLayer[0] && _isRayPacketAABBIntersectWithTopLevelNode(rayPacketAABBData, child2Node);
      }

      workgroupBarrier();

      // 加入兩個子節點到stack中
      if (LocalInvocationIndex == 0) {
        if (isAddChild1) {
          stackContainer[stackSize] = child1Node;

          stackSize += 1;
        }

        if (isAddChild2) {
          stackContainer[stackSize] = child2Node;

          stackSize += 1;
        }
      }
    }

    workgroupBarrier();
  }

  return intersectResult;
}

“創建Packet的AABB”代碼說明

如上圖所示，紅色方塊為Packet的2D AABB，由64個射線的pointInScreen組成
圖中藍色方塊為一個pointInScreen，從中心點開始

pointInScreen對應的局部單位序號是從紅色方塊左下角開始，朝著右上方增加。
部分局部單位序號在圖中用數字標注了出來

所以Packet的2D AABB的min即為0號局部單位對應的pointInScreen，而max則為63號局部單位對應的pointInScreen

創建Packet的AABB的相關代碼如下：

  var pointInScreen = ray.rayTarget;

  // 用兩個局部單位并行創建Packet的AABB
  if (LocalInvocationIndex == 0) {
    rayPacketAABBData[0] = pointInScreen.x;
    rayPacketAABBData[1] = pointInScreen.y;
  }
  if (LocalInvocationIndex == 63) {
    rayPacketAABBData[2] = pointInScreen.x;
    rayPacketAABBData[3] = pointInScreen.y;
  }

并行計算的優化點

因為現在需要對一個Packet中64條射線并行計算，所以需要了解GPU的架構和特點，從而進行相應的優化

GPU是以warp為單位，每個wrap包含32個線程。所以我們這里的一個Packet應該使用了兩個wrap，其中一個wrap中的一個線程對應一個局部單位

并行計算的優化點為：

只有當整個wrap中所有線程都不執行某個操作時，這個wrap才不會被執行，從而FPS會提高。只要wrap中至少有一個線程要執行某個操作，那么即使其它所有線程不執行該操作，它們也會在執行"workgroupBarrier()"時等待）
GPU中Memory IO是個瓶頸，所以應該減少內存讀寫操作
GPU內存類型分為寄存器、共享、局部、常量、紋理、全局，讀寫速度依次遞減
內存模型如下圖所示：

內存特性如下圖所示：

應該減少速度慢的局部、常量、紋理、全局內存的讀寫

減少bank conflict
共享內存是由多個bank組成，對于同一個bank的不同地址的讀寫會造成bank conflict
因此盡量讀寫共享內存的連續地址
減少wrap diverse
應該減少一個Packet中不同的局部單位進入不同的if分支的情況(這會造成局部單位阻塞)

為什么沒有提升FPS？

之前的Demo使用串行算法，而現在的Demo使用并行算法

如果都不使用traverse order優化（即判斷節點如果在射線的后面，則不遍歷它），則FPS提升了50%；
如果都使用traverse order優化，則FPS沒有變化；

這說明traverse order優化對于串行算法的提升更大。這是因為對于并行算法而言，只有當節點在一個Packet的所有射線的后面時，才不會遍歷節點（可以參考之前的提到的優化點1）；
而對于串行算法而言，只要當節點在一條射線的后面，就不會遍歷節點

為什么不使用改進的Ranged traversal算法？

通過Large Ray Packets for Real-time Whitted Ray Tracing論文，得知現在的并行算法屬于文中提到的“Masked traversal”算法
文中還介紹了改進版的Ranged traversal算法，具體就是指一個Packet增加first active和last active標志，從而使一個Packet中只有first到last之間的射線進行相交檢測，減少了相交檢測的射線數量

但是這個算法應用到本文的并行算法中并不會提高FPS！因為本文的并行算法中的一個Packet至少會有一條射線會與葉節點進行相交檢測，所以根據之前提到的優化點1可知FPS不會提高

具體案例分析

現在我們對代碼進行一些具體的分析：

寄存器、共享、局部、全局內存分析

參考代碼與存儲的映射關系：

以及參考gpu cpu 共享內存提高傳輸速度_GPU編程3--GPU內存深入了解，我們來分析下本文部分代碼中的內存使用情況：

fn _intersectScene(ray: Ray, LocalInvocationIndex : u32) -> RingIntersect {
    //intersectResult是較大結構體，應該位于局部內存
    var intersectResult: RingIntersect;

    intersectResult.isClosestHit = false;
    intersectResult.layer = 0;

    //topLevel.topLevels是storage buffer數據，位于全局內存

    //rootNode是較大結構體，應該位于局部內存
    var rootNode = topLevel.topLevels[0];

    //pointInScreen位于寄存器
    var pointInScreen = ray.rayTarget;

    if (LocalInvocationIndex == 0) {
        //rayPacketAABBData位于共享內存
        rayPacketAABBData[0] = pointInScreen.x;
        rayPacketAABBData[1] = pointInScreen.y;
    }
    if (LocalInvocationIndex == 63) {
        rayPacketAABBData[2] = pointInScreen.x;
        rayPacketAABBData[3] = pointInScreen.y;
    }

    ...
}

GPU Memory IO優化

根據之前提到的優化點2，我們知道應該減少GPU內存的讀寫，特別是全局內存的讀寫
我們來對照代碼分析一下全局內存讀寫的優化，看下面的代碼：

var rootNode = topLevel.topLevels[0];

...


if (LocalInvocationIndex == 1) {
    stackSize = 1;

    stackContainer[0] = rootNode;
}

這里一個Packet中所有局部單位都從全局內存中讀取第一個元素為rootNode，寫到本地內存中（這里進行了64次全局內存讀操作）；
然后在1號局部單位中，從本地內存中讀取rootNode，寫到共享內存中

如果將代碼改為：

if (LocalInvocationIndex == 1) {
    stackSize = 1;

    stackContainer[0] = topLevel.topLevels[0];
}

那么就應該只進行1次全局內存讀操作，從而提高FPS

但是實際上卻降低了10%左右FPS，這是為什么呢？

這是因為GPU會將對全局內存同一地址或者相鄰地址的讀操作合并為一次操作（寫操作也是一樣），所以修改前后的代碼對全局內存的讀操作都是1次。

那么FPS應該不變，但為什么下降了呢？

這是因為在進行內存操作時，需要加上事務（進行鎖之類的同步操作）。如果一個wrap中的所有線程同時對全局內存的一個地址進行讀，則合并后的該次操作只需要一個事務；而如果是一個wrap中的部分線程進行讀，則合并后的該次操作需要更多的事務（如4個），從而需要更多時間開銷，降低了FPS

參考資料：
Nvidia GPU simultaneous access to a single location in global memory
初識事務內存（Transactional Memory）

我們再看下面的代碼：

    else {
      var child1Node = topLevel.topLevels[u32(currentNode.child1Index)];
      var child2Node = topLevel.topLevels[u32(currentNode.child2Index)];

      var child1NodeMaxLayer = _getMaxLayer(u32(child1Node.leafInstanceCountAndMaxLayer));
      var child2NodeMaxLayer = _getMaxLayer(u32(child2Node.leafInstanceCountAndMaxLayer));

這里只需要一次合并后的讀操作，從全局內存中讀出child1Node、child2Node

如果代碼改為：

    else {
      var child1Node = topLevel.topLevels[u32(currentNode.child1Index)];
      var child1NodeMaxLayer = _getMaxLayer(u32(child1Node.leafInstanceCountAndMaxLayer));

      var child2Node = topLevel.topLevels[u32(currentNode.child2Index)];
      var child2NodeMaxLayer = _getMaxLayer(u32(child2Node.leafInstanceCountAndMaxLayer));

因為要切換內存（全局和局部內存之間切換），所以就不能合并了，而需要二次全局內存的讀操作，分別讀出child1Node和child2Node，所以FPS會下降10%左右

parallel reduction優化并沒有提高FPS

下面的代碼使用了parallel reduction的優化版本來并行計算Packet中所有射線的最小layer：

      if (LocalInvocationIndex < 32) {
        _minForRayPacketRingIntersectLayer(LocalInvocationIndex, LocalInvocationIndex + 32);
      }
      workgroupBarrier();
      if (LocalInvocationIndex < 16) {
        _minForRayPacketRingIntersectLayer(LocalInvocationIndex, LocalInvocationIndex + 16);
      }
      workgroupBarrier();
      if (LocalInvocationIndex < 8) {
        _minForRayPacketRingIntersectLayer(LocalInvocationIndex, LocalInvocationIndex + 8);
      }
      workgroupBarrier();
      if (LocalInvocationIndex < 4) {
        _minForRayPacketRingIntersectLayer(LocalInvocationIndex, LocalInvocationIndex + 4);
      }
      workgroupBarrier();
      if (LocalInvocationIndex < 2) {
        _minForRayPacketRingIntersectLayer(LocalInvocationIndex, LocalInvocationIndex + 2);
      }
      workgroupBarrier();

關于parallel reduction優化可以參考啥是Parallel Reduction

優化后的代碼并不比下面的原始版本的代碼快:

for (var s: u32 = 1; s < 64; s = s * 2) {
  if (LocalInvocationIndex % (2 * s) == 0) {
    _minForRayPacketRingIntersectLayer(LocalInvocationIndex, LocalInvocationIndex + s);
  }
  workgroupBarrier();
}

估計是GPU幫我們優化了

wrap diverse優化

在之前使用串行算法的Demo代碼中，我們在遍歷葉節點包含的所有圓環時，加入了“如果找到了最前面的相交圓環，則退出葉節點的遍歷”的優化，相關代碼為：

      if (_isLeafNode(leafInstanceCount)) {
        ...
        while (leafInstanceCount > 0) {
          ...
          if (_isPointIntersectWithAABB(point, bottomLevel.worldMin, bottomLevel.worldMax)) {
            ...
            if (_isIntersectWithRing(point, instance, geometry)) {
              if (!intersectResult.isClosestHit || layer > intersectResult.layer) {
                ...

                if(layer == maxLayer){
                  break;
                }
              }
            }
          }

          ...
        }

而在本文的使用并行算法的代碼，卻刪除了這個優化，FPS反而提升了10%。這是因為Packet中只有部分的局部單位會退出遍歷，造成其它局部單位阻塞等待

注：如果并行算法的代碼要加入這個優化，那也不能直接break，而是改為設置isBreak=true，在if中判斷isBreak，這樣才不會有同步的問題

bank conflict優化

現代顯卡實現了broadcast機制，從而對同一個bank的相同地址的讀寫不會造成bank conflict。但由于我的測試顯卡是老顯卡，還沒有該機制，所以在這種情況下會造成bank conflict
如下面代碼所示：

      if (LocalInvocationIndex == 0) {
        isAddChild1 = _isRayPacketAABBIntersectWithTopLevelNode(rayPacketAABBData, child1Node);
      }
      if (LocalInvocationIndex == 1) {
        isAddChild2 = _isRayPacketAABBIntersectWithTopLevelNode(rayPacketAABBData, child2Node);
      }

兩個局部單位同時讀共享變量rayPacketAABBData，從而造成bank conflict，FPS下降5%左右