Skip to content

Commit 167faff

Browse files
mvaligurskyMartin Valigursky
mvaligursky
and
Martin Valigursky
authored
Adds a GPU-based radix sort implementation using fragment shaders with mipmap-based prefix sums (playcanvas#8255)
* Adds a GPU-based radix sort implementation using fragment shaders with mipmap-based prefix sums * lint * test * reverted test --------- Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>
1 parent 9fc20f1 commit 167faff

18 files changed

+2116
-0
lines changed

‎examples/src/examples/test/radix-sort.controls.mjs‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
/**
2+
* @param {import('../../app/components/Example.mjs').ControlOptions} options - The options.
3+
* @returns {JSX.Element} The returned JSX Element.
4+
*/
5+
export const controls = ({ observer, ReactPCUI, jsx, fragment }) => {
6+
const { BindingTwoWay, LabelGroup, Panel, SliderInput } = ReactPCUI;
7+
8+
return fragment(
9+
jsx(
10+
Panel,
11+
{ headerText: 'Radix Sort' },
12+
jsx(
13+
LabelGroup,
14+
{ text: 'Elements (K)' },
15+
jsx(SliderInput, {
16+
binding: new BindingTwoWay(),
17+
link: { observer, path: 'options.elementsK' },
18+
value: 1000,
19+
min: 1,
20+
max: 10000,
21+
precision: 0
22+
})
23+
),
24+
jsx(
25+
LabelGroup,
26+
{ text: 'Bits' },
27+
jsx(SliderInput, {
28+
binding: new BindingTwoWay(),
29+
link: { observer, path: 'options.bits' },
30+
value: 16,
31+
min: 4,
32+
max: 24,
33+
precision: 0
34+
})
35+
)
36+
)
37+
);
38+
};

‎examples/src/examples/test/radix-sort.example.mjs‎

Lines changed: 334 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,334 @@
1+
// @config DESCRIPTION Test example for RenderPassRadixSort - GPU radix sort using mipmap binary search
2+
// @config HIDDEN
3+
import files from 'examples/files';
4+
import { data } from 'examples/observer';
5+
import { deviceType, rootPath } from 'examples/utils';
6+
import * as pc from 'playcanvas';
7+
8+
const canvas = /** @type {HTMLCanvasElement} */ (document.getElementById('application-canvas'));
9+
window.focus();
10+
11+
const gfxOptions = {
12+
deviceTypes: [deviceType],
13+
glslangUrl: `${rootPath}/static/lib/glslang/glslang.js`,
14+
twgslUrl: `${rootPath}/static/lib/twgsl/twgsl.js`
15+
};
16+
17+
const device = await pc.createGraphicsDevice(canvas, gfxOptions);
18+
device.maxPixelRatio = Math.min(window.devicePixelRatio, 2);
19+
20+
const createOptions = new pc.AppOptions();
21+
createOptions.graphicsDevice = device;
22+
23+
createOptions.componentSystems = [pc.RenderComponentSystem, pc.CameraComponentSystem, pc.LightComponentSystem];
24+
createOptions.resourceHandlers = [pc.TextureHandler];
25+
26+
const app = new pc.AppBase(canvas);
27+
app.init(createOptions);
28+
app.start();
29+
30+
app.setCanvasFillMode(pc.FILLMODE_FILL_WINDOW);
31+
app.setCanvasResolution(pc.RESOLUTION_AUTO);
32+
33+
const resize = () => app.resizeCanvas();
34+
window.addEventListener('resize', resize);
35+
app.on('destroy', () => {
36+
window.removeEventListener('resize', resize);
37+
});
38+
39+
// State
40+
/** @type {number} */
41+
let currentNumElements = 1000 * 1000; // 1M default
42+
/** @type {number} */
43+
let currentNumBits = 16;
44+
/** @type {pc.Texture|null} */
45+
let keysTexture = null;
46+
/** @type {pc.RenderPassRadixSort|null} */
47+
let radixSort = null;
48+
/** @type {number[]} */
49+
let originalValues = [];
50+
/** @type {boolean} */
51+
let needsRegen = true;
52+
53+
// Create render pass instance once
54+
radixSort = new pc.RenderPassRadixSort(device);
55+
56+
// ==================== MATERIALS ====================
57+
58+
// Create unsorted visualization material
59+
const unsortedMaterial = new pc.ShaderMaterial({
60+
uniqueName: 'UnsortedVizMaterial',
61+
vertexGLSL: files['vert.glsl'],
62+
fragmentGLSL: files['unsorted.glsl.frag'],
63+
vertexWGSL: files['vert.wgsl'],
64+
fragmentWGSL: files['unsorted.wgsl.frag'],
65+
attributes: {
66+
aPosition: pc.SEMANTIC_POSITION,
67+
aUv0: pc.SEMANTIC_TEXCOORD0
68+
}
69+
});
70+
71+
// Create sorted visualization material
72+
const sortedMaterial = new pc.ShaderMaterial({
73+
uniqueName: 'SortedVizMaterial',
74+
vertexGLSL: files['vert.glsl'],
75+
fragmentGLSL: files['sorted.glsl.frag'],
76+
vertexWGSL: files['vert.wgsl'],
77+
fragmentWGSL: files['sorted.wgsl.frag'],
78+
attributes: {
79+
aPosition: pc.SEMANTIC_POSITION,
80+
aUv0: pc.SEMANTIC_TEXCOORD0
81+
}
82+
});
83+
84+
// ==================== SCENE SETUP ====================
85+
86+
// Create camera entity
87+
const camera = new pc.Entity('camera');
88+
camera.addComponent('camera', {
89+
clearColor: new pc.Color(0.1, 0.1, 0.1),
90+
projection: pc.PROJECTION_ORTHOGRAPHIC,
91+
orthoHeight: 1
92+
});
93+
camera.setPosition(0, 0, 1);
94+
app.root.addChild(camera);
95+
96+
// Create unsorted visualization plane (top half)
97+
const unsortedPlane = new pc.Entity('unsortedPlane');
98+
unsortedPlane.addComponent('render', {
99+
type: 'plane',
100+
material: unsortedMaterial,
101+
castShadows: false,
102+
receiveShadows: false
103+
});
104+
unsortedPlane.setLocalPosition(0, 0.5, 0);
105+
unsortedPlane.setLocalScale(2, 1, 1);
106+
unsortedPlane.setEulerAngles(90, 0, 0);
107+
app.root.addChild(unsortedPlane);
108+
109+
// Create sorted visualization plane (bottom half)
110+
const sortedPlane = new pc.Entity('sortedPlane');
111+
sortedPlane.addComponent('render', {
112+
type: 'plane',
113+
material: sortedMaterial,
114+
castShadows: false,
115+
receiveShadows: false
116+
});
117+
sortedPlane.setLocalPosition(0, -0.5, 0);
118+
sortedPlane.setLocalScale(2, 1, 1);
119+
sortedPlane.setEulerAngles(90, 0, 0);
120+
app.root.addChild(sortedPlane);
121+
122+
// Create spinning cube for visual frame rate indicator
123+
const cube = new pc.Entity('cube');
124+
cube.addComponent('render', {
125+
type: 'box'
126+
});
127+
cube.setLocalPosition(0, 0, 0.3);
128+
cube.setLocalScale(0.15, 0.15, 0.15);
129+
app.root.addChild(cube);
130+
131+
// Create directional light for the cube
132+
const light = new pc.Entity('light');
133+
light.addComponent('light');
134+
light.setEulerAngles(45, 30, 0);
135+
app.root.addChild(light);
136+
137+
// ==================== HELPER FUNCTIONS ====================
138+
139+
/**
140+
* Calculates the optimal texture size for storing N elements.
141+
*
142+
* @param {number} numElements - Number of elements.
143+
* @returns {{width: number, height: number}} Texture dimensions.
144+
*/
145+
function calcTextureSize(numElements) {
146+
const pixels = Math.ceil(numElements);
147+
const size = Math.ceil(Math.sqrt(pixels));
148+
return { width: size, height: size };
149+
}
150+
151+
/**
152+
* Recreates the keys texture and generates random data.
153+
*/
154+
function regenerateData() {
155+
const numElements = currentNumElements;
156+
const numBits = currentNumBits;
157+
const maxValue = (1 << numBits) - 1;
158+
159+
// Calculate non-POT texture size
160+
const { width, height } = calcTextureSize(numElements);
161+
162+
// Destroy old texture
163+
if (keysTexture) {
164+
keysTexture.destroy();
165+
}
166+
167+
// Create new source keys texture
168+
keysTexture = new pc.Texture(device, {
169+
name: 'SourceKeys',
170+
width: width,
171+
height: height,
172+
format: pc.PIXELFORMAT_R32U,
173+
mipmaps: false,
174+
minFilter: pc.FILTER_NEAREST,
175+
magFilter: pc.FILTER_NEAREST,
176+
addressU: pc.ADDRESS_CLAMP_TO_EDGE,
177+
addressV: pc.ADDRESS_CLAMP_TO_EDGE
178+
});
179+
180+
// Generate random test data directly into texture (linear layout)
181+
const texData = keysTexture.lock();
182+
183+
// also keep original values for verification
184+
originalValues = [];
185+
186+
for (let i = 0; i < numElements; i++) {
187+
const value = Math.floor(Math.random() * maxValue);
188+
texData[i] = value;
189+
originalValues.push(value);
190+
}
191+
192+
// Fill remaining pixels with max value + 1 (will sort to end)
193+
for (let i = numElements; i < texData.length; i++) {
194+
texData[i] = maxValue + 1;
195+
}
196+
197+
keysTexture.unlock();
198+
199+
needsRegen = false;
200+
}
201+
202+
/**
203+
* Runs the GPU sort.
204+
*
205+
* @param {boolean} [verify] - Whether to verify results.
206+
*/
207+
function runSort(verify = false) {
208+
if (!keysTexture || !radixSort) return;
209+
210+
// Execute the GPU sort and get the sorted indices texture
211+
const sortedIndices = radixSort.sort(keysTexture, currentNumElements, currentNumBits);
212+
213+
// Update materials with the sorted texture
214+
updateMaterialParameters(sortedIndices);
215+
216+
// Verify results if requested
217+
if (verify) {
218+
verifyResults(sortedIndices);
219+
}
220+
}
221+
222+
/**
223+
* Updates material parameters after sort completes or data changes.
224+
*
225+
* @param {pc.Texture} sortedIndices - The sorted indices texture.
226+
*/
227+
function updateMaterialParameters(sortedIndices) {
228+
if (!keysTexture || !sortedIndices) {
229+
return;
230+
}
231+
232+
// Update unsorted material
233+
unsortedMaterial.setParameter('keysTexture', keysTexture);
234+
unsortedMaterial.setParameter('maxValue', (1 << currentNumBits) - 1);
235+
unsortedMaterial.setParameter('elementCount', currentNumElements);
236+
unsortedMaterial.setParameter('textureSize', [keysTexture.width, keysTexture.height]);
237+
unsortedMaterial.setParameter('debugMode', 0.0);
238+
unsortedMaterial.update();
239+
240+
// Update sorted material
241+
sortedMaterial.setParameter('sortedIndices', sortedIndices);
242+
sortedMaterial.setParameter('keysTexture', keysTexture);
243+
sortedMaterial.setParameter('elementCount', currentNumElements);
244+
sortedMaterial.setParameter('textureWidth', sortedIndices.width);
245+
sortedMaterial.setParameter('maxValue', (1 << currentNumBits) - 1);
246+
sortedMaterial.setParameter('sourceTextureSize', [keysTexture.width, keysTexture.height]);
247+
sortedMaterial.setParameter('debugMode', 0.0);
248+
sortedMaterial.update();
249+
}
250+
251+
/**
252+
* Downloads and verifies the sorted results. Only logs if validation fails.
253+
*
254+
* @param {pc.Texture} sortedIndices - The sorted indices texture to verify.
255+
*/
256+
async function verifyResults(sortedIndices) {
257+
if (!sortedIndices) {
258+
console.error('No sorted indices texture available');
259+
return;
260+
}
261+
262+
// Capture current state before async operation (use slice() for large arrays, not spread)
263+
const capturedOriginalValues = originalValues.slice();
264+
const capturedNumElements = currentNumElements;
265+
const width = sortedIndices.width;
266+
267+
// Read the sorted indices texture (R32U)
268+
const indicesResult = await sortedIndices.read(0, 0, width, width, {
269+
immediate: true
270+
});
271+
272+
// Extract sorted indices (stored in linear order)
273+
const sortedIndicesArray = [];
274+
for (let i = 0; i < capturedNumElements; i++) {
275+
sortedIndicesArray.push(indicesResult[i]);
276+
}
277+
278+
// Get sorted values by looking up original values (using captured copy)
279+
const sortedValues = sortedIndicesArray.map(idx => capturedOriginalValues[idx]);
280+
281+
// Verify sorting - only log on failure
282+
let errorCount = 0;
283+
for (let i = 1; i < sortedValues.length; i++) {
284+
if (sortedValues[i] < sortedValues[i - 1]) {
285+
if (errorCount < 5) {
286+
console.error(`Sort error at index ${i}: ${sortedValues[i - 1]} > ${sortedValues[i]}`);
287+
}
288+
errorCount++;
289+
}
290+
}
291+
292+
if (errorCount > 0) {
293+
console.error(`✗ [${device.deviceType}] Array is NOT correctly sorted (${errorCount} errors, ${(errorCount / capturedNumElements * 100).toFixed(2)}%)`);
294+
}
295+
}
296+
297+
// Initialize control values in the data observer
298+
data.set('options', {
299+
elementsK: 1000,
300+
bits: 16
301+
});
302+
303+
// Handle control changes from data binding
304+
data.on('*:set', (/** @type {string} */ path, /** @type {any} */ value) => {
305+
if (path === 'options.elementsK') {
306+
const newElements = value * 1000;
307+
if (newElements !== currentNumElements) {
308+
currentNumElements = newElements;
309+
needsRegen = true;
310+
}
311+
} else if (path === 'options.bits') {
312+
if (value !== currentNumBits) {
313+
currentNumBits = value;
314+
needsRegen = true;
315+
}
316+
}
317+
});
318+
319+
// Update loop - continuously sorts every frame
320+
app.on('update', (/** @type {number} */ dt) => {
321+
// Rotate the cube for visual frame rate indication
322+
cube.rotate(10 * dt, 20 * dt, 30 * dt);
323+
324+
// Regenerate data when parameters change
325+
const verify = needsRegen;
326+
if (needsRegen) {
327+
regenerateData();
328+
}
329+
330+
// Sort every frame, verify only after regeneration
331+
runSort(verify);
332+
});
333+
334+
export { app };

0 commit comments

Comments
 (0)