WebGPU Pathtracer Discussion Of Implementation Details

[TheBlek]

Disclaimer

This PR is not meant to be merged, but rather to open discussion on some of the implementation choices made here before they are extracted in a separate branch for a clean merge.

This is a pathtracer implementation that I worked on and off for the past month or so. It deviates from current webgpu-pathtracer branch quite a bit. Yet I think some of the implementation details can be merged. They are discussed separately below.

Transmissive materials

In WebGL transmissive materials are currently implemented use Lambert BTDF for rough materials and some modification of it for glossy materials. This seems to not be how its usually done. Everywhere I've seen, they are implemented using a GGX distribution (same as specular) that support roughness parameterization. From what I understand it better models how light behaves on refractions. However, this distribution is not energy preserving - it reflects less light than receives - especially with higher roughness. Here we can create a texture similar to current Turquin texture with extra dimension - for ior.

Care was taken to handle total internal reflection:

1. It returns 0 if microsurface normal forces refraction to be TIR
2. Proper dielectric fresnel is taken into account when calculating lobe weights to take TIR into account for reflection and refraction probabilities

Additionally, I believe there is some kind of bug with refraction at glancing angles since its overly bright (lean too heavily into the surface's color). Initially I was thinking this is due to ggx energy compensation not being implemented, but it seems like an integration bug.

WebGL	This PR

Filter Glossy

Blender was the reference for filter glossy implementation. Formula here is based on the minimal pdf value of scatters in current path. I think it would be good to give user a parameter that works as in Blender. What do you think?

Optimisations

I've seen that you noted bad BVH traversal performance in #768. Most of those are algorithmic or heavy changes though great and should be tried. Things discussed here are small changes made after profiling and experiments. Now I'm curious to see if you see those mentioned slowdowns in this branch.

Benchmarking utility

When talking about performance its essential to be able to measure it, hence the benchmarking utility at scripts/benchmark-scenes.js. It runs pathtracer at a number of sample counts, measuring time and MSE/RMSE/PSNR error of generated images in comparison with the golden one.

Some of the features missing:

1. Multiple iterations with result averaging for better reproducibility
2. Warmup iterations
3. Automatically adjust iteration per frame value to fully utilize gpu. Right now it requires hand tuning to find best results
4. Calculated metrics are hard to interpret, so using something like Flip should be more informative. Maybe we could automate calculation of its metric/comparison images if user has the tool installed?
5. Cleaning up. It was purely vibecoded :)

When I started working on performance, WebGPU version was significantly slower on bigger scenes (Sponza) specifically because of ray tracing kernel.

Storing material in local memory

Turns out storing whole material struct inside private thread memory is not such a good idea, since it occupies precious vector registers and decreases possible parallelism. So storing just an index of the current material makes things better. See diff to PathtracerBVHComputeData.js file for this. (Couldn't figure out how to create a proper link)

Double BVH stack cost

Allocating 60 slots for the traversal stack on each BVH level is a lot of memory and takes up a lot of vector registers forcing kernel to spill them into scratch memory, which is slow. To mediate that I did 2 things:

1. Cut each traversal stack in half to 30 since there are two of them now. In practice, I think we can make lower level stack bigger at the cost of object-level stack. Or, another idea, determine maximal traversal depth when scene changes and adjust the values
2. Put one of the stack into workgroup memory since it is a separate resource. This shifts memory allocation from one place to another and moving second stack proved to be useless since it limited parallelism because of limited workgroup space.

Wavefront architecture considerations

In my opinion, proper comparison between megakernel and wavefront architecture (and in general) pathtracers would compare how fast they converge on a correct image. To inspect that I plotted MSE over time of different aproaches:

What's different in new wavefront implementation? Its basically an implementation of this paper with material stage unified in one kernel. This unification was done because WebGPU does not allow us to run kernels from that stage simultaneously (because it can't prove we wont be accessing same data in path data array). And it saves us the overhead of couple extra queues.

Overall changes:

• A fixed number of paths are traced at a time, If path terminates, a new is generated on next step.
• Almost all data lies in one array with one entry corresponding to a path being traced

Data flow in kernels is quite hard to wrap your head around because it has a material stage before ray tracing stage. I'll try to explain it here as best as I can.

First iteration:

1. Logic kernel detects that path is terminated, marks it for generation and updates result texture; This causes only a single black write per pixel since all paths terminated later have a valid color.
2. Material kernel generates a new ray, places it into the ray tracing queue, remembers index of the ray to fetch results
3. Ray intersection kernel

Second iteration:

1. Logic kernel fetches intersection result;
   1. If object is hit, sample a light direction, save that path and pdf for later in the data array, save intersection point for later
   2. If no object is hit, mark ray as terminated and sample environment/background map, updates result texture
2. Material kernel
   1. If path is terminated, generate a new ray and places it into the ray tracing queue
   2. If path is not terminated, calculate material reaction at the point of ray intersection, place new direction ray into the ray tracing queue, calculate and save pdfs for mis, place shadow ray (generated in logic kernel) into shadow ray intersection queue.
3. All queued rays are traced

Third iteration and so on is basically the same.

I meant to fix black pixel write quirk somehow but never got around. Maybe initialize resultColor to have alpha = 0 and have this write a noop?

Why is new wavefront implementation performs better? My thinking is following:

Ray tracing code is tightly isolated in their own kernels to minimize vector register usage. Its critical for this code as its latency bound because of a tree traversal as gpu can't predict next node index. At least I believe so. And these kind of workflows benefit most from executing multiple warps on a single ALU. Lowering register count on ray tracing unit allowed running almost 2x more warps simultaneously on my machine. Which in turn improved performance.

While I'm not saying that new implementation is the one to merge, but it proves we clearly can do better than the current one. I would like to try isolating the performance benefit changes part while keeping the complexity low.

What do you think of this new kernel structure?

Custom wgsl structs

WGSLStructTypeNode allows for custom wgsl structs to be treated as TSL-native structs. This allows us to bundle queue data with meta information (length, capacity, etc.) and therefore allows us to get rid of queueSizes buffer and reduce used buffers count. This was required for NEE implementation in current wavefront architecture.

Better sampling

When researching topic of performance I stumbled upon a couple papers describing faster VNDF sampling algorithms in comparison with the current implemented:

1. https://arxiv.org/pdf/2306.05044
2. https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (seems to be an improvement on the first one)

Scene size problem

After testing this pathtracer on a variety of lego models it is apparent that supported scene size is much smaller than WebGL's. This is due to 128Mb buffer limit (and perhaps two level bvh structure?). Here, I think are a couple ways to ease the problem none of which are implemented:

1. Pack attribute data. UVs are now vec4 which is totally unnecessary and I suspect can be packed inside 32 bits without any major problems in quality. This article briefly discusses using 16 bit float and int values in context of vector register optimization. Moreover, it seems that the normal vector can also be packed: to 8 bytes without losing precision (if thought of in spherical coordinates) or even 3 bytes. Though I'm pretty sure renders will be affected by the latter.
2. Pack bvh node data. This paper demonstrates that it is possible to compress BVH description. And performance will even increase. In general, 8-wide bvh adoption will be beneficial in terms of memory and time.
3. Conditional support for bigger buffers. Probably the easiest to implement.

[gkjohnson]

Thanks! I'll leave my comments below - let me know if this is the kind of feedback you're looking for. I haven't looked through all the code but I've looked at some of the pieces relevant to the comments.

And let me know how you'd like to handle some of these changes. There's a lot of good improvements listed here - I can plan to add some of the new thoughts to the plans list in #713, as well. And some of these I can take a look at when looking into the BVH traversal improvements

Transmission

In WebGL transmissive materials are currently implemented use Lambert BTDF for rough materials and some modification of it for glossy materials. This seems to not be how its usually done.

Yes this was not really rigorously handled in the WebGL version of the path tracer. It would be good to get a better method working, which it sounds like you've been looking in to.

However, this distribution is not energy preserving - it reflects less light than receives - especially with higher roughness. Here we can create a texture similar to current Turquin texture with extra dimension - for ior.

I'm less familiar with all the details of the energy preserving logic but this sounds like a fine solution. I think I recall seeing in the original paper that this was suggested.

Additionally, I believe there is some kind of bug with refraction at glancing angles since its overly bright (lean too heavily into the surface's color). Initially I was thinking this is due to ggx energy compensation not being implemented, but it seems like an integration bug.

I'm not sure if this is exactly what you're referring to, but one thing to bear in mind is that the WebGLPathTracer is sampling the "background" texture rather than the environment map if the path has been fully transmissive. The WebGPU version looks like it's sampling the environment map, which will the glass look brighter. See here how the color showing through the surface is black rather than the env map in the WebGL screen shot:

FilterGlossyFactor

Blender was the reference for filter glossy implementation. Formula here is based on the minimal pdf value of scatters in current path. I think it would be good to give user a parameter that works as in Blender. What do you think?

This sounds good. I think generally following what Blender does is a good direction. When I initially implemented this I hadn't gone digging through the blender code quite as much.

Optimizations

re: Storing material in local memory

Good to know! The material fix seems like an easy fix. Have you seen tangible performance improvements from this?

re: Double BVH stack cost

This is something I figured might be an issue. I think it should be easy enough to share a single stack between them, which I'll take a look at while I try to fix the overlapping bounds.

re: Wavefront Optimization

I'll have to take a deeper look at this later but I'm noticing some new render bugs when enabling the wavefront mode. Should I wait for these to be addressed before taking a look at the code?

re: Custom wgsl structs

Freeing up some storage buffers is great! Nice work. Is there a reason we can't use the built-in "StructTypeNode", though? For example:

const hitQueueStruct = new StructTypeNode( {
  start: 'u32',
  end: 'u32',
  _padding: 'array<u32, 2>',

  elements: `array<${ queuedHitStruct.name }>`,
}, 'HitQueue' );

This will work as long as we don't call something like getMemberLayout or getLength.

re: Better sampling

Faster sampling functions always sound great.

Scene size problem

This is something I think would be good to tackle later. Right now we can fallback to requesting increased memory limits on the WebGPU context to load higher detail models. Tackling CWBVH and packed attributes is something on my list but I'm thinking it's best to get a stable compute data class working, first.

Benchmarking utility

Benchmarking is always great and I think this will be useful to be able to run on PRs, assuming we can get the WebGPU path tracer running in Github CI. I'm wondering if this code can be pulled out into a command line rendering utility? Then it can handle benchmarking and serve as a headless rendering utility.

[TheBlek]

WebGLPathTracer is sampling the "background" texture rather than the environment map if the path has been fully transmissive.

Oh, I will look into this, maybe that's causing the glass brightness on the octopus scene.

Have you seen tangible performance improvements from this?

Yes, traversal optimisations listed took webgpu version from being significantly slower that webgl to being ~15% faster on Sponza.

I think it should be easy enough to share a single stack between them, which I'll take a look at while I try to fix the overlapping bounds.

Sharing a single stack sounds great for simplifying the code. However notice that in my implementation one of the stacks (so half of the memory for stacks) is allocated in workgroup memory. This is done specifically to reduce the number of vector registers allocated. This is a significant part of the optimisations.

And some of these I can take a look at when looking into the BVH traversal improvements.

Yeah, since you will be looking into BVH traversal improvements I think its only natural you could try optimisation listed here. Let me know you would like a better explanation of how they work.

Is there a reason we can't use the built-in "StructTypeNode", though?

Well, I couldn't get them to work at the time of writing. I believe it's still not possible as per this issue: mrdoob/three.js#33041 (see point 7)

I'm noticing some new render bugs when enabling the wavefront mode

Can you point out those bugs? I believe it should work as expected beside the first black frame. I will get around to fix this as soon as I can.

I'm wondering if this code can be pulled out into a command line rendering utility? Then it can handle benchmarking and serve as a headless rendering utility.

Yes, in fact, it already supports headless execution. However, I'm not sure we will be able to run this inside Github CI, at least for free, as it does not seem to provide GPU on free runners. For now, making a simple utility to generate a diff report between current version and a specific branch would be enough for now I think. Just like it works in three-mesh-bvh.

I want to tackle benchmarking utility next. Let me know if you think it should be handled in a different order.

[gkjohnson]

Yes, traversal optimisations listed took webgpu version from being significantly slower that webgl to being ~15% faster on Sponza.

🎉

in my implementation one of the stacks (so half of the memory for stacks) is allocated in workgroup memory

This is good to know - I didn't realize moving to workgroup memory would have such a big impact.

One thing worth noting is that three.js recently changed variables to be declared as global (and therefore likely "private") when declared via TSL (see mrdoob/three.js#33302). This won't impact this project so much I don't think since we're using WGSL directly but I'm wondering if that change would have a negative impact on other three.js shaders? Maybe worth mentioning this potential performance impact at the three.js repo.

Yeah, since you will be looking into BVH traversal improvements I think its only natural you could try optimisation listed here. Let me know you would like a better explanation of how they work.

To run down what I'll look into in my updates, then:

• Faster traversal using a an improved BVH traversal.
• Store material index rather than full material struct in global variable.
• Use a pre-allocated workgroup-space array rather than locally declared array in the function.
• (sometime later) use compressed-wide BVH, pack BVH memory further
• (sometime later) improve packing of attribute data.

Is there anything I'm missing?

Well, I couldn't get them to work at the time of writing. I believe it's still not possible as per this issue: mrdoob/three.js#33041 (see point 7)

We're using arrays and other struct references successfully in our struct definitions here - the caveats are that "getLength" needs to be overwritten because the original implementation will choke on array and struct definitions. And the struct definition will not be implicitly included so it needs to be declared as a dependency some other way. I'd have to see exactly what issue you're running in to but maybe this helps explain things. It's definitely not ideal, at the moment.

Can you point out those bugs? I believe it should work as expected beside the first black frame. I will get around to fix this as soon as I can.

In "furnace_test" I'm seeing that the bottom row of the screen is noticeably darker than then rest after enabling wavefront. It's a bit hard to catch in a screenshot but you can see the artifacts in the bottom:

And in the index.html demo I'm seeing the floor not render properly - the fade out seems cutoff in a spiral type shape:

megakernel	wavefront

However, I'm not sure we will be able to run this inside Github CI, at least for free, as it does not seem to provide GPU on free runners

Might be worth a try. Three.js is rendering WebGPU with (I believe) free-tier Github CI. As far as I understand puppeteer falls back to a CPU-based implementation of WebGPU - so our mileage (and performance) may vary but we can see. It wouldn't be great for performance but help us do a per-pixel diffs for regression testing.

I want to tackle benchmarking utility next. Let me know if you think it should be handled in a different order.

Sounds great to me! If it's possible to make cli-enable rendering utility for users in the same pass that would be a cool bonus.

[gkjohnson]

Can you point out those bugs? I believe it should work as expected beside the first black frame. I will get around to fix this as soon as I can.

I'm also seeing that the wavefront path tracer is running with a slower framerate than the megkernel one - which is the opposite of what I'd expect and what I'm seeing in the examples branch I worked on from #768.

[TheBlek]

We're using arrays and other struct references successfully in our struct definitions here

Oh, yeah, I will try out with those. Perhaps I was using three.js's default struct node instead. Or maybe it was because of unbounded arrays? Anyway, I'll look into it.

And in the index.html demo I'm seeing the floor not render properly - the fade out seems cutoff in a spiral type shape:

Seeing this too. Seems like a ray traversal bug or scene data corruption of some sort. Let me look into this and polish things.

I'm also seeing that the wavefront path tracer is running with a slower framerate than the megakernel one

Different architectures are a hard thing to compare since they are computing different work per frame. New wavefront runs a fixed number of rays per frame (limited only by 128Mb webgpu buffer for state management) and megakernel workload depends on tile size. This is why I tried to compare performance by plotting error over time to see how fast wavefront converges. You can see in the PR description that its not exactly faster than megakernel implementation even though tracing ~10-15% more rays:

I didn't realize moving to workgroup memory would have such a big impact.

This is because stack is really big. I'm pretty sure it will have almost to none effect on moving one float variable, for example. Regarding the three.js PR - my context is limited - but making every variable global (in private scope) could indeed harm performance in big shaders. This also will depend on the GPU and driver compiler. But I'm no expert on this, so would want to be careful affecting a big project.

To run down what I'll look into in my updates, then:
Faster traversal using a an improved BVH traversal.
Store material index rather than full material struct in global variable.
Use a pre-allocated workgroup-space array rather than locally declared array in the function.
(sometime later) use compressed-wide BVH, pack BVH memory further
(sometime later) improve packing of attribute data.
Is there anything I'm missing?

Yes, that's about it. Additionally, in this implementation those stacks are 2x smaller. I figured if we previously had one 60 deep stack, having two 30 deep stack would pretty much be on par.

[gkjohnson]

Different architectures are a hard thing to compare since they are computing different work per frame.

I understand but more than convergence time the big benefit of Wavefront path tracer is that the per-frame performance can remain the same regardless of the bounce count, etc, which has never been possible with a megakernel variant. It may even just be a matter of tuning ray counts but I want to make sure this is accounted for in whatever architectural decisions we make so we can avoid locking up the users machine, as has always happened with the WebGL version.

I had originally adjusted the settings of the wavefront path tracer to keep a high framerate when rendering (admitted tuned for my machine, 2021 M1 Pro Macbook) so that's what I'm referring to here. For comparison here's a list of framerates from this branch and the webgpu/examples (algorithm unchanged from webgpu-pathtracer branch) with the "Imaginary Friend Room" demo model with a 1.0 resolution ratio:

	megakernel	wavefront
this branch	~27fps	~35fps
webgpu/examples	~45fps	120+fps (browser limited)

So that's some nice improvement with the megakernel branch but a 70+% degradation in what amounts to browser & UI responsiveness with the wavefront. Again, this could be a matter of just adjusting some knobs but I think we should aim for defaults that will prioritize framerate, even if that likely means the image won't converge as quickly as it can. That doesn't mean we shouldn't evaluate the wavefront path tracer in a "full power" mode with maximum throughput settings cranked up, as well. We should account for both.

--

One other thing I noticed while testing is that the "sample counter" in the bottom left seems to go up extremely fast - it "feels" like 1 sample per frame? It doesn't look like the image is resolving at the same rate, though. Is this right?

[TheBlek]

@gkjohnson I've fixed visual bugs on wavefront implementation in this branch, so you take a look at the kernels. The main thing there is that ray tracing kernel is isolated as much as possible to allow more parallel execution on one unit. Ordering is a bit weird because MaterialKernel needs to calculate pdf values for both regular and shadow rays.

I had originally adjusted the settings of the wavefront path tracer to keep a high framerate when rendering (admitted tuned for my machine, 2021 M1 Pro Macbook) so that's what I'm referring to here.

I understand that. Imaginary Friend Room does not seem to load uvs for me, but I'm getting 144 fps for both megakernel and old wavefront while ~110 for new wavefront. Perhaps difference is less on my side because I have a bit wider gpu. I believe number of simultaneous rays processed per frame should adjust automatically somehow or be a configurable parameter to achieve high frame rates. This just a matter of settings, I believe.

This setting was not taken into account when benchmarking though, so old wavefront could be faster and data needs to be regenerated. I hope to build an easy to run tool so we can get comparable and reproducible data on this from different machines and scenes.

One other thing I noticed while testing is that the "sample counter" in the bottom left seems to go up extremely fast - it "feels" like 1 sample per frame? It doesn't look like the image is resolving at the same rate, though. Is this right?

Yep, sample counter in the new wavefront reflects how many steps were run without estimating how much samples per pixel are actually computed. Data fetched from the gpu later in the run should be accurate though.

[gkjohnson]

What do you think of this new kernel structure?

I think overall it makes sense to move in the direction you've suggested if I'm understanding things correctly. But I'm having a harder time understanding the changes and tracking down where the performance difference is coming from in the PR with so many additional lines & features like NEE. Would it be possible to make to make a PR with some of these changes for the WaveFront architecture changes specifically?