Fix issues

src/pipelines/diffusion/diffuse.wgsl

@@ -1,34 +1,28 @@
-struct VertexOutput {
-  @builtin(position) Position : vec4<f32>,
-  @location(0) fragUV : vec2<f32>,
-}
-
-@vertex
-fn vertex(@builtin(vertex_index) i : u32) -> VertexOutput {
-  var pos = array<vec2<f32>, 4>(
-    vec2(-1.0, 1.0),
-    vec2(-1.0, -1.0),
-    vec2(1.0, 1.0),
-    vec2(1.0, -1.0),
-  );
-
-  var output : VertexOutput;
-  output.Position = vec4<f32>(pos[i], 0.0, 1.0);
-  output.fragUV = output.Position.xy * 0.5 + 0.5;
-  return output;
-}
-
-
-@group(0) @binding(0) var mySampler: sampler;
-@group(0) @binding(1) var TargetTexture : texture_2d<f32>;
+struct Settings {
+  size : vec2<f32>,
+  diffusionRate : f32,
+  decayRate : f32,
+  deltaTime : f32,
+};
 
+@group(0) @binding(0) var<uniform> settings : Settings;
+@group(0) @binding(1) var Sampler: sampler;
+@group(0) @binding(2) var trailMap : texture_2d<f32>;
 
 @fragment
-fn fragment(@location(0) fragUV: vec2<f32>) -> @location(0) vec4<f32> {
-  // return vec4(1.0, 0.0, 0.0, 1.0);
-  return mix(
-    vec4(textureSample(TargetTexture, mySampler, fragUV).rgb, 0.1),
-    vec4(1.0, 1.0, 1.0, 1.0), 
-    0.01
+fn fragment(@location(0) uv: vec2<f32>) -> @location(0) vec4<f32> {
+  let current = textureSample(trailMap, Sampler, uv).rgb;
+
+  let neighbours: vec3<f32> = (
+      textureSample(trailMap, Sampler, uv + vec2<f32>(0, 1) / settings.size).rgb
+    + textureSample(trailMap, Sampler, uv + vec2<f32>(0, -1) / settings.size).rgb
+    + textureSample(trailMap, Sampler, uv + vec2<f32>(-1, 0) / settings.size).rgb
+    + textureSample(trailMap, Sampler, uv + vec2<f32>(1, 0) / settings.size).rgb
   );
+
+  return vec4(mix(
+    current,
+    neighbours / 4.0,
+    settings.diffusionRate
+  ) * (1.0 - settings.decayRate), 1.0);
 }

src/pipelines/diffusion/diffusion-pipeline.ts

@@ -1,19 +1,23 @@
+import { setUpFullScreenQuad } from '../../utils/full-screen-quad';
 import shader from './diffuse.wgsl';
 
 export class DiffusionPipeline {
+  private static readonly UNIFORM_COUNT = 6;
+
   private readonly pipeline: GPURenderPipeline;
+  private readonly uniforms: GPUBuffer;
+  private readonly quadVertexBuffer: GPUBuffer;
+
   private bindGroup?: GPUBindGroup;
   private previousTrailMapIn?: GPUTexture;
 
   public constructor(private readonly device: GPUDevice) {
+    const { buffer, vertex } = setUpFullScreenQuad(device);
+    this.quadVertexBuffer = buffer;
+
     this.pipeline = device.createRenderPipeline({
       layout: 'auto',
-      vertex: {
-        module: device.createShaderModule({
-          code: shader,
-        }),
-        entryPoint: 'vertex',
-      },
+      vertex,
       fragment: {
         module: device.createShaderModule({
           code: shader,
@@ -29,6 +33,31 @@ export class DiffusionPipeline {
         topology: 'triangle-strip',
       },
     });
+
+    this.uniforms = this.device.createBuffer({
+      size: DiffusionPipeline.UNIFORM_COUNT * Float32Array.BYTES_PER_ELEMENT,
+      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
+    });
+  }
+
+  public setParameters({
+    width,
+    height,
+    diffusionRate,
+    decayRate,
+    deltaTime,
+  }: {
+    width: number;
+    height: number;
+    diffusionRate: number;
+    decayRate: number;
+    deltaTime: number;
+  }) {
+    this.device.queue.writeBuffer(
+      this.uniforms,
+      0,
+      new Float32Array([width, height, diffusionRate, decayRate, deltaTime])
+    );
   }
 
   public execute(
@@ -49,11 +78,12 @@ export class DiffusionPipeline {
       ],
     };
 
-    const renderPassEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
-    renderPassEncoder.setBindGroup(0, this.bindGroup!);
-    renderPassEncoder.setPipeline(this.pipeline);
-    renderPassEncoder.draw(4, 1);
-    renderPassEncoder.end();
+    const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
+    passEncoder.setPipeline(this.pipeline);
+    passEncoder.setVertexBuffer(0, this.quadVertexBuffer);
+    passEncoder.setBindGroup(0, this.bindGroup);
+    passEncoder.draw(4, 1);
+    passEncoder.end();
   }
 
   private ensureBindGroupExists(trailMapIn: GPUTexture) {
@@ -63,13 +93,19 @@ export class DiffusionPipeline {
         entries: [
           {
             binding: 0,
+            resource: {
+              buffer: this.uniforms,
+            },
+          },
+          {
+            binding: 1,
             resource: this.device.createSampler({
               magFilter: 'linear',
               minFilter: 'linear',
             }),
           },
           {
-            binding: 1,
+            binding: 2,
             resource: trailMapIn.createView(),
           },
         ],