diff --git a/apps/typegpu-docs/src/examples/simulation/slime-mold-3d/index.ts b/apps/typegpu-docs/src/examples/simulation/slime-mold-3d/index.ts index 2a6e2cc00..bbb2b0e86 100644 --- a/apps/typegpu-docs/src/examples/simulation/slime-mold-3d/index.ts +++ b/apps/typegpu-docs/src/examples/simulation/slime-mold-3d/index.ts @@ -6,9 +6,7 @@ import { randf } from '@typegpu/noise'; import * as m from 'wgpu-matrix'; const root = await tgpu.init({ - device: { - optionalFeatures: ['float32-filterable'], - }, + device: { optionalFeatures: ['float32-filterable'] }, }); const canFilter = root.enabledFeatures.has('float32-filterable'); const device = root.device; @@ -32,8 +30,8 @@ const CAMERA_FOV_DEGREES = 60; const CAMERA_DISTANCE_MULTIPLIER = 1.5; const CAMERA_INITIAL_ANGLE = Math.PI / 4; -const RAYMARCH_STEPS = 128; -const DENSITY_MULTIPLIER = 0.05; +const RAYMARCH_STEPS = 48; +const DENSITY_MULTIPLIER = 0.1; const RANDOM_DIRECTION_WEIGHT = 0.3; const CENTER_BIAS_WEIGHT = 0.7; @@ -105,15 +103,19 @@ const Params = d.struct({ evaporationRate: d.f32, }); -const agentsData = root.createMutable(d.arrayOf(Agent, NUM_AGENTS)); +const agentsDataBuffers = [0, 1].map(() => + root.createBuffer(d.arrayOf(Agent, NUM_AGENTS)).$usage('storage') +); +const mutableAgentsDataBuffers = agentsDataBuffers.map((b) => b.as('mutable')); root['~unstable'].createGuardedComputePipeline((x) => { 'use gpu'; randf.seed(x / NUM_AGENTS); const pos = randf.inUnitSphere().mul(resolution.x / 4).add(resolution.div(2)); const center = resolution.div(2); const dir = std.normalize(center.sub(pos)); - agentsData.$[x] = Agent({ position: pos, direction: dir }); + mutableAgentsDataBuffers[0].$[x] = Agent({ position: pos, direction: dir }); + mutableAgentsDataBuffers[1].$[x] = Agent({ position: pos, direction: dir }); }).dispatchThreads(NUM_AGENTS); const params = root.createUniform(Params, { @@ -136,8 +138,16 @@ const textures = [0, 1].map(() => ); const computeLayout = tgpu.bindGroupLayout({ + oldAgents: { storage: d.arrayOf(Agent), access: 'readonly' }, oldState: { storageTexture: d.textureStorage3d('r32float', 'read-only') }, + newAgents: { storage: d.arrayOf(Agent), access: 'mutable' }, + newState: { storageTexture: d.textureStorage3d('r32float', 'write-only') }, +}); + +const blurLayout = tgpu.bindGroupLayout({ + oldState: { texture: d.texture3d() }, newState: { storageTexture: d.textureStorage3d('r32float', 'write-only') }, + sampler: { sampler: 'filtering' }, }); const renderLayout = tgpu.bindGroupLayout({ @@ -223,7 +233,7 @@ const updateAgents = tgpu['~unstable'].computeFn({ const dims = std.textureDimensions(computeLayout.$.oldState); const dimsf = d.vec3f(dims); - const agent = agentsData.$[gid.x]; + const agent = computeLayout.$.oldAgents[gid.x]; const random = randf.sample(); let direction = std.normalize(agent.direction); @@ -292,7 +302,7 @@ const updateAgents = tgpu['~unstable'].computeFn({ ); } - agentsData.$[gid.x] = Agent({ + computeLayout.$.newAgents[gid.x] = Agent({ position: newPos, direction, }); @@ -306,52 +316,48 @@ const updateAgents = tgpu['~unstable'].computeFn({ ); }); +const sampler = root['~unstable'].createSampler({ + magFilter: canFilter ? 'linear' : 'nearest', + minFilter: canFilter ? 'linear' : 'nearest', +}); + const blur = tgpu['~unstable'].computeFn({ in: { gid: d.builtin.globalInvocationId }, workgroupSize: BLUR_WORKGROUP_SIZE, })(({ gid }) => { - const dims = std.textureDimensions(computeLayout.$.oldState); + const dims = std.textureDimensions(blurLayout.$.oldState); if (gid.x >= dims.x || gid.y >= dims.y || gid.z >= dims.z) return; + const uv = d.vec3f(gid).add(0.5).div(d.vec3f(dims)); + const offset = d.vec3f(1).div(d.vec3f(dims)); + let sum = d.f32(); - let count = d.f32(); - - for (let offsetZ = -1; offsetZ <= 1; offsetZ++) { - for (let offsetY = -1; offsetY <= 1; offsetY++) { - for (let offsetX = -1; offsetX <= 1; offsetX++) { - const samplePos = d.vec3i(gid.xyz).add( - d.vec3i(offsetX, offsetY, offsetZ), - ); - const dimsi = d.vec3i(dims); - - if ( - samplePos.x >= 0 && samplePos.x < dimsi.x && - samplePos.y >= 0 && samplePos.y < dimsi.y && - samplePos.z >= 0 && samplePos.z < dimsi.z - ) { - const value = - std.textureLoad(computeLayout.$.oldState, d.vec3u(samplePos)).x; - sum = sum + value; - count = count + 1; - } - } + + for (let axis = 0; axis < 3; axis++) { + for (let sign = -1; sign <= 1; sign += 2) { + const offsetVec = d.vec3f( + std.select(0, offset.x * sign, axis === 0), + std.select(0, offset.y * sign, axis === 1), + std.select(0, offset.z * sign, axis === 2), + ); + sum = sum + std.textureSampleLevel( + blurLayout.$.oldState, + blurLayout.$.sampler, + uv.add(offsetVec), + 0, + ).x; } } - const blurred = sum / count; + const blurred = sum / 6.0; const newValue = std.saturate(blurred - params.$.evaporationRate); std.textureStore( - computeLayout.$.newState, + blurLayout.$.newState, gid.xyz, d.vec4f(newValue, 0, 0, 1), ); }); -const sampler = root['~unstable'].createSampler({ - magFilter: canFilter ? 'linear' : 'nearest', - minFilter: canFilter ? 'linear' : 'nearest', -}); - // Ray-box intersection const rayBoxIntersection = ( rayOrigin: d.v3f, @@ -375,6 +381,7 @@ const fragmentShader = tgpu['~unstable'].fragmentFn({ in: { uv: d.vec2f }, out: d.vec4f, })(({ uv }) => { + randf.seed2(uv); const ndc = d.vec2f(uv.x * 2 - 1, 1 - uv.y * 2); const ndcNear = d.vec4f(ndc, -1, 1); const ndcFar = d.vec4f(ndc, 1, 1); @@ -393,11 +400,23 @@ const fragmentShader = tgpu['~unstable'].fragmentFn({ return d.vec4f(); } - // March params - const tStart = std.max(isect.tNear, 0); + const jitter = randf.sample() * 20; + const tStart = std.max(isect.tNear + jitter, jitter); const tEnd = isect.tFar; - const numSteps = RAYMARCH_STEPS; - const stepSize = (tEnd - tStart) / numSteps; + + const intersectionLength = tEnd - tStart; + const baseStepsPerUnit = d.f32(0.3); + const minSteps = d.i32(8); + const maxSteps = d.i32(RAYMARCH_STEPS); + + const adaptiveSteps = std.clamp( + d.i32(intersectionLength * baseStepsPerUnit), + minSteps, + maxSteps, + ); + + const numSteps = adaptiveSteps; + const stepSize = intersectionLength / d.f32(numSteps); const thresholdLo = d.f32(0.06); const thresholdHi = d.f32(0.25); @@ -411,11 +430,8 @@ const fragmentShader = tgpu['~unstable'].fragmentFn({ const TMin = d.f32(1e-3); - for (let i = 0; i < numSteps; i++) { - if (transmittance <= TMin) { - break; - } - + let i = d.i32(0); + while (i < numSteps && transmittance > TMin) { const t = tStart + (d.f32(i) + 0.5) * stepSize; const pos = rayOrigin.add(rayDir.mul(t)); const texCoord = pos.div(resolution); @@ -433,6 +449,8 @@ const fragmentShader = tgpu['~unstable'].fragmentFn({ accum = accum.add(contrib.mul(transmittance)); transmittance = transmittance * (1 - alphaSrc); + + i += 1; } const alpha = 1 - transmittance; @@ -454,8 +472,18 @@ const blurPipeline = root['~unstable'] const bindGroups = [0, 1].map((i) => root.createBindGroup(computeLayout, { + oldAgents: agentsDataBuffers[i], + oldState: textures[i], + newAgents: agentsDataBuffers[1 - i], + newState: textures[1 - i], + }) +); + +const blurBindGroups = [0, 1].map((i) => + root.createBindGroup(blurLayout, { oldState: textures[i], newState: textures[1 - i], + sampler: sampler, }) ); @@ -476,7 +504,7 @@ function frame() { params.writePartial({ deltaTime }); blurPipeline - .with(bindGroups[currentTexture]) + .with(blurBindGroups[currentTexture]) .dispatchWorkgroups( Math.ceil(resolution.x / BLUR_WORKGROUP_SIZE[0]), Math.ceil(resolution.y / BLUR_WORKGROUP_SIZE[1]), diff --git a/packages/typegpu/tests/examples/individual/slime-mold-3d.test.ts b/packages/typegpu/tests/examples/individual/slime-mold-3d.test.ts index b01863174..ca622f4f0 100644 --- a/packages/typegpu/tests/examples/individual/slime-mold-3d.test.ts +++ b/packages/typegpu/tests/examples/individual/slime-mold-3d.test.ts @@ -59,30 +59,35 @@ describe('slime mold 3d example', () => { direction: vec3f, } - @group(0) @binding(1) var agentsData_10: array; + @group(0) @binding(1) var item_10: array; + + @group(0) @binding(2) var item_12: array; fn wrappedCallback_2(x: u32, _arg_1: u32, _arg_2: u32) { randSeed_3((f32(x) / 8e+5f)); var pos = ((randInUnitSphere_6() * 64.) + vec3f(128)); var center = vec3f(128); var dir = normalize((center - pos)); - agentsData_10[x] = Agent_11(pos, dir); + item_10[x] = Agent_11(pos, dir); + item_12[x] = Agent_11(pos, dir); } - struct mainCompute_Input_12 { + struct mainCompute_Input_13 { @builtin(global_invocation_id) id: vec3u, } - @compute @workgroup_size(256, 1, 1) fn mainCompute_0(in: mainCompute_Input_12) { + @compute @workgroup_size(256, 1, 1) fn mainCompute_0(in: mainCompute_Input_13) { if (any(in.id >= sizeUniform_1)) { return; } wrappedCallback_2(in.id.x, in.id.y, in.id.z); } - @group(1) @binding(0) var oldState_1: texture_storage_3d; + @group(1) @binding(0) var oldState_1: texture_3d; + + @group(1) @binding(2) var sampler_2: sampler; - struct Params_3 { + struct Params_4 { deltaTime: f32, moveSpeed: f32, sensorAngle: f32, @@ -91,37 +96,31 @@ describe('slime mold 3d example', () => { evaporationRate: f32, } - @group(0) @binding(0) var params_2: Params_3; + @group(0) @binding(0) var params_3: Params_4; - @group(1) @binding(1) var newState_4: texture_storage_3d; + @group(1) @binding(1) var newState_5: texture_storage_3d; - struct blur_Input_5 { + struct blur_Input_6 { @builtin(global_invocation_id) gid: vec3u, } - @compute @workgroup_size(4, 4, 4) fn blur_0(_arg_0: blur_Input_5) { + @compute @workgroup_size(4, 4, 4) fn blur_0(_arg_0: blur_Input_6) { var dims = textureDimensions(oldState_1); if ((((_arg_0.gid.x >= dims.x) || (_arg_0.gid.y >= dims.y)) || (_arg_0.gid.z >= dims.z))) { return; } + var uv = ((vec3f(_arg_0.gid) + 0.5) / vec3f(dims)); + var offset = (vec3f(1) / vec3f(dims)); var sum = 0f; - var count = 0f; - for (var offsetZ = -1; (offsetZ <= 1i); offsetZ++) { - for (var offsetY = -1; (offsetY <= 1i); offsetY++) { - for (var offsetX = -1; (offsetX <= 1i); offsetX++) { - var samplePos = (vec3i(_arg_0.gid.xyz) + vec3i(offsetX, offsetY, offsetZ)); - var dimsi = vec3i(dims); - if (((((((samplePos.x >= 0i) && (samplePos.x < dimsi.x)) && (samplePos.y >= 0i)) && (samplePos.y < dimsi.y)) && (samplePos.z >= 0i)) && (samplePos.z < dimsi.z))) { - var value = textureLoad(oldState_1, vec3u(samplePos)).x; - sum = (sum + value); - count = (count + 1f); - } - } + for (var axis = 0; (axis < 3i); axis++) { + for (var sign = -1; (sign <= 1i); sign += 2i) { + var offsetVec = vec3f(select(0f, (offset.x * f32(sign)), (axis == 0i)), select(0f, (offset.y * f32(sign)), (axis == 1i)), select(0f, (offset.z * f32(sign)), (axis == 2i))); + sum = (sum + textureSampleLevel(oldState_1, sampler_2, (uv + offsetVec), 0).x); } } - var blurred = (sum / count); - var newValue = saturate((blurred - params_2.evaporationRate)); - textureStore(newState_4, _arg_0.gid.xyz, vec4f(newValue, 0f, 0f, 1f)); + var blurred = (sum / 6f); + var newValue = saturate((blurred - params_3.evaporationRate)); + textureStore(newState_5, _arg_0.gid.xyz, vec4f(newValue, 0f, 0f, 1f)); } var seed_3: vec2f; @@ -134,14 +133,14 @@ describe('slime mold 3d example', () => { seed_2(seed); } - @group(1) @binding(0) var oldState_4: texture_storage_3d; + @group(1) @binding(1) var oldState_4: texture_storage_3d; struct Agent_6 { position: vec3f, direction: vec3f, } - @group(0) @binding(0) var agentsData_5: array; + @group(1) @binding(0) var oldAgents_5: array; fn item_8() -> f32 { var a = dot(seed_3, vec2f(23.140779495239258, 232.6168975830078)); @@ -180,7 +179,7 @@ describe('slime mold 3d example', () => { evaporationRate: f32, } - @group(0) @binding(1) var params_11: Params_12; + @group(0) @binding(0) var params_11: Params_12; struct SenseResult_13 { weightedDir: vec3f, @@ -231,20 +230,22 @@ describe('slime mold 3d example', () => { return (sign(alignment) * value); } - @group(1) @binding(1) var newState_18: texture_storage_3d; + @group(1) @binding(2) var newAgents_18: array; - struct updateAgents_Input_19 { + @group(1) @binding(3) var newState_19: texture_storage_3d; + + struct updateAgents_Input_20 { @builtin(global_invocation_id) gid: vec3u, } - @compute @workgroup_size(64) fn updateAgents_0(_arg_0: updateAgents_Input_19) { + @compute @workgroup_size(64) fn updateAgents_0(_arg_0: updateAgents_Input_20) { if ((_arg_0.gid.x >= 800000u)) { return; } randSeed_1(((f32(_arg_0.gid.x) / 8e+5f) + 0.1f)); var dims = textureDimensions(oldState_4); var dimsf = vec3f(dims); - var agent = agentsData_5[_arg_0.gid.x]; + var agent = oldAgents_5[_arg_0.gid.x]; var random = randFloat01_7(); var direction = normalize(agent.direction); var senseResult = sense3D_9(agent.position, direction); @@ -289,10 +290,10 @@ describe('slime mold 3d example', () => { var toCenter = normalize((center - newPos)); direction = normalize(((randomDir * 0.3) + (toCenter * 0.7))); } - agentsData_5[_arg_0.gid.x] = Agent_6(newPos, direction); + newAgents_18[_arg_0.gid.x] = Agent_6(newPos, direction); var oldState = textureLoad(oldState_4, vec3u(newPos)).x; var newState = (oldState + 1f); - textureStore(newState_18, vec3u(newPos), vec4f(newState, 0f, 0f, 1f)); + textureStore(newState_19, vec3u(newPos), vec4f(newState, 0f, 0f, 1f)); } struct fullScreenTriangle_Input_1 { @@ -311,21 +312,31 @@ describe('slime mold 3d example', () => { return fullScreenTriangle_Output_2(vec4f(pos[in.vertexIndex], 0, 1), uv[in.vertexIndex]); } - struct Camera_5 { + var seed_6: vec2f; + + fn seed2_5(value: vec2f) { + seed_6 = value; + } + + fn randSeed2_4(seed: vec2f) { + seed2_5(seed); + } + + struct Camera_8 { viewProj: mat4x4f, invViewProj: mat4x4f, position: vec3f, } - @group(0) @binding(0) var cameraData_4: Camera_5; + @group(0) @binding(0) var cameraData_7: Camera_8; - struct RayBoxResult_7 { + struct RayBoxResult_10 { tNear: f32, tFar: f32, hit: bool, } - fn rayBoxIntersection_6(rayOrigin: vec3f, rayDir: vec3f, boxMin: vec3f, boxMax: vec3f) -> RayBoxResult_7 { + fn rayBoxIntersection_9(rayOrigin: vec3f, rayDir: vec3f, boxMin: vec3f, boxMax: vec3f) -> RayBoxResult_10 { var invDir = (vec3f(1) / rayDir); var t0 = ((boxMin - rayOrigin) * invDir); var t1 = ((boxMax - rayOrigin) * invDir); @@ -334,58 +345,76 @@ describe('slime mold 3d example', () => { var tNear = max(max(tmin.x, tmin.y), tmin.z); var tFar = min(min(tmax.x, tmax.y), tmax.z); var hit = ((tFar >= tNear) && (tFar >= 0f)); - return RayBoxResult_7(tNear, tFar, hit); + return RayBoxResult_10(tNear, tFar, hit); + } + + fn item_12() -> f32 { + var a = dot(seed_6, vec2f(23.140779495239258, 232.6168975830078)); + var b = dot(seed_6, vec2f(54.47856521606445, 345.8415222167969)); + seed_6.x = fract((cos(a) * 136.8168f)); + seed_6.y = fract((cos(b) * 534.7645f)); + return seed_6.y; } - @group(1) @binding(0) var state_8: texture_3d; + fn randFloat01_11() -> f32 { + return item_12(); + } + + @group(1) @binding(0) var state_13: texture_3d; - @group(0) @binding(1) var sampler_9: sampler; + @group(0) @binding(1) var sampler_14: sampler; - struct fragmentShader_Input_10 { + struct fragmentShader_Input_15 { @location(0) uv: vec2f, } - @fragment fn fragmentShader_3(_arg_0: fragmentShader_Input_10) -> @location(0) vec4f { + @fragment fn fragmentShader_3(_arg_0: fragmentShader_Input_15) -> @location(0) vec4f { + randSeed2_4(_arg_0.uv); var ndc = vec2f(((_arg_0.uv.x * 2f) - 1f), (1f - (_arg_0.uv.y * 2f))); var ndcNear = vec4f(ndc, -1, 1f); var ndcFar = vec4f(ndc, 1f, 1f); - var worldNear = (cameraData_4.invViewProj * ndcNear); - var worldFar = (cameraData_4.invViewProj * ndcFar); + var worldNear = (cameraData_7.invViewProj * ndcNear); + var worldFar = (cameraData_7.invViewProj * ndcFar); var rayOrigin = (worldNear.xyz / worldNear.w); var rayEnd = (worldFar.xyz / worldFar.w); var rayDir = normalize((rayEnd - rayOrigin)); var boxMin = vec3f(); var boxMax = vec3f(256); - var isect = rayBoxIntersection_6(rayOrigin, rayDir, boxMin, boxMax); + var isect = rayBoxIntersection_9(rayOrigin, rayDir, boxMin, boxMax); if (!isect.hit) { return vec4f(); } - var tStart = max(isect.tNear, 0f); + var jitter = (randFloat01_11() * 20f); + var tStart = max((isect.tNear + jitter), jitter); var tEnd = isect.tFar; - var numSteps = 128; - var stepSize = ((tEnd - tStart) / f32(numSteps)); + var intersectionLength = (tEnd - tStart); + var baseStepsPerUnit = 0.30000001192092896f; + var minSteps = 8i; + var maxSteps = 48i; + var adaptiveSteps = clamp(i32((intersectionLength * baseStepsPerUnit)), minSteps, maxSteps); + var numSteps = adaptiveSteps; + var stepSize = (intersectionLength / f32(numSteps)); var thresholdLo = 0.05999999865889549f; var thresholdHi = 0.25f; var gamma = 1.399999976158142f; - var sigmaT = 0.05000000074505806f; + var sigmaT = 0.10000000149011612f; var albedo = vec3f(0.5699999928474426, 0.4399999976158142, 0.9599999785423279); var transmittance = 1f; var accum = vec3f(); var TMin = 0.0010000000474974513f; - for (var i = 0; (i < numSteps); i++) { - if ((transmittance <= TMin)) { - break; - } + var i = 0i; + while (((i < numSteps) && (transmittance > TMin))) { var t = (tStart + ((f32(i) + 0.5f) * stepSize)); var pos = (rayOrigin + (rayDir * t)); var texCoord = (pos / vec3f(256)); - var sampleValue = textureSampleLevel(state_8, sampler_9, texCoord, 0).x; + var sampleValue = textureSampleLevel(state_13, sampler_14, texCoord, 0).x; var d0 = smoothstep(thresholdLo, thresholdHi, sampleValue); var density = pow(d0, gamma); var alphaSrc = (1f - exp(((-sigmaT * density) * stepSize))); var contrib = (albedo * alphaSrc); accum = (accum + (contrib * transmittance)); transmittance = (transmittance * (1f - alphaSrc)); + i += 1i; } var alpha = (1f - transmittance); return vec4f(accum, alpha);