Skip to content

Commit 8548081

Browse files
bmschmidtbmschmidt
committed
webGPU pipelines
1 parent 830811c commit 8548081

File tree

5 files changed

+385
-307
lines changed

5 files changed

+385
-307
lines changed

‎package-lock.json‎

Lines changed: 4 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@
4545
},
4646
"homepage": "https://github.com/nomic-ai/deepscatter#readme",
4747
"peerDependencies": {
48-
"apache-arrow": ">=11.0.0"
48+
"apache-arrow": "^17.0.0"
4949
},
5050
"dependencies": {
5151
"d3-array": "^3.2.4",

‎src/webGPU/buffertools.ts‎

Lines changed: 113 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -1,127 +1,133 @@
11
import { isTypedArray, type TypedArray } from 'webgpu-utils';
22
import { BufferSet } from '../regl_rendering';
33
import { WebGPUBufferLocation } from '../types';
4-
// I track locations on buffers like this.
5-
// We keep track of both size -- the number of meaningful data bytes
6-
// and paddedSize -- the number of bytes including 256-byte padding.
7-
8-
export class WebGPUBufferSet extends BufferSet<GPUBuffer, WebGPUBufferLocation> {
9-
// Copied with alterations from deepscatter
4+
import { Some, TupleMap } from '../utilityFunctions';
105

11-
// An abstraction creating an expandable set of buffers that can be subdivided
12-
// to put more than one variable on the same
13-
// block of memory. Reusing buffers this way can have performance benefits over allocating
14-
// multiple different buffers for each small block used.
6+
// Unlike in webgl, we keep track of both size -- the number of meaningful data bytes
7+
// and paddedSize -- the number of bytes including 256-byte padding.
158

16-
// The general purpose here is to call 'allocate_block' that releases a block of memory
17-
// to use in creating a new array to be passed to regl.
9+
export class WebGPUBufferSet extends BufferSet<
10+
GPUBuffer,
11+
WebGPUBufferLocation
12+
> {
13+
public device: GPUDevice;
14+
private stagingBuffer: GPUBuffer;
15+
public usage: number;
1816

19-
public device: GPUDevice;
20-
private stagingBuffer: GPUBuffer;
21-
public usage: number;
17+
public store: TupleMap<string, WebGPUBufferLocation> = new TupleMap();
2218

23-
public store: Map<string, WebGPUBufferLocation> = new Map();
19+
/**
20+
*
21+
* @param regl the Regl context we're using.
22+
* @param buffer_size The number of bytes on each strip of memory that we'll ask for.
23+
*/
2424

25-
/**
26-
*
27-
* @param regl the Regl context we're using.
28-
* @param buffer_size The number of bytes on each strip of memory that we'll ask for.
29-
*/
25+
constructor(
26+
device: GPUDevice,
27+
buffer_size: number,
28+
usage: number = GPUBufferUsage.STORAGE |
29+
GPUBufferUsage.COPY_DST |
30+
GPUBufferUsage.COPY_SRC,
31+
) {
32+
super(buffer_size);
33+
this.device = device;
34+
// Track the ends in case we want to allocate smaller items.
35+
this.usage = usage;
36+
this.generate_new_buffer();
37+
this.stagingBuffer = device.createBuffer({
38+
size: buffer_size,
39+
usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE,
40+
mappedAtCreation: false, // saves a little trouble in the passThrough function
41+
});
42+
}
3043

31-
constructor(
32-
device: GPUDevice,
33-
buffer_size: number,
34-
usage: number = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC
35-
) {
36-
super(buffer_size)
37-
this.device = device;
38-
// Track the ends in case we want to allocate smaller items.
39-
this.usage = usage;
40-
this.generate_new_buffer();
41-
this.stagingBuffer = device.createBuffer({
42-
size: buffer_size,
43-
usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE,
44-
mappedAtCreation: false // saves a little trouble in the passThrough function
45-
});
46-
}
44+
private async passThroughStagingBuffer(
45+
values: Uint32Array,
46+
bufferLocation: WebGPUBufferLocation,
47+
) {
48+
// WebGPU
49+
const { buffer, offset, paddedSize } = bufferLocation;
50+
while (this.stagingBuffer.mapState !== 'unmapped') {
51+
// Wait in line for a millisecond.
52+
// Would be better to hold a queue and apply more than one of these at once.
53+
await new Promise((resolve) => setTimeout(resolve, 1));
54+
}
55+
await this.stagingBuffer.mapAsync(GPUMapMode.WRITE, 0, paddedSize);
56+
new Uint32Array(
57+
this.stagingBuffer.getMappedRange(0, values.byteLength),
58+
).set(values);
59+
this.stagingBuffer.unmap();
60+
const commandEncoder = this.device.createCommandEncoder();
61+
commandEncoder.copyBufferToBuffer(
62+
this.stagingBuffer,
63+
0,
64+
buffer,
65+
offset,
66+
paddedSize,
67+
);
68+
this.device.queue.submit([commandEncoder.finish()]);
69+
}
4770

48-
private async passThroughStagingBuffer(values: Uint32Array, bufferLocation: WebGPUBufferLocation) {
49-
// WebGPU
50-
const { buffer, offset, paddedSize } = bufferLocation;
51-
while (this.stagingBuffer.mapState !== 'unmapped') {
52-
// Wait in line for a millisecond.
53-
// Would be better to hold a queue and apply more than one of these at once.
54-
await new Promise((resolve) => setTimeout(resolve, 1));
55-
}
56-
await this.stagingBuffer.mapAsync(GPUMapMode.WRITE, 0, paddedSize);
57-
new Uint32Array(this.stagingBuffer.getMappedRange(0, values.byteLength)).set(values);
58-
this.stagingBuffer.unmap();
59-
const commandEncoder = this.device.createCommandEncoder();
60-
commandEncoder.copyBufferToBuffer(this.stagingBuffer, 0, buffer, offset, paddedSize);
61-
this.device.queue.submit([commandEncoder.finish()]);
62-
}
71+
register(k: Some<string>, v: WebGPUBufferLocation) {
72+
this.store.set(k, v);
73+
}
6374

64-
register(k: string, v: WebGPUBufferLocation) {
65-
this.store.set(k, v);
66-
}
75+
async set(key: Some<string>, value: TypedArray) {
76+
if (this.store.has(key)) {
77+
throw new Error(`Key ${key.join(', ')} already exists in buffer set.`);
78+
}
79+
const size = value.byteLength;
80+
const paddedSize = Math.ceil(size / 256) * 256;
6781

68-
async set(key: string, value: TypedArray) {
69-
if (this.store.has(key)) {
70-
throw new Error(`Key ${key} already exists in buffer set.`);
71-
}
72-
const size = value.byteLength;
73-
const paddedSize = Math.ceil(size / 256) * 256;
82+
const { buffer, offset } = this.allocate_block(paddedSize);
7483

75-
const { buffer, offset } = this.allocate_block(paddedSize);
84+
// If it's a typed array, we can just copy it directly.
85+
// cast it to uint32array
86+
const v2 = value;
87+
const data = new Uint32Array(v2.buffer, v2.byteOffset, v2.byteLength / 4);
88+
const description = { buffer, offset, size, paddedSize };
89+
await this.passThroughStagingBuffer(data, description);
90+
this.register(key, description);
91+
}
7692

77-
// If it's a typed array, we can just copy it directly.
78-
// cast it to uint32array
79-
const v2 = value;
80-
const data = new Uint32Array(v2.buffer, v2.byteOffset, v2.byteLength / 4);
81-
const description = { buffer, offset, size, paddedSize };
82-
await this.passThroughStagingBuffer(data, description);
83-
this.register(key, description);
84-
}
93+
_create_buffer(): GPUBuffer {
94+
return this.device.createBuffer({
95+
size: this.buffer_size,
96+
usage: this.usage,
97+
mappedAtCreation: false,
98+
});
99+
}
85100

86-
_create_buffer() : GPUBuffer {
87-
return this.device.createBuffer({
88-
size: this.buffer_size,
89-
usage: this.usage,
90-
mappedAtCreation: false
91-
})
92-
}
93-
94-
_create_leftover_buffer() : WebGPUBufferLocation {
95-
return {
96-
buffer: this.buffers[0],
97-
offset: this.pointer,
98-
stride: 4, // meaningless here.
99-
byte_size: this.buffer_size - this.pointer,
100-
paddedSize: this.buffer_size - this.pointer
101-
}
102-
}
101+
_create_leftover_buffer(): WebGPUBufferLocation {
102+
return {
103+
buffer: this.buffers[0],
104+
offset: this.pointer,
105+
stride: 4, // meaningless here.
106+
byte_size: this.buffer_size - this.pointer,
107+
paddedSize: this.buffer_size - this.pointer,
108+
};
109+
}
103110
}
104111

105-
106112
export function createSingletonBuffer(
107-
device: GPUDevice,
108-
data: Uint32Array | Int32Array | Float32Array | ArrayBuffer,
109-
usage: number
113+
device: GPUDevice,
114+
data: Uint32Array | Int32Array | Float32Array | ArrayBuffer,
115+
usage: number,
110116
): GPUBuffer {
111-
// Creates a disposable singleton buffer.
112-
// ReadonlyBufferSet ought to provide better performance; but
113-
// this allows more different buffer sizes and easier destruction.
114-
const buffer = device.createBuffer({
115-
size: data.byteLength,
116-
usage,
117-
mappedAtCreation: true
118-
});
119-
const mappedRange = buffer.getMappedRange();
120-
if (isTypedArray(data)) {
121-
new Uint32Array(mappedRange).set(data as TypedArray);
122-
} else {
123-
new Uint32Array(mappedRange).set(new Uint32Array(data as ArrayBuffer));
124-
}
125-
buffer.unmap();
126-
return buffer;
117+
// Creates a disposable singleton buffer.
118+
// ReadonlyBufferSet ought to provide better performance; but
119+
// this allows more different buffer sizes and easier destruction.
120+
const buffer = device.createBuffer({
121+
size: data.byteLength,
122+
usage,
123+
mappedAtCreation: true,
124+
});
125+
const mappedRange = buffer.getMappedRange();
126+
if (isTypedArray(data)) {
127+
new Uint32Array(mappedRange).set(data as TypedArray);
128+
} else {
129+
new Uint32Array(mappedRange).set(new Uint32Array(data as ArrayBuffer));
130+
}
131+
buffer.unmap();
132+
return buffer;
127133
}

0 commit comments

Comments
 (0)