发布于 由admin

JS `WebGPU` `Ray Tracing` (`BVH` 结构) 实现与性能调优

各位观众老爷们，欢迎来到今天的 WebGPU Ray Tracing 性能调优脱口秀（误）。今天咱们聊聊怎么用 WebGPU 搞定光线追踪，还必须得是带 BVH 加速的那种，最后再顺手调调性能。保证各位听完之后，以后面试再也不怕被问到 WebGPU 光追了！

一、WebGPU 光追：Hello, World! 之后该干啥？

首先，咱们得明确一点：WebGPU 本身并没有内置的光线追踪 API。所以，咱们得自己手动实现。这听起来吓人，但其实也就那么回事。核心思路就是：

  1. 光线生成 (Ray Generation)： 从摄像机发出光线，穿过屏幕上的每个像素。
  2. 场景相交测试 (Scene Intersection)： 检查光线是否与场景中的物体相交。
  3. 着色 (Shading)： 如果光线相交，根据材质属性和光照计算像素颜色。

最简单的光线追踪，那就是遍历场景中的每个三角形，看看光线是不是跟它相交。这效率嘛，懂得都懂，简直慢到姥姥家了。所以，我们需要 BVH (Bounding Volume Hierarchy) 这种神器来加速相交测试。

二、BVH：光追加速的秘密武器

BVH，简单来说，就是把场景中的物体（通常是三角形）组织成一棵树。每个节点都包含一个包围盒 (Bounding Box)，这个包围盒包裹了它所有子节点包含的物体。这样，我们在进行光线相交测试时，就可以先测试光线是否与节点的包围盒相交。如果没相交，那这个节点下的所有物体就都可以直接跳过了，大大减少了需要测试的物体数量。

BVH 的构建方式有很多种，比如：

  • Top-Down (自顶向下): 从整个场景开始，递归地将节点分裂成更小的子节点，直到每个叶子节点包含少量物体。
  • Bottom-Up (自底向上): 从单个物体开始，逐步将它们合并成更大的包围盒，直到形成根节点。

这里我们选择最常见的 Top-Down 方式，使用 Surface Area Heuristic (SAH) 来决定如何分裂节点。SAH 的目标是最小化光线穿过整个 BVH 的平均开销。

三、WebGPU + WGSL：代码时间到！

接下来，就是激动人心的代码环节了！咱们先定义一些基本的数据结构：

// 包围盒
class BoundingBox {
  constructor(min, max) {
    this.min = min;
    this.max = max;
  }

  intersects(rayOrigin, rayDirection) {
    let tMin = (this.min[0] - rayOrigin[0]) / rayDirection[0];
    let tMax = (this.max[0] - rayOrigin[0]) / rayDirection[0];

    if (tMin > tMax) {
      [tMin, tMax] = [tMax, tMin]; // Swap
    }

    let tyMin = (this.min[1] - rayOrigin[1]) / rayDirection[1];
    let tyMax = (this.max[1] - rayOrigin[1]) / rayDirection[1];

    if (tyMin > tyMax) {
      [tyMin, tyMax] = [tyMax, tyMin]; // Swap
    }

    if ((tMin > tyMax) || (tyMin > tMax)) {
      return false;
    }

    tMin = Math.max(tMin, tyMin);
    tMax = Math.min(tMax, tyMax);

    let tzMin = (this.min[2] - rayOrigin[2]) / rayDirection[2];
    let tzMax = (this.max[2] - rayOrigin[2]) / rayDirection[2];

    if (tzMin > tzMax) {
      [tzMin, tzMax] = [tzMax, tzMin]; // Swap
    }

    if ((tMin > tzMax) || (tzMin > tMax)) {
      return false;
    }

    tMin = Math.max(tMin, tzMin);
    tMax = Math.min(tMax, tzMax);

    return tMax >= 0; // Important: Check for intersection behind the ray origin
  }
}

// BVH 节点
class BVHNode {
    constructor(boundingBox, left, right, triangleIndex) {
        this.boundingBox = boundingBox;
        this.left = left;
        this.right = right;
        this.triangleIndex = triangleIndex; // Only for leaf nodes
    }
}

// 三角形
class Triangle {
  constructor(v0, v1, v2) {
    this.v0 = v0;
    this.v1 = v1;
    this.v2 = v2;
  }

