rayzee 7.5.0 → 7.5.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "rayzee",
-  "version": "7.5.0",
+  "version": "7.5.1",
   "type": "module",
   "description": "Real-time WebGPU path tracing engine built on Three.js",
   "main": "dist/rayzee.umd.js",

package/src/EngineDefaults.js

@@ -88,6 +88,10 @@ export const ENGINE_DEFAULTS = {
 	performanceModeAdaptive: 'medium',
 
 	fireflyThreshold: 3.0,
+	// Wavefront material-coherence sort: global counting-sort of entering rays by material before
+	// Shade (material-pure workgroups), under dynamic dispatch. Measured −8% at 1024²/8b. Gated on
+	// material count > 8; the histogram bin count is sized per-scene to the material count.
+	wavefrontSortMaterials: true,
 	renderLimitMode: 'frames',
 	renderTimeLimit: 30,
 	renderMode: 0,

package/src/Processor/KernelManager.js

@@ -12,8 +12,7 @@
 const WORKGROUP_SIZES = {
 	generate: [ 16, 16, 1 ], // 2D screen-space
 	extend: [ 256, 1, 1 ], // 1D ray-parallel
-	sort: [ 256, 1, 1 ], // 1D ray-parallel
-	shade: [ 256, 1, 1 ], // 1D ray-parallel (sorted)
+	shade: [ 256, 1, 1 ], // 1D ray-parallel (material-sorted when sort enabled)
 	connect: [ 256, 1, 1 ], // 1D shadow-ray-parallel
 	accumulate: [ 256, 1, 1 ], // 1D shadow-ray-parallel
 	compact: [ 256, 1, 1 ], // 1D ray-parallel

package/src/Processor/QueueManager.js

@@ -40,6 +40,9 @@ export class QueueManager {
 		// A/B alternate: one read by current bounce, other written by compaction
 		this.activeIndices = null;
 		this.activeIndicesRO = null;
+		this.sortedIndices = null;
+		this.sortedIndicesRO = null;
+		this.sortGlobalHistogram = null;
 		this.pingPong = 0; // 0 = read A / write B, 1 = read B / write A
 
 		if ( maxRays > 0 ) {
@@ -83,11 +86,23 @@ export class QueueManager {
 			b: storage( attrB, 'uint' ).toReadOnly(),
 		};
 
+		// Material-sort output: a material-reordered permutation of the active index
+		// list, written by the global material sort and read by Shade in place of activeIndices.
+		const sortAttr = new StorageInstancedBufferAttribute( new Uint32Array( capacity ), 1 );
+		this._sortAttr = sortAttr;
+		this.sortedIndices = storage( sortAttr, 'uint' );
+		this.sortedIndicesRO = storage( sortAttr, 'uint' ).toReadOnly();
+
+		// Global material-sort histogram, sized to the bin cap (SORT_GLOBAL_MAX_BINS=256); kernels use
+		// only the first `bins` entries (= per-scene material count). 256 × 4B = 1KB.
+		this._sortGlobalHistAttr = new StorageInstancedBufferAttribute( new Uint32Array( 256 ), 1 );
+		this.sortGlobalHistogram = storage( this._sortGlobalHistAttr, 'uint' ).toAtomic();
+
 		this.pingPong = 0;
 
 		const totalBytes = (
 			COUNTER.COUNT * 4 +
-			capacity * 4 * 2
+			capacity * 4 * 3
 		);
 
 		console.log(
@@ -131,6 +146,24 @@ export class QueueManager {
 
 	}
 
+	getSortedRW() {
+
+		return this.sortedIndices;
+
+	}
+
+	getSortedRO() {
+
+		return this.sortedIndicesRO;
+
+	}
+
+	getSortGlobalHistogram() {
+
+		return this.sortGlobalHistogram;
+
+	}
+
 	// raw attribute for `renderer.getArrayBufferAsync(...)` readback
 	getCountersAttribute() {
 
@@ -167,6 +200,9 @@ export class QueueManager {
 		this.counters = null;
 		this.activeIndices = null;
 		this.activeIndicesRO = null;
+		this.sortedIndices = null;
+		this.sortedIndicesRO = null;
+		this.sortGlobalHistogram = null;
 		this.capacity = 0;
 
 	}

package/src/Stages/PathTracer.js

@@ -17,6 +17,11 @@ import { buildShadeKernel, SHADE_WG_SIZE } from '../TSL/ShadeKernel.js';
 import { buildCompactKernel, buildCompactSubgroupKernel, COMPACT_WG_SIZE } from '../TSL/CompactKernel.js';
 import { buildFinalWriteKernel, FINALWRITE_WG_SIZE } from '../TSL/FinalWriteKernel.js';
 import { buildDebugKernel, DEBUG_WG_SIZE } from '../TSL/DebugKernel.js';
+import {
+	buildResetGlobalHistKernel, buildGlobalHistKernel, buildGlobalPrefixKernel, buildGlobalScatterKernel,
+	SORT_GLOBAL_WG_SIZE, SORT_GLOBAL_MAX_BINS,
+} from '../TSL/SortGlobalKernels.js';
+import { ENGINE_DEFAULTS } from '../EngineDefaults.js';
 import {
 	Fn, uint, atomicStore, atomicLoad, instanceIndex, If, Return,
 } from 'three/tsl';
@@ -43,6 +48,11 @@ export class PathTracer extends PathTracerStage {
 		// CPU sizes per-bounce kernels from last frame's survivor curve; kernels bound on ENTERING_COUNT so over-sizing is safe. (indirect dispatch not viable — three.js doesn't sync compute-written indirect buffers across submissions)
 		this._useDynamicDispatch = true;
 
+		// Global material-coherence sort: set per-build from ENGINE_DEFAULTS + material count (>8).
+		// Reorders entering rays into material-pure workgroups before Shade; runs under dynamic dispatch
+		// (compact reads the unsorted active list, so the survivor set is unchanged). Measured −8% at 1024²/8b.
+		this._sortMaterials = false;
+
 		// Flag-gated off: perf-neutral vs atomic-append and adds a 'subgroups' feature dependency.
 		this._useSubgroupCompact = false;
 
@@ -90,7 +100,7 @@ export class PathTracer extends PathTracerStage {
 			if ( ! qm ) return null;
 			return [
 				qm._countersAttr, qm._bounceCountsAttr,
-				qm._attrA, qm._attrB,
+				qm._attrA, qm._attrB, qm._sortAttr,
 			];
 
 		} );
@@ -250,6 +260,7 @@ export class PathTracer extends PathTracerStage {
 			this._wfCurrentBounce.value = bounce;
 
 			// Functional-compaction path (dynamic dispatch): copyback keeps the read buffer dense, kernels sized to live survivors. Dynamic-off uses the full path (ENTERING=maxRays, identity buffer).
+			// Material sort is compatible: shade reads sortedIndices while compact still reads the UNSORTED active list (getActiveReadRO), so the survivor set is unchanged.
 			const useFunctionalCompaction = this._useDynamicDispatch;
 			if ( useFunctionalCompaction ) {
 
@@ -281,6 +292,8 @@ export class PathTracer extends PathTracerStage {
 				const wg = [ Math.ceil( sized / 256 ), 1, 1 ];
 				km.setDispatchCount( 'extend', wg );
 				km.setDispatchCount( 'shade', wg );
+				km.setDispatchCount( 'globalHist', wg );
+				km.setDispatchCount( 'globalScatter', wg );
 				km.setDispatchCount( 'compact', wg );
 				km.setDispatchCount( 'compactCopyback', wg );
 
@@ -291,11 +304,23 @@ export class PathTracer extends PathTracerStage {
 				km.setDispatchCount( 'extend', full );
 				km.setDispatchCount( 'shade', full );
 				km.setDispatchCount( 'compact', full );
+				km.setDispatchCount( 'globalHist', full );
+				km.setDispatchCount( 'globalScatter', full );
 
 			}
 
 			// Extend/Shade kept separate (not fused): a fused kernel's register pressure drops occupancy more than fusion saves.
 			km.dispatch( 'extend' );
+			if ( this._sortMaterials ) {
+
+				// Global material counting sort (material-pure workgroups): reset, histogram, prefix-sum, scatter.
+				km.dispatch( 'resetGlobalHist' );
+				km.dispatch( 'globalHist' );
+				km.dispatch( 'globalPrefix' );
+				km.dispatch( 'globalScatter' );
+
+			}
+
 			km.dispatch( 'shade' );
 
 			km.dispatch( 'resetActiveCounter' );
@@ -660,6 +685,10 @@ export class PathTracer extends PathTracerStage {
 			displacementMaps: freshDisplacementMaps,
 		};
 
+		// Material-coherence sort gate (experiment): only worthwhile above a few materials.
+		this._sortMaterials = ( ENGINE_DEFAULTS.wavefrontSortMaterials ?? false )
+			&& ( this.materialData?.materialCount ?? 0 ) > 8;
+
 		const extFn = buildExtendKernel( {
 			bvhBuffer: freshBvh,
 			triangleBuffer: freshTri,
@@ -677,6 +706,50 @@ export class PathTracer extends PathTracerStage {
 			)
 		);
 
+		// Material-coherence sort: reorder the entering-ray indices by material between
+		// Extend and Shade. Histogram is workgroup-shared (patches.js §4); Shade reads the output.
+		if ( this._sortMaterials ) {
+
+			const sgHist = qm.getSortGlobalHistogram();
+			const sortBins = Math.min( SORT_GLOBAL_MAX_BINS, this.materialData?.materialCount ?? SORT_GLOBAL_MAX_BINS );
+			this._kernelManager.register( 'resetGlobalHist',
+				buildResetGlobalHistKernel( { sortGlobalHistogram: sgHist, bins: sortBins } )().compute(
+					[ 1, 1, 1 ], [ SORT_GLOBAL_WG_SIZE, 1, 1 ]
+				)
+			);
+			this._kernelManager.register( 'globalHist',
+				buildGlobalHistKernel( {
+					hitBufferRO: pb.hitBuffer.ro,
+					activeIndicesReadRO: qm.getActiveReadRO(),
+					sortGlobalHistogram: sgHist,
+					counters,
+					bins: sortBins,
+				} )().compute(
+					[ Math.ceil( maxRays / SORT_GLOBAL_WG_SIZE ), 1, 1 ],
+					[ SORT_GLOBAL_WG_SIZE, 1, 1 ]
+				)
+			);
+			this._kernelManager.register( 'globalPrefix',
+				buildGlobalPrefixKernel( { sortGlobalHistogram: sgHist, bins: sortBins } )().compute(
+					[ 1, 1, 1 ], [ 1, 1, 1 ]
+				)
+			);
+			this._kernelManager.register( 'globalScatter',
+				buildGlobalScatterKernel( {
+					hitBufferRO: pb.hitBuffer.ro,
+					activeIndicesReadRO: qm.getActiveReadRO(),
+					sortedIndicesRW: qm.getSortedRW(),
+					sortGlobalHistogram: sgHist,
+					counters,
+					bins: sortBins,
+				} )().compute(
+					[ Math.ceil( maxRays / SORT_GLOBAL_WG_SIZE ), 1, 1 ],
+					[ SORT_GLOBAL_WG_SIZE, 1, 1 ]
+				)
+			);
+
+		}
+
 		const shadeFn = buildShadeKernel( {
 			gBufferRW,
 			envCompensationDelta: this.envCompensationDelta,
@@ -689,7 +762,7 @@ export class PathTracer extends PathTracerStage {
 			rngBufferRW: pb.rngBuffer.rw,
 			hitBufferRO: pb.hitBuffer.ro,
 			counters,
-			activeIndicesRO: qm.getActiveReadRO(),
+			activeIndicesRO: this._sortMaterials ? qm.getSortedRO() : qm.getActiveReadRO(),
 			albedoMaps: freshAlbedoMaps,
 			normalMaps: freshNormalMaps,
 			bumpMaps: freshBumpMaps,

package/src/TSL/SortGlobalKernels.js

@@ -0,0 +1,135 @@
+/**
+ * SortGlobalKernels.js — GLOBAL material counting sort of the entering rays.
+ *
+ * Unlike a per-workgroup sort (locally coherent only), this orders rays by material
+ * ACROSS ALL workgroups, so each material's rays occupy one contiguous region of sortedIndices.
+ * For materials with many rays that means whole material-PURE workgroups/subgroups → real shading
+ * coherence. MEASURED −8% vs no-sort at 1024²/8b (per-WG sort was neutral).
+ *
+ * Four passes (counting sort): reset → histogram → exclusive prefix sum → scatter.
+ *   - histogram uses a per-workgroup local histogram (workgroup-shared atomics, patches.js §4)
+ *     flushed to the global histogram once per bin per workgroup, cutting global-atomic contention
+ *     from O(rays) to O(workgroups × bins).
+ *   - the prefix-summed histogram doubles as the per-bin running cursor in scatter.
+ *
+ * `bins` is chosen per-scene by the caller (= material count, capped). Materials ≥ bins clamp into
+ * the last bin (graceful coherence loss). Capped at SORT_GLOBAL_MAX_BINS so a single 256-thread
+ * workgroup can zero/flush exactly one bin per thread (no per-thread loop needed).
+ */
+
+import {
+	Fn, uint,
+	If, Loop,
+	instanceIndex, localId,
+	workgroupBarrier,
+	atomicAdd, atomicLoad, atomicStore,
+} from 'three/tsl';
+
+import { workgroupAtomicArray } from './patches.js';
+import { readHitMaterialIndex } from '../Processor/PackedRayBuffer.js';
+import { COUNTER } from '../Processor/QueueManager.js';
+
+const WG_SIZE = 256;
+export const SORT_GLOBAL_MAX_BINS = 256;
+
+// Pass 1 — zero the global histogram.
+export function buildResetGlobalHistKernel( { sortGlobalHistogram, bins = SORT_GLOBAL_MAX_BINS } ) {
+
+	return Fn( () => {
+
+		If( instanceIndex.lessThan( uint( bins ) ), () => {
+
+			atomicStore( sortGlobalHistogram.element( instanceIndex ), uint( 0 ) );
+
+		} );
+
+	} );
+
+}
+
+// Pass 2 — histogram: per-workgroup local count (on-chip), flushed to global once per bin/WG.
+export function buildGlobalHistKernel( { hitBufferRO, activeIndicesReadRO, sortGlobalHistogram, counters, bins = SORT_GLOBAL_MAX_BINS } ) {
+
+	return Fn( () => {
+
+		const tid = instanceIndex;
+		const lid = localId.x;
+		const local = workgroupAtomicArray( 'uint', bins );
+
+		If( lid.lessThan( uint( bins ) ), () => {
+
+			atomicStore( local.element( lid ), uint( 0 ) );
+
+		} );
+		workgroupBarrier();
+
+		const entering = atomicLoad( counters.element( uint( COUNTER.ENTERING_COUNT ) ) );
+		If( tid.lessThan( entering ), () => {
+
+			const rayID = activeIndicesReadRO.element( tid );
+			const bin = uint( readHitMaterialIndex( hitBufferRO, rayID ) ).clamp( uint( 0 ), uint( bins - 1 ) );
+			atomicAdd( local.element( bin ), uint( 1 ) );
+
+		} );
+		workgroupBarrier();
+
+		If( lid.lessThan( uint( bins ) ), () => {
+
+			const c = atomicLoad( local.element( lid ) );
+			If( c.greaterThan( uint( 0 ) ), () => {
+
+				atomicAdd( sortGlobalHistogram.element( lid ), c );
+
+			} );
+
+		} );
+
+	} );
+
+}
+
+// Pass 3 — exclusive prefix sum over global bins (single thread; dispatch [1,1,1]).
+export function buildGlobalPrefixKernel( { sortGlobalHistogram, bins = SORT_GLOBAL_MAX_BINS } ) {
+
+	return Fn( () => {
+
+		If( instanceIndex.equal( uint( 0 ) ), () => {
+
+			const running = uint( 0 ).toVar();
+			Loop( { start: uint( 0 ), end: uint( bins ), type: 'uint' }, ( { i } ) => {
+
+				const slot = sortGlobalHistogram.element( i );
+				const c = atomicLoad( slot );
+				atomicStore( slot, running );
+				running.addAssign( c );
+
+			} );
+
+		} );
+
+	} );
+
+}
+
+// Pass 4 — scatter: atomicAdd on the prefix-summed histogram gives each ray its global slot
+// within its material's contiguous region.
+export function buildGlobalScatterKernel( { hitBufferRO, activeIndicesReadRO, sortedIndicesRW, sortGlobalHistogram, counters, bins = SORT_GLOBAL_MAX_BINS } ) {
+
+	return Fn( () => {
+
+		const tid = instanceIndex;
+		const entering = atomicLoad( counters.element( uint( COUNTER.ENTERING_COUNT ) ) );
+		If( tid.lessThan( entering ), () => {
+
+			const rayID = activeIndicesReadRO.element( tid );
+			const bin = uint( readHitMaterialIndex( hitBufferRO, rayID ) ).clamp( uint( 0 ), uint( bins - 1 ) );
+			const pos = atomicAdd( sortGlobalHistogram.element( bin ), uint( 1 ) );
+			sortedIndicesRW.element( pos ).assign( rayID );
+
+		} );
+
+	} );
+
+}
+
+export { WG_SIZE as SORT_GLOBAL_WG_SIZE };

package/src/TSL/patches.js

@@ -61,6 +61,10 @@ WebGPUBackend.prototype.createNodeBuilder = function ( object, renderer ) {
 		configurable: true,
 	} );
 
+	// Install the workgroup-atomic-array codegen patch (section 4) lazily off the
+	// first builder, so we don't need to import WGSLNodeBuilder directly.
+	_installScopedArrayAtomicPatch( builder );
+
 	return builder;
 
 };
@@ -220,3 +224,89 @@ export function struct( members, name = null ) {
 	return wrappedFactory;
 
 }
+
+// ---------------------------------------------------------------------------
+// 4. Workgroup-scoped atomic arrays — var<workgroup> array<atomic<T>, N>
+// ---------------------------------------------------------------------------
+// TSL's `workgroupArray(type, count)` emits `var<workgroup> name: array<T, N>`
+// (`WGSLNodeBuilder.getScopedArrays`, r185 ~line 1773) with a NON-atomic element
+// type, so `atomicAdd(arr.element(i), v)` fails WGSL validation. AtomicFunctionNode
+// is itself address-space-agnostic — it emits `atomicAdd(&name[i], v)`, valid for
+// BOTH storage and workgroup pointers — so the ONLY missing piece is the
+// declaration. This patch adds `workgroupAtomicArray()` (a `workgroupArray` tagged
+// atomic) and overrides `getScopedArrays` to wrap the element type in `atomic<…>`
+// for tagged arrays, enabling true on-chip workgroup-shared atomics (e.g. fast
+// per-workgroup histograms) instead of the slow global-storage-atomic fallback.
+// Access tagged arrays ONLY via atomic ops (atomicAdd/atomicLoad/atomicStore).
+
+import { workgroupArray } from 'three/tsl';
+
+const _WorkgroupInfoNode = workgroupArray( 'uint', 1 ).constructor;
+
+// Tag the node so the patched getScopedArrays emits an atomic element type. The
+// array's WGSL name is only known after generate() runs, so we mark the builder's
+// scopedArrays entry there (idempotent across analyze/generate passes).
+const _origWorkgroupGenerate = _WorkgroupInfoNode.prototype.generate;
+
+_WorkgroupInfoNode.prototype.generate = function ( builder ) {
+
+	const name = _origWorkgroupGenerate.call( this, builder );
+	if ( this.isAtomicArray === true ) {
+
+		const entry = builder.scopedArrays && builder.scopedArrays.get( name );
+		if ( entry ) entry.isAtomic = true;
+
+	}
+
+	return name;
+
+};
+
+/**
+ * Like `workgroupArray(type, count)` but declares the workgroup buffer with an
+ * atomic element type: `var<workgroup> name: array<atomic<type>, count>`.
+ * Elements MUST be accessed only via atomic ops (atomicAdd/atomicLoad/atomicStore).
+ *
+ * @param {string} type - Element type (e.g. 'uint').
+ * @param {number} count - Number of elements.
+ * @returns {WorkgroupInfoNode} The tagged workgroup array node.
+ */
+export function workgroupAtomicArray( type, count ) {
+
+	const node = workgroupArray( type, count );
+	node.isAtomicArray = true;
+	return node;
+
+}
+
+let _scopedArraysPatched = false;
+
+// Override getScopedArrays on the builder's prototype (compute-stage WGSL
+// assembly) to emit `atomic<T>` element types for tagged workgroup arrays.
+// Installed once, lazily, off the first node builder created.
+function _installScopedArrayAtomicPatch( builder ) {
+
+	if ( _scopedArraysPatched ) return;
+	const proto = Object.getPrototypeOf( builder );
+	if ( ! proto || typeof proto.getScopedArrays !== 'function' ) return;
+
+	proto.getScopedArrays = function ( shaderStage ) {
+
+		if ( shaderStage !== 'compute' ) return;
+
+		const snippets = [];
+		for ( const { name, scope, bufferType, bufferCount, isAtomic } of this.scopedArrays.values() ) {
+
+			const type = this.getType( bufferType );
+			const elementType = ( isAtomic === true ) ? `atomic< ${ type } >` : type;
+			snippets.push( `var<${ scope }> ${ name }: array< ${ elementType }, ${ bufferCount } >;` );
+
+		}
+
+		return snippets.join( '\n' );
+
+	};
+
+	_scopedArraysPatched = true;
+
+}