ParticleSort.ts

examples/mpm/ParticleSort.ts · raw · open the example · all examples · site home

Exports: MpmParticleSort

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
// Reorders the MPM particle buffers by the Morton code of each particle's
// grid cell, so that neighbouring threads work on neighbouring particles:
// P2G's atomics and G2P's gathers then hit the same cache lines, and a
// workgroup's particles span a small box of cells (what the tiled P2G needs).
// The particles drift apart as the water moves; MpmSim re-sorts every
// `sortEvery` steps.
//
// One sort: a kernel writes (Morton key, index) per particle, WgpuRadixSort
// sorts the pairs, two gather kernels write the particles in sorted order to
// scratch buffers (two, to stay within 8 storage buffers per stage), and the
// scratch buffers are copied back.
import type { WebGPURenderer } from "three/webgpu";
import { TSL, WgpuBuffer, WgpuCompute, WgpuRadixSort, type SortKeyBits, type TslNode } from "@sparkjsdev/radix";
import type { MpmSim } from "./MpmSim.js";

const { If, clamp, floor, ivec3, uint } = TSL;

// The low 10 bits of v with two zero bits between each (3D Morton spreading)
function spreadBits(v: TslNode): TslNode {
  const x = v.bitAnd(uint(0x3ff)).toVar();
  x.assign(x.bitOr(x.shiftLeft(uint(16))).bitAnd(uint(0x030000ff)));
  x.assign(x.bitOr(x.shiftLeft(uint(8))).bitAnd(uint(0x0300f00f)));
  x.assign(x.bitOr(x.shiftLeft(uint(4))).bitAnd(uint(0x030c30c3)));
  x.assign(x.bitOr(x.shiftLeft(uint(2))).bitAnd(uint(0x09249249)));
  return x;
}

export class MpmParticleSort {
  sim: MpmSim;
  renderer: WebGPURenderer;
  // Key bits the radix sort looks at (3 x the bits of the largest grid dimension)
  readonly bits: SortKeyBits;
  private keys: WgpuBuffer<"uint32">;
  private values: WgpuBuffer<"uint32">;
  private radix: WgpuRadixSort;
  // Sorted copies of pos, vel, c0, c1, c2
  private scratch: WgpuBuffer<"float32x4">[];
  private keyKernel: WgpuCompute;
  private gatherKernels: WgpuCompute[];

  constructor(sim: MpmSim) {
    this.sim = sim;
    const { renderer, res } = sim;
    this.renderer = renderer;
    const keyBits = 3 * Math.ceil(Math.log2(Math.max(res.x, res.y, res.z)));
    if (keyBits > 30) {
      throw new Error(`MpmParticleSort: grid ${res.x}x${res.y}x${res.z} is too large for 10-bit Morton axes`);
    }
    this.bits = keyBits <= 16 ? 16 : keyBits <= 24 ? 24 : 32;
    this.keys = new WgpuBuffer({ renderer, format: "uint32", count: sim.count });
    this.values = new WgpuBuffer({ renderer, format: "uint32", count: sim.count });
    this.radix = new WgpuRadixSort({ renderer, keys: this.keys, values: this.values });
    this.scratch = this.sources().map(() => new WgpuBuffer({ renderer, format: "float32x4", count: sim.count }));
    this.keyKernel = this.makeKeys();
    const [pos, vel, c0, c1, c2] = this.sources();
    const [sPos, sVel, sC0, sC1, sC2] = this.scratch;
    this.gatherKernels = [this.makeGather([pos, vel, c0], [sPos, sVel, sC0]), this.makeGather([c1, c2], [sC1, sC2])];
  }

  private sources() {
    const { pos, vel, c0, c1, c2 } = this.sim;
    return [pos, vel, c0, c1, c2];
  }

  private makeKeys() {
    const { sim, keys, values } = this;
    const { res } = sim;
    return new WgpuCompute({
      label: "mpmSortKeys",
      body: ({ index }) => {
        If(index.lessThan(sim.uniforms.count), () => {
          const cell = clamp(ivec3(floor(sim.pos.element(index).xyz)), ivec3(0), ivec3(res.x - 1, res.y - 1, res.z - 1)).toVar();
          const key = spreadBits(uint(cell.x)).shiftLeft(uint(2))
            .bitOr(spreadBits(uint(cell.y)).shiftLeft(uint(1)))
            .bitOr(spreadBits(uint(cell.z)));
          keys.element(index).assign(key);
          values.element(index).assign(index);
        });
      },
    });
  }

  // dst[k][i] = src[k][values[i]]
  private makeGather(src: WgpuBuffer<"float32x4">[], dst: WgpuBuffer<"float32x4">[]) {
    const { sim, values } = this;
    return new WgpuCompute({
      label: "mpmSortGather",
      body: ({ index }) => {
        If(index.lessThan(sim.uniforms.count), () => {
          const from = values.element(index).toVar();
          src.forEach((buffer, k) => {
            dst[k].element(index).assign(buffer.element(from));
          });
        });
      },
    });
  }

  sort() {
    const { renderer, sim } = this;
    const count = sim.count;
    if (count <= 0) {
      return;
    }
    this.keys.ensureCapacity(count);
    this.values.ensureCapacity(count);
    for (const buffer of this.scratch) {
      buffer.ensureCapacity(count);
    }
    this.keyKernel.dispatch(renderer, count);
    this.radix.sort(count, { bits: this.bits });
    for (const kernel of this.gatherKernels) {
      kernel.dispatch(renderer, count);
    }
    // Same queue as the dispatches above, so ordered after them
    this.sources().forEach((buffer, k) => buffer.copyFrom(this.scratch[k], { count }));
  }

  dispose() {
    this.radix.dispose();
    for (const buffer of [this.keys, this.values, ...this.scratch]) {
      buffer.dispose();
    }
  }
}