  intersects(rayOrigin, rayDirection) {
      // Moller-Trumbore intersection algorithm
      const EPSILON = 0.000001;
      const edge1 = [this.v1[0] - this.v0[0], this.v1[1] - this.v0[1], this.v1[2] - this.v0[2]];
      const edge2 = [this.v2[0] - this.v0[0], this.v2[1] - this.v0[1], this.v2[2] - this.v0[2]];

      const h = cross(rayDirection, edge2);
      const a = dot(edge1, h);

      if (a > -EPSILON && a < EPSILON)
          return null;    // This ray is parallel to this triangle.

      const f = 1.0 / a;
      const s = [rayOrigin[0] - this.v0[0], rayOrigin[1] - this.v0[1], rayOrigin[2] - this.v0[2]];
      const u = f * dot(s, h);

      if (u < 0.0 || u > 1.0)
          return null;

      const q = cross(s, edge1);
      const v = f * dot(rayDirection, q);

      if (v < 0.0 || u + v > 1.0)
          return null;

      // At this stage we can compute t to find ray intersection
      const t = f * dot(edge2, q);

      if (t > EPSILON) // ray intersection
          return t;
      else // This means that there is a line intersection but not a ray intersection
          return null;
  }
}

// 辅助函数
function cross(a, b) {
  return [
    a[1] * b[2] - a[2] * b[1],
    a[2] * b[0] - a[0] * b[2],
    a[0] * b[1] - a[1] * b[0]
  ];
}

function dot(a, b) {
  return a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
}
// 构建 BVH (简化版)
function buildBVH(triangles) {
    function calculateBoundingBox(triangles) {
        let minX = Infinity, minY = Infinity, minZ = Infinity;
        let maxX = -Infinity, maxY = -Infinity, maxZ = -Infinity;

        for (const triangle of triangles) {
            minX = Math.min(minX, triangle.v0[0], triangle.v1[0], triangle.v2[0]);
            minY = Math.min(minY, triangle.v0[1], triangle.v1[1], triangle.v2[1]);
            minZ = Math.min(minZ, triangle.v0[2], triangle.v1[2], triangle.v2[2]);

            maxX = Math.max(maxX, triangle.v0[0], triangle.v1[0], triangle.v2[0]);
            maxY = Math.max(maxY, triangle.v0[1], triangle.v1[1], triangle.v2[1]);
            maxZ = Math.max(maxZ, triangle.v0[2], triangle.v1[2], triangle.v2[2]);
        }

        return new BoundingBox([minX, minY, minZ], [maxX, maxY, maxZ]);
    }

    function buildNode(triangles, triangleIndices) {
        if (triangleIndices.length === 0) {
            return null;
        }

        // Create bounding box for the triangles in this node
        const boundingBox = calculateBoundingBox(triangleIndices.map(index => triangles[index]));

        // If the number of triangles is small enough, create a leaf node
        if (triangleIndices.length <= 4) { // Arbitrary threshold
            return new BVHNode(boundingBox, null, null, triangleIndices); // Store indices in leaf node
        }

        // Determine split axis (e.g., longest axis of the bounding box)
        const extentX = boundingBox.max[0] - boundingBox.min[0];
        const extentY = boundingBox.max[1] - boundingBox.min[1];
        const extentZ = boundingBox.max[2] - boundingBox.min[2];

        let splitAxis = 0; // X axis
        if (extentY > extentX && extentY > extentZ) {
            splitAxis = 1; // Y axis
        } else if (extentZ > extentX && extentZ > extentY) {
            splitAxis = 2; // Z axis
        }

        // Calculate split position (e.g., midpoint along the split axis)
        let splitPosition = (boundingBox.min[splitAxis] + boundingBox.max[splitAxis]) / 2;

        // Partition triangles based on the split position
        const leftTriangleIndices = [];
        const rightTriangleIndices = [];

        for (const triangleIndex of triangleIndices) {
            const triangle = triangles[triangleIndex];
            let centroid = [(triangle.v0[0] + triangle.v1[0] + triangle.v2[0]) / 3,
                            (triangle.v0[1] + triangle.v1[1] + triangle.v2[1]) / 3,
                            (triangle.v0[2] + triangle.v1[2] + triangle.v2[2]) / 3];

            if (centroid[splitAxis] < splitPosition) {
                leftTriangleIndices.push(triangleIndex);
            } else {
                rightTriangleIndices.push(triangleIndex);
            }
        }

        // Recursively build child nodes
        const leftChild = buildNode(triangles, leftTriangleIndices);
        const rightChild = buildNode(triangles, rightTriangleIndices);

        return new BVHNode(boundingBox, leftChild, rightChild, null);
    }

    const triangleIndices = Array.from({ length: triangles.length }, (_, i) => i); // Indices of all triangles
    return buildNode(triangles, triangleIndices);
}

// 光线与 BVH 相交测试
function traverseBVH(rayOrigin, rayDirection, bvhNode, triangles) {
    if (!bvhNode) {
        return null;
    }

    if (!bvhNode.boundingBox.intersects(rayOrigin, rayDirection)) {
        return null; // No intersection with this node
    }

    if (bvhNode.triangleIndex) { // Leaf node: Intersect with triangles
        let closestIntersection = null;
        let closestDistance = Infinity;

        for (const triangleIndex of bvhNode.triangleIndex) {
            const triangle = triangles[triangleIndex];
            const intersectionDistance = triangle.intersects(rayOrigin, rayDirection);

            if (intersectionDistance !== null && intersectionDistance < closestDistance) {
                closestDistance = intersectionDistance;
                closestIntersection = {
                    distance: intersectionDistance,
                    triangle: triangle
                };
            }
        }

        return closestIntersection;
    }

    // Internal node: Traverse children
    const leftIntersection = traverseBVH(rayOrigin, rayDirection, bvhNode.left, triangles);
    const rightIntersection = traverseBVH(rayOrigin, rayDirection, bvhNode.right, triangles);

    if (leftIntersection && rightIntersection) {
        return leftIntersection.distance < rightIntersection.distance ? leftIntersection : rightIntersection;
    } else if (leftIntersection) {
        return leftIntersection;
    } else {
        return rightIntersection;
    }
}

这段代码构建了一个简单的 BVH，并提供了光线相交测试的函数。注意，这只是一个简化版本，省略了很多优化细节。

接下来，我们需要 WGSL (WebGPU Shading Language) 来编写光线追踪的着色器代码。

struct Ray {
    origin: vec3f,
    direction: vec3f,
};

struct HitPayload {
    hit: bool,
    distance: f32,
    normal: vec3f,
};

// Assuming BVH is stored in a texture or buffer
// This is a placeholder, the actual implementation depends on how you store the BVH

@group(0) @binding(0) var<uniform> cameraPosition: vec3f;
@group(0) @binding(1) var<storage, read_write> outputTexture: array<vec4f>;

// Placeholder for BVH traversal function
fn traverseBVH(ray: Ray) -> HitPayload {
    var payload: HitPayload;
    payload.hit = false;
    payload.distance = 100000.0; // Large distance
    payload.normal = vec3f(0.0);

    // Implement your BVH traversal logic here
    // This is a simplified example
    // For a real implementation, you would need to fetch BVH nodes from a texture or buffer
    // and perform intersection tests with bounding boxes and triangles

    return payload;
}

@compute @workgroup_size(8, 8, 1)
fn main(@builtin(global_invocation_id) global_id: vec3u) {
    let width = 512.0; // Replace with your actual width
    let height = 512.0; // Replace with your actual height

    let uv = vec2f(f32(global_id.x) / width, f32(global_id.y) / height);

    // Create a ray from the camera
    let rayDirection = normalize(vec3f(uv.x * 2.0 - 1.0, uv.y * 2.0 - 1.0, -1.0)); // Simplified perspective projection
    let ray: Ray = Ray(cameraPosition, rayDirection);

    // Perform ray tracing
    let hitPayload = traverseBVH(ray);

    // Shade the pixel based on the hit information
    var color: vec4f;
    if (hitPayload.hit) {
        color = vec4f(hitPayload.normal * 0.5 + vec3f(0.5), 1.0); // Simple normal shading
    } else {
        color = vec4f(0.0, 0.0, 0.0, 1.0); // Black background
    }

    // Store the result in the output texture
    let index = global_id.y * u32(width) + global_id.x;
    outputTexture[index] = color;
}

这段 WGSL 代码只是一个框架，你需要根据你存储 BVH 的方式来实现 traverseBVH 函数。这里只是简单地模拟了光线追踪的过程，并没有真正进行场景相交测试。

四、性能调优：让光追跑得飞起！

光线追踪的性能调优是一个永恒的话题。这里给出一些常见的优化技巧：

  • BVH 构建优化： 选择合适的分裂策略 (SAH)，优化内存布局，减少 BVH 的深度。
  • WGSL 代码优化： 避免不必要的计算，使用 select 函数代替 if 语句，尽量使用向量化操作。
  • 数据传输优化： 减少 CPU 和 GPU 之间的数据传输，尽量将数据存储在 GPU 内存中。
  • 并行化优化： 充分利用 GPU 的并行计算能力，合理设置 workgroup_size。

咱们来具体聊聊：

优化方向 优化手段 效果
BVH 构建 使用 Surface Area Heuristic (SAH) 来决定节点分裂的位置，而不是简单地使用中点。 显著减少光线穿过 BVH 的平均开销。SAH 考虑了分裂后左右子树的表面积和包含的三角形数量，选择最优的分裂方案。
BVH 构建 使用空间中位数分割 (Spatial Median Split)。 不仅考虑对象的质心，还考虑对象占据的空间范围。 这在对象分布不均匀时，可以创建更平衡的 BVH。 显著减少光线穿过 BVH 的平均开销,尤其是在对象分布不均匀的时候。
WGSL 使用 select 函数代替 if 语句。 select 函数是向量化的，可以并行执行，而 if 语句可能会导致分支发散，降低性能。 提高着色器代码的执行效率，减少分支发散带来的性能损失。
WGSL 避免不必要的计算。 例如，如果只需要判断光线是否与三角形相交，而不需要计算交点坐标，那么就不要计算交点坐标。 减少着色器代码的计算量，提高执行效率。
WGSL 尽量使用向量化操作。 WebGPU 的着色器是基于 SIMD (Single Instruction, Multiple Data) 架构的，可以同时处理多个数据。 因此，尽量使用向量化操作，例如 vec3f 和 vec4f，而不是标量操作。 提高着色器代码的执行效率，充分利用 GPU 的并行计算能力。
数据传输 将 BVH 数据存储在 GPU 内存中，例如使用纹理 (Texture) 或存储缓冲区 (Storage Buffer)。 避免每次光线追踪都从 CPU 传输 BVH 数据到 GPU。 减少 CPU 和 GPU 之间的数据传输，提高光线追踪的效率。
并行化 合理设置 workgroup_size。 workgroup_size 决定了每个工作组 (Workgroup) 中有多少个工作项 (Workitem)。 选择合适的 workgroup_size 可以充分利用 GPU 的并行计算能力。 提高光线追踪的并行度，充分利用 GPU 的计算资源。通常需要根据具体的硬件和场景来调整 workgroup_size。
内存布局 优化 BVH 节点的内存布局，例如使用结构体数组 (Array of Structures, AOS) 或结构体数组的数组 (Array of Array of Structures, AoAoS) 来提高内存访问的局部性。 提高内存访问的效率，减少缓存未命中 (Cache Miss) 的概率。
硬件加速 某些 GPU 提供了硬件加速的光线追踪功能，例如 NVIDIA 的 RTX 和 AMD 的 Radeon RX。 如果你的 GPU 支持硬件加速的光线追踪，那么可以考虑使用这些功能来提高光线追踪的性能。 不过 WebGPU 目前还没有暴露这些硬件加速的接口。 理论上可以显著提高光线追踪的性能，但需要使用特定的 API 和驱动程序。

五、总结：光追之路，道阻且长，行则将至！

WebGPU 光线追踪的实现和性能调优是一个复杂的过程，需要深入理解光线追踪的原理，以及 WebGPU 的 API 和着色器语言。希望今天的脱口秀（误）能帮助大家入门 WebGPU 光线追踪，并在未来的学习和实践中取得更大的进步！

记住，光追之路，道阻且长，行则将至！ 只要坚持不懈，就能最终征服光线追踪这座大山！

感谢各位的观看，咱们下期再见！ (挥手)

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注