Skip to content

WebGPUPathTracer: Convert kernels to use wgsl tag fn #748

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gkjohnson merged 7 commits into from
Mar 14, 2026
Merged

WebGPUPathTracer: Convert kernels to use wgsl tag fn #748

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
880d656
WGSLTagFnNode: normalize arguments on build rather than immediately
gkjohnson Mar 10, 2026
57eb9ca
Simplify megakernel
gkjohnson Mar 10, 2026
b0e9d93
Fix 2 wavefront kernels, normalize variable naming
gkjohnson Mar 10, 2026
00931af
Adjust wavefront ray generation kernel
gkjohnson Mar 10, 2026
13320ea
Update last two wavefront kernels
gkjohnson Mar 10, 2026
b498527
Fix zero out buffer kernel
gkjohnson Mar 10, 2026
a0d9744
Remove unneeded arguments
gkjohnson Mar 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/webgpu/PathTracerBackend.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,8 @@ export class PathTracerBackend {
this.sampleCountTarget.name = 'Sample Count';
this.sampleCountTarget.generateMipmaps = false;

this.sampleCountClearKernel = new ZeroOutKernel( { textureType: 'r32uint' } ).setWorkgroupSize( 8, 8, 1 );
this.outputTargetClearKernel = new ZeroOutKernel( { textureType: 'rgba32float' } ).setWorkgroupSize( 8, 8, 1 );
this.sampleCountClearKernel = new ZeroOutKernel().setWorkgroupSize( 8, 8, 1 );
this.outputTargetClearKernel = new ZeroOutKernel().setWorkgroupSize( 8, 8, 1 );

}

Expand Down
62 changes: 29 additions & 33 deletions src/webgpu/compute/PathTracerMegaKernel.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,17 @@ import { DataTexture, Matrix3, Matrix4, Vector2, StorageTexture } from 'three/we
import { ndcToCameraRay } from '../lib/wgsl/common.wgsl.js';
import { ComputeKernel } from './ComputeKernel.js';
import { texture, sampler, uniform, globalId, textureStore } from 'three/tsl';
import { pcgRand2, pcgRand3, pcgInit } from '../nodes/random.wgsl.js';
import { pcgRand2, pcgInit } from '../nodes/random.wgsl.js';
import { getSurfaceRecordFunc, lambertBsdfFunc } from '../nodes/material.wgsl.js';
import { sampleEnvironmentFn, weightedAlphaBlendFn } from '../nodes/sampling.wgsl.js';
import { proxy } from '../lib/nodes/NodeProxy.js';
import { proxy, proxyFn } from '../lib/nodes/NodeProxy.js';
import { wgslTagFn } from '../lib/nodes/WGSLTagFnNode.js';

export class PathTracerMegaKernel extends ComputeKernel {

constructor() {

const parameters = {
const params = {
bvhData: { value: null },

prevOutputTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadOnly(),
Expand Down Expand Up @@ -47,6 +47,9 @@ export class PathTracerMegaKernel extends ComputeKernel {
globalId: globalId,
};

const raycastFirstHitFn = proxyFn( 'bvhData.value.fns.raycastFirstHit', params );
const sampleTrianglePointFn = proxyFn( 'bvhData.value.fns.sampleTrianglePoint', params );

const shader = wgslTagFn/* wgsl */`

fn compute(
Expand Down Expand Up @@ -81,6 +84,9 @@ export class PathTracerMegaKernel extends ComputeKernel {

) -> void {

let transforms = &${ proxy( 'bvhData.value.storage.transforms', params ) };
let materials = &${ proxy( 'bvhData.value.storage.materials', params ) };

let envInfo = EnvironmentInfo(
envMapRotation,
envMapIntensity,
Expand All @@ -102,7 +108,7 @@ export class PathTracerMegaKernel extends ComputeKernel {

// to screen coordinates
let indexUV = offset + globalId.xy;
let targetDimensions = textureDimensions( ${ parameters.outputTarget } );
let targetDimensions = textureDimensions( ${ params.outputTarget } );
if ( indexUV.x >= targetDimensions.x || indexUV.y >= targetDimensions.y ) {

return;
Expand All @@ -112,35 +118,35 @@ export class PathTracerMegaKernel extends ComputeKernel {
let uv = vec2f( indexUV ) / vec2f( targetDimensions );
let ndc = uv * 2.0 - vec2f( 1.0 );

pcgInitialize( indexUV, seed );
${ pcgInit }( indexUV, seed );

// scene ray
var jitter = 2.0 * pcgRand2() / vec2f( targetDimensions.xy );
var ray = ndcToCameraRay( ndc + jitter, cameraToModelMatrix * inverseProjectionMatrix );
var jitter = 2.0 * ${ pcgRand2 }() / vec2f( targetDimensions.xy );
var ray = ${ ndcToCameraRay }( ndc + jitter, cameraToModelMatrix * inverseProjectionMatrix );
ray.direction = normalize( ray.direction );

var resultColor = vec4f( 0, 0, 0, 1 );
var throughputColor = vec3f( 1.0 );

for ( var bounce = 0u; bounce < bounces; bounce ++ ) {

let hitResult = bvh_RaycastFirstHit( ray );
let hitResult = ${ raycastFirstHitFn }( ray );
if ( hitResult.didHit ) {

let object = bvh_transforms.value[ hitResult.objectIndex ];
var material = bvh_materials.value[ object.materialIndex ];
let object = transforms[ hitResult.objectIndex ];
var material = materials[ object.materialIndex ];

// apply per-object colors
material.color *= object.color.rgb;
material.opacity *= object.color.a;

var vertexData = bvh_sampleTrianglePoint( hitResult.barycoord, hitResult.indices.xyz );
var vertexData = ${ sampleTrianglePointFn }( hitResult.barycoord, hitResult.indices.xyz );
vertexData.normal = normalize( transpose( object.inverseMatrixWorld ) * vertexData.normal );
vertexData.position = object.matrixWorld * vertexData.position;

let surface = getSurfaceRecord( material, vertexData, hitResult.side, hitResult.normal, textures, textureSampler );
let surface = ${ getSurfaceRecordFunc }( material, vertexData, hitResult.side, hitResult.normal, textures, textureSampler );

let scatterRec = bsdfSample( - ray.direction, surface );
let scatterRec = ${ lambertBsdfFunc }( - ray.direction, surface );

// white diffuse surface
throughputColor *= scatterRec.color / scatterRec.pdf;
Expand All @@ -152,11 +158,11 @@ export class PathTracerMegaKernel extends ComputeKernel {

if ( bounce > 0u ) {

resultColor = sampleEnvironment( envMap, envMapSampler, envInfo, ray.direction, pcgRand2() ) * vec4f( throughputColor, 1.0 );
resultColor = ${ sampleEnvironmentFn }( envMap, envMapSampler, envInfo, ray.direction, pcgRand2() ) * vec4f( throughputColor, 1.0 );

} else {

resultColor = sampleEnvironment( background, backgroundSampler, backgroundInfo, ray.direction, pcgRand2() );
resultColor = ${ sampleEnvironmentFn }( background, backgroundSampler, backgroundInfo, ray.direction, pcgRand2() );

}

Expand All @@ -166,27 +172,17 @@ export class PathTracerMegaKernel extends ComputeKernel {

}

let sampleCount = textureLoad( ${ parameters.sampleCountTarget }, indexUV ).r + 1;
let prevColor = textureLoad( ${ parameters.prevOutputTarget }, indexUV );
let blendedColor = weightedAlphaBlend( prevColor, resultColor, 1.0 / f32( sampleCount ) );
textureStore( ${ parameters.sampleCountTarget }, indexUV, vec4( sampleCount ) );
textureStore( ${ parameters.outputTarget }, indexUV, blendedColor );

}
let sampleCount = textureLoad( ${ params.sampleCountTarget }, indexUV ).r + 1;
let prevColor = textureLoad( ${ params.prevOutputTarget }, indexUV );
let blendedColor = ${ weightedAlphaBlendFn }( prevColor, resultColor, 1.0 / f32( sampleCount ) );
textureStore( ${ params.sampleCountTarget }, indexUV, vec4( sampleCount ) );
textureStore( ${ params.outputTarget }, indexUV, blendedColor );

${ [
proxy( 'bvhData.value.storage.materials', parameters ),
proxy( 'bvhData.value.structs.material', parameters ),
proxy( 'bvhData.value.structs.transform', parameters ),
proxy( 'bvhData.value.fns.raycastFirstHit', parameters ),
proxy( 'bvhData.value.fns.sampleTrianglePoint', parameters ),
ndcToCameraRay, pcgRand2, pcgRand3, pcgInit, lambertBsdfFunc,
sampleEnvironmentFn, getSurfaceRecordFunc, weightedAlphaBlendFn,
] }`;
}`;

super( shader( parameters ) );
super( shader( params ) );

this.defineUniformAccessors( parameters );
this.defineUniformAccessors( params );

}

Expand Down
17 changes: 7 additions & 10 deletions src/webgpu/compute/ZeroOutBufferKernel.js
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import { IndirectStorageBufferAttribute } from 'three/webgpu';
import { ComputeKernel } from './ComputeKernel.js';
import { storage, wgslFn, globalId } from 'three/tsl';
import { storage, globalId } from 'three/tsl';
import { wgslTagFn } from '../lib/nodes/WGSLTagFnNode.js';

export class ZeroOutBufferKernel extends ComputeKernel {

Expand All @@ -15,19 +16,15 @@ export class ZeroOutBufferKernel extends ComputeKernel {
outputTarget: storage( new IndirectStorageBufferAttribute( 1, 1 ), type ),
};

const fn = wgslFn( /* wgsl */`
const fn = wgslTagFn/* wgsl */`
fn compute( globalId: vec3u ) -> void {

fn compute(
globalId: vec3u,
outputTarget: ptr<storage, array<u32>, read_write>,
) -> void {

outputTarget[ globalId.x ] = 0;
${ params.outputTarget }[ globalId.x ] = 0;

}
` )( params );
`;

super( fn );
super( fn( params ) );

this.defineUniformAccessors( {
target: params.outputTarget,
Expand Down
23 changes: 12 additions & 11 deletions src/webgpu/compute/wavefront/PrimeRayGenerationDispatchKernel.js
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
import { Vector2, Vector3 } from 'three';
import { IndirectStorageBufferAttribute } from 'three/webgpu';
import { wgslFn, uniform, storage } from 'three/tsl';
import { uniform, storage } from 'three/tsl';
import { ComputeKernel } from '../ComputeKernel.js';
import { queuedRayStruct } from './structs.js';
import { wgslTagFn } from '../../lib/nodes/WGSLTagFnNode.js';

export class PrimeRayGenerationDispatchKernel extends ComputeKernel {

Expand All @@ -18,24 +19,24 @@ export class PrimeRayGenerationDispatchKernel extends ComputeKernel {
rayQueue: storage( new IndirectStorageBufferAttribute( 1, queuedRayStruct.getLength() ), queuedRayStruct ).toReadOnly(),
rayQueueSize: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ),

outputTileIndex: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ),
outputDispatch: storage( new IndirectStorageBufferAttribute( 3, 1 ), 'u32' ),
outputTileIndex: storage( new IndirectStorageBufferAttribute( 2, 1 ), 'u32' ).setName( 'outputTileIndex' ),
outputDispatch: storage( new IndirectStorageBufferAttribute( 3, 1 ), 'u32' ).setName( 'outputDispatch' ),
};

const kernel = wgslFn( /* wgsl */`
const fn = wgslTagFn/* wgsl */`
fn compute(
rayWorkGroupSize: vec3u,

tileSize: vec2u,
tileCount: vec2u,
tileOffset: u32,
) -> void {

rayQueue: ptr<storage, array<QueuedRay>, read>,
rayQueueSize: ptr<storage, array<u32>, read_write>,
let rayQueue = &${ params.rayQueue };
let rayQueueSize = &${ params.rayQueueSize };

outputTileIndex: ptr<storage, array<u32>, read_write>,
outputDispatch: ptr<storage, array<u32>, read_write>,
) -> void {
let outputTileIndex = &${ params.outputTileIndex };
let outputDispatch = &${ params.outputDispatch };

// keep the queue index small
let queueCapacity = arrayLength( rayQueue );
Expand Down Expand Up @@ -79,9 +80,9 @@ export class PrimeRayGenerationDispatchKernel extends ComputeKernel {
}

}
`, [ queuedRayStruct ] )( params );
`;

super( kernel );
super( fn( params ) );

this.defineUniformAccessors( params );

Expand Down
74 changes: 31 additions & 43 deletions src/webgpu/compute/wavefront/ProcessHitsKernel.js
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,17 +1,17 @@
import { IndirectStorageBufferAttribute, StorageTexture, DataTexture } from 'three/webgpu';
import { ComputeKernel } from '../ComputeKernel.js';
import { uniform, storage, wgslFn, textureStore, globalId, texture, sampler } from 'three/tsl';
import { pcgRand3, pcgInit } from '../../nodes/random.wgsl.js';
import { uniform, storage, textureStore, globalId, texture, sampler } from 'three/tsl';
import { getSurfaceRecordFunc, lambertBsdfFunc } from '../../nodes/material.wgsl.js';
import { queuedRayStruct, queuedHitStruct } from './structs.js';
import { proxy } from '../../lib/nodes/NodeProxy.js';
import { proxy, proxyFn } from '../../lib/nodes/NodeProxy.js';
import { weightedAlphaBlendFn } from '../../nodes/sampling.wgsl.js';
import { wgslTagFn } from '../../lib/nodes/WGSLTagFnNode.js';

export class ProcessHitsKernel extends ComputeKernel {

constructor() {

const parameters = {
const params = {
bvhData: { value: null },

prevOutputTarget: textureStore( new StorageTexture( 1, 1 ) ).toReadOnly(),
Expand All @@ -35,32 +35,30 @@ export class ProcessHitsKernel extends ComputeKernel {
globalId: globalId,
};

const fn = wgslFn( /* wgsl */`
const sampleTrianglePointFn = proxyFn( 'bvhData.value.fns.sampleTrianglePoint', params );

fn compute(
// indices and target
prevOutputTarget: texture_storage_2d<rgba32float, read>,
outputTarget: texture_storage_2d<rgba32float, write>,
sampleCountTarget: texture_storage_2d<r32uint, read_write>,
const fn = wgslTagFn/* wgsl */`

fn compute(
// settings
smoothNormals: u32,
bounces: u32,

// rays
rayQueue: ptr<storage, array<QueuedRay>, read_write>,
rayQueueSize: ptr<storage, array<atomic<u32>>, read_write>,

// hits
hitQueue: ptr<storage, array<QueuedHit>, read_write>,
hitQueueSize: ptr<storage, array<u32>, read_write>,

textures: texture_2d_array<f32>,
textureSampler: sampler,

globalId: vec3u
) -> void {

let rayQueue = &${ params.rayQueue };
let rayQueueSize = &${ params.rayQueueSize };

let hitQueue = &${ params.hitQueue };
let hitQueueSize = &${ params.hitQueueSize };

let materials = &${ proxy( 'bvhData.value.storage.materials', params ) };
let transforms = &${ proxy( 'bvhData.value.storage.transforms', params ) };

// skip any rays invocations beyond the ray count
let hitQueueCapacity = arrayLength( hitQueue );
let hitIndex = ( globalId.x + hitQueueSize[ 0 ] );
Expand All @@ -77,29 +75,29 @@ export class ProcessHitsKernel extends ComputeKernel {

g_state.s0 = input.pcgStateS0;

let object = bvh_transforms.value[ input.objectIndex ];
var material = bvh_materials.value[ object.materialIndex ];
let object = transforms[ input.objectIndex ];
var material = materials[ object.materialIndex ];

// apply per-object colors
material.color *= object.color.rgb;
material.opacity *= object.color.a;

var vertexData = bvh_sampleTrianglePoint( input.barycoord, input.indices.xyz );
var vertexData = ${ sampleTrianglePointFn }( input.barycoord, input.indices.xyz );
vertexData.normal = normalize( transpose( object.inverseMatrixWorld ) * vertexData.normal );
vertexData.position = object.matrixWorld * vertexData.position;

let surface = getSurfaceRecord( material, vertexData, input.side, input.normal, textures, textureSampler );
let surface = ${ getSurfaceRecordFunc }( material, vertexData, input.side, input.normal, textures, textureSampler );

let scatterRec = bsdfSample( input.view, surface );
let scatterRec = ${ lambertBsdfFunc }( input.view, surface );

if ( input.currentBounce >= bounces ) {

// terminate ray, write color
let sampleCount = ( textureLoad( sampleCountTarget, indexUV ).r & ( ~ ACTIVE_FLAG ) ) + 1;
let prevColor = textureLoad( prevOutputTarget, indexUV );
let blendedColor = weightedAlphaBlend( prevColor, vec4f( 0, 0, 0, 1 ), 1.0 / f32( sampleCount ) );
textureStore( sampleCountTarget, indexUV, vec4( sampleCount ) );
textureStore( outputTarget, indexUV, blendedColor );
let sampleCount = ( textureLoad( ${ params.sampleCountTarget }, indexUV ).r & ( ~ ACTIVE_FLAG ) ) + 1;
let prevColor = textureLoad( ${ params.prevOutputTarget }, indexUV );
let blendedColor = ${ weightedAlphaBlendFn }( prevColor, vec4f( 0, 0, 0, 1 ), 1.0 / f32( sampleCount ) );
textureStore( ${ params.sampleCountTarget }, indexUV, vec4( sampleCount ) );
textureStore( ${ params.outputTarget }, indexUV, blendedColor );

} else {

Expand All @@ -114,21 +112,11 @@ export class ProcessHitsKernel extends ComputeKernel {

}

}
`, [
proxy( 'bvhData.value.structs.material', parameters ),
proxy( 'bvhData.value.structs.transform', parameters ),
proxy( 'bvhData.value.storage.materials', parameters ),
proxy( 'bvhData.value.storage.transforms', parameters ),
proxy( 'bvhData.value.fns.sampleTrianglePoint', parameters ),
queuedRayStruct, lambertBsdfFunc, getSurfaceRecordFunc,
pcgRand3, pcgInit, queuedHitStruct,
weightedAlphaBlendFn,
] );

super( fn( parameters ) );

this.defineUniformAccessors( parameters );
}`;

super( fn( params ) );

this.defineUniformAccessors( params );

}

Expand Down
Loading