WgpuSplatSorter.ts

src/wgpu/WgpuSplatSorter.ts · raw · all examples · site home

Exports: WgpuSplatSorter

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
import type { WebGPURenderer } from "three/webgpu";
import { SortKeys, type SortKeyBits } from "../core/sort";
import { GpuArrayReader } from "../gpu/GpuArrayReader";
import type { GpuArrays } from "../gpu/GpuArrays";
import { GpuSortKeys } from "../gpu/sortKey.tsl";
import { TSL, type TslNode } from "../gpu/tsl";
import { WgpuBuffer } from "./WgpuBuffer";
import { WgpuCompact } from "./WgpuCompact";
import { WgpuCompute, requireWebGPU } from "./WgpuCompute";
import { WgpuRadixSort } from "./WgpuRadixSort";

const { uint } = TSL;

// The draw order of projected splats, sorted on the GPU in the same frame
// (WebGPU backend only): compaction drops the invalid keys of the
// projection's `sortKeys` bundle, a radix sort orders the rest back to front, and
// a one-thread kernel writes the drawIndexedIndirect arguments. Nothing is
// read back.
export class WgpuSplatSorter {
  // Splats per instance of the indirect draw: SplatRenderer's quads
  static readonly QUADS_PER_INSTANCE = 1024;

  readonly renderer: WebGPURenderer;
  // Sorted splat indices: values.element(i) is the i-th splat drawn
  readonly values: WgpuBuffer<"uint32">;
  // The number of them (a TSL node, the sort's visible count)
  readonly count: TslNode;
  // drawIndexedIndirect arguments: the indices of QUADS_PER_INSTANCE
  // quads, enough instances for the visible count, first index, base
  // vertex and first instance 0
  readonly indirect: WgpuBuffer<"uint32">;

  private keys = new GpuArrayReader(SortKeys.KEYS);
  private compact: WgpuCompact;
  private radix: WgpuRadixSort;
  private args: WgpuCompute;

  constructor(renderer: WebGPURenderer) {
    requireWebGPU(renderer, "WgpuSplatSorter");
    this.renderer = renderer;
    this.compact = new WgpuCompact({ renderer, key: (i) => GpuSortKeys.unpack(this.keys.load("key", i)) });
    this.radix = new WgpuRadixSort({ renderer, keys: this.compact.keys, values: this.compact.values });
    this.values = this.compact.values;

    const indirect = this.indirect = new WgpuBuffer({ renderer, format: "uint32", count: 5, indirect: true, label: "radix_splatDrawArgs" });
    const count = this.count = this.compact.count.element(uint(0));
    const quads = WgpuSplatSorter.QUADS_PER_INSTANCE;
    this.args = new WgpuCompute({
      label: "radix_splatDrawArgs",
      workgroupSize: [1],
      body: () => {
        indirect.element(uint(0)).assign(uint(quads * 6));
        indirect.element(uint(1)).assign(count.add(uint(quads - 1)).div(uint(quads)));
        indirect.element(uint(2)).assign(uint(0));
        indirect.element(uint(3)).assign(uint(0));
        indirect.element(uint(4)).assign(uint(0));
      },
    });
  }

  // Sort the first `count` splats of `projected` (arrays with a sortKeys
  // bundle of SortKeys.KEYS, as a SplatRenderer's projected arrays) by their keys of
  // width `bits`
  sort(projected: GpuArrays, count: number, bits: SortKeyBits) {
    this.keys.bind(projected, "sortKeys");
    this.compact.invalid.value = SortKeys.invalid(bits);

    this.compact.compact(count);
    this.radix.sortIndirect(this.compact.count, { bits });
    this.args.dispatch(this.renderer, [1]);
  }

  // The visible count of the latest sort (a readback, for tests and stats)
  async readCount(): Promise<number> {
    return (await this.compact.count.read(0, 1))[0];
  }

  dispose() {
    this.compact.dispose();
    this.radix.dispose();
    this.args.dispose();
    this.indirect.dispose();
  }
}