JS `WebGPU` `Ray Tracing` (光线追踪) 的实现与性能挑战

各位靓仔靓女，晚上好！我是今晚的主讲人，很高兴能和大家一起聊聊 WebGPU 中的 Ray Tracing（光线追踪）。

今天咱们的主题是 “JS WebGPU Ray Tracing：实现与性能挑战”。听到这个题目，是不是感觉有点 high-tech？别怕，咱们今天就把这玩意儿给它掰开了揉碎了，用最接地气的方式，让大家都能听明白，甚至能上手撸两行代码。

开场白：为啥要搞 WebGPU Ray Tracing？

首先，咱们得明白一个问题：为啥要在 WebGPU 里面搞 Ray Tracing？答案很简单：因为它炫酷啊！开玩笑的，当然是因为它能带来更逼真的渲染效果。传统的 rasterization (光栅化) 技术，虽然速度快，但模拟光照效果时，总是会遇到各种各样的问题，比如阴影不真实、反射不准确等等。而 Ray Tracing 就不一样了，它通过模拟光线的传播路径，能够更真实地模拟光照效果，从而产生更加逼真的图像。

更重要的是，WebGPU 给了我们一个在 Web 上实现高性能图形渲染的机会。以前在 Web 上搞复杂图形应用，那简直就是噩梦。现在有了 WebGPU，我们可以充分利用 GPU 的并行计算能力，实现以前想都不敢想的效果。

Ray Tracing 的核心概念：一句话概括

Ray Tracing 的核心概念其实很简单：就是从摄像机发射出无数条光线，然后追踪这些光线的路径，直到它们击中场景中的物体，或者到达最大追踪深度。通过计算光线与物体的交点颜色，就可以得到最终的图像。

WebGPU Ray Tracing 的基本流程：四步走

在 WebGPU 中实现 Ray Tracing，大致可以分为以下几个步骤：

构建 Acceleration Structure (加速结构)： 这是 Ray Tracing 的关键。加速结构可以帮助我们快速找到光线与场景中物体的交点。常用的加速结构有 BVH (Bounding Volume Hierarchy) 和 KD-Tree。
编写 Ray Generation Shader (光线生成着色器)： 这个着色器的作用是生成从摄像机发射的光线。
编写 Intersection Shader (相交着色器)： 这个着色器用于计算光线与物体的交点。
编写 Ray Tracing Pipeline (光线追踪管线)： 将上面的各个部分组合起来，形成一个完整的 Ray Tracing 管线。

代码示例：一个简单的 WebGPU Ray Tracing Demo

接下来，咱们来撸一段简单的代码，演示一下如何在 WebGPU 中实现 Ray Tracing。为了简单起见，咱们只渲染一个球体。

// 1. 初始化 WebGPU 设备
async function initWebGPU() {
  if (!navigator.gpu) {
    throw new Error("WebGPU is not supported on this browser.");
  }

  const adapter = await navigator.gpu.requestAdapter();
  if (!adapter) {
    throw new Error("No appropriate GPUAdapter found.");
  }

  const device = await adapter.requestDevice();
  return device;
}

// 2. 创建 Acceleration Structure (简化版本，只包含一个球体)
function createAccelerationStructure(device) {
  const vertexBuffer = device.createBuffer({
    size: 12, // 3 vertices * 4 bytes/vertex
    usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
    mappedAtCreation: true,
  });
  new Float32Array(vertexBuffer.getMappedRange()).set([
    0.0, 0.5, 0.0,
    -0.5, -0.5, 0.0,
    0.5, -0.5, 0.0,
  ]);
  vertexBuffer.unmap();

  const indexBuffer = device.createBuffer({
    size: 6, // 2 triangles * 3 indices
    usage: GPUBufferUsage.INDEX | GPUBufferUsage.COPY_DST,
    mappedAtCreation: true,
  });
  new Uint16Array(indexBuffer.getMappedRange()).set([
    0, 1, 2
  ]);
  indexBuffer.unmap();

  // 这里我们简化处理，直接返回顶点和索引数据。
  return {
    vertexBuffer,
    indexBuffer,
    vertexCount: 3,
    indexCount: 3
  };
}

// 3. 创建 Ray Generation Shader
const rayGenShaderCode = `
@group(0) @binding(0) var<storage, read_write> outputBuffer: array<f32>;

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

fn createRay(uv : vec2f) -> Ray {
    let aspectRatio = 1.0; //假设画面是正方形
    let cameraOrigin = vec3f(0.0, 0.0, -3.0);
    let cameraDirection = normalize(vec3f(uv.x * aspectRatio, uv.y, 1.0));
    return Ray(cameraOrigin, cameraDirection);
}

fn hitSphere(center : vec3f, radius : f32, ray : Ray) -> f32 {
  let oc = ray.origin - center;
  let a = dot(ray.direction, ray.direction);
  let b = 2.0 * dot(oc, ray.direction);
  let c = dot(oc, oc) - radius * radius;
  let discriminant = b * b - 4.0 * a * c;
  if (discriminant < 0.0) {
      return -1.0; // no hit
  } else {
      return (-b - sqrt(discriminant)) / (2.0 * a);
  }
}

@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) global_id : vec3u) {
  let width = 512;
  let height = 512;
  let uv = vec2f(f32(global_id.x) / f32(width) * 2.0 - 1.0, f32(global_id.y) / f32(height) * 2.0 - 1.0);
  let ray = createRay(uv);
  let t = hitSphere(vec3f(0.0, 0.0, 0.0), 1.0, ray);

  var color = vec3f(0.0, 0.0, 0.0);
  if (t > 0.0) {
      let normal = normalize(ray.origin + ray.direction * t);
      color = normal * 0.5 + 0.5;
  } else {
      let t = 0.5 * (ray.direction.y + 1.0);
      color = mix(vec3f(1.0, 1.0, 1.0), vec3f(0.5, 0.7, 1.0), t);
  }

  let index = (global_id.y * width + global_id.x) * 4;
  outputBuffer[index] = color.x;
  outputBuffer[index + 1] = color.y;
  outputBuffer[index + 2] = color.z;
  outputBuffer[index + 3] = 1.0;
}
`;

// 4. 创建 Compute Pipeline
async function createComputePipeline(device) {
  const computeShaderModule = device.createShaderModule({
    code: rayGenShaderCode,
  });

  const computePipeline = device.createComputePipeline({
    layout: 'auto',
    compute: {
      module: computeShaderModule,
      entryPoint: 'main',
    },
  });

  return computePipeline;
}

// 5. 创建 Buffer 用于存储结果
async function createOutputBuffer(device, width, height) {
  const buffer = device.createBuffer({
    size: width * height * 4 * 4, // width * height * 4 (RGBA) * 4 bytes (f32)
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
  });
  return buffer;
}

// 主函数
async function main() {
  const device = await initWebGPU();
  const accelerationStructure = createAccelerationStructure(device);
  const computePipeline = await createComputePipeline(device);

  const width = 512;
  const height = 512;
  const outputBuffer = await createOutputBuffer(device, width, height);

  // 创建 Bind Group
  const bindGroup = device.createBindGroup({
    layout: computePipeline.getBindGroupLayout(0),
    entries: [
      {
        binding: 0,
        resource: {
          buffer: outputBuffer,
        },
      },
    ],
  });

  // 创建 Command Encoder
  const commandEncoder = device.createCommandEncoder();

  // 创建 Compute Pass
  const computePass = commandEncoder.beginComputePass();
  computePass.setPipeline(computePipeline);
  computePass.setBindGroup(0, bindGroup);
  computePass.dispatchWorkgroups(width / 8, height / 8); // 根据 workgroup_size 调整
  computePass.end();

  // 将结果从 GPU 拷贝到 CPU
  const readBuffer = device.createBuffer({
    size: width * height * 4 * 4,
    usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
  });
  commandEncoder.copyBufferToBuffer(outputBuffer, 0, readBuffer, 0, width * height * 4 * 4);

  // 提交命令
  device.queue.submit([commandEncoder.finish()]);

  // 读取数据
  await readBuffer.mapAsync(GPUMapMode.READ);
  const data = new Float32Array(readBuffer.getMappedRange());

  // 创建 Canvas 并显示图像
  const canvas = document.createElement('canvas');
  canvas.width = width;
  canvas.height = height;
  document.body.appendChild(canvas);
  const ctx = canvas.getContext('2d');
  const imageData = ctx.createImageData(width, height);

  for (let i = 0; i < width * height; i++) {
    imageData.data[i * 4] = data[i * 4] * 255;
    imageData.data[i * 4 + 1] = data[i * 4 + 1] * 255;
    imageData.data[i * 4 + 2] = data[i * 4 + 2] * 255;
    imageData.data[i * 4 + 3] = data[i * 4 + 3] * 255;
  }

  ctx.putImageData(imageData, 0, 0);

  readBuffer.unmap();
}

main();

代码解释：逐行分析

上面的代码虽然简单，但是包含了 WebGPU Ray Tracing 的基本要素。咱们来逐行分析一下：

initWebGPU(): 初始化 WebGPU 设备。这是所有 WebGPU 程序的第一步。
createAccelerationStructure(): 创建 Acceleration Structure。在这个例子中，我们简化了处理，没有真正构建 BVH 或 KD-Tree，只是简单地返回了球体的顶点和索引数据。真正的 Ray Tracing 应用中，需要使用更复杂的加速结构。
rayGenShaderCode: 这是 Ray Generation Shader 的代码。它使用 WGSL 编写，用于生成光线，并计算光线与球体的交点颜色。
createComputePipeline(): 创建 Compute Pipeline。将 Ray Generation Shader 编译成 Compute Pipeline。
createOutputBuffer(): 创建 Buffer 用于存储结果。
main(): 主函数。在这个函数中，我们创建 Bind Group，Command Encoder，Compute Pass，然后提交命令，并将结果从 GPU 拷贝到 CPU，最后显示在 Canvas 上。

性能挑战：优化之路漫漫

虽然 WebGPU 给了我们一个在 Web 上实现高性能图形渲染的机会，但是 Ray Tracing 的性能仍然是一个巨大的挑战。优化 Ray Tracing 的性能，需要从多个方面入手：

Acceleration Structure 的构建： 选择合适的加速结构，并优化其构建算法，可以显著提高 Ray Tracing 的性能。
Shader 的优化： 优化 Ray Generation Shader 和 Intersection Shader 的代码，减少计算量。
并行计算： 充分利用 GPU 的并行计算能力，将计算任务分配到多个 workgroup 中。
降噪： Ray Tracing 的结果通常会包含很多噪点。使用降噪算法，可以减少噪点，提高图像质量。

一些常用的优化技巧：

优化技巧	描述	适用场景

发表回复 取消回复

发表回复取消回复