ThreeJS WebGPU & Neighbor Search ( 近傍探索 )
前回の WebGPU に続いて、ThreeJS WebGPU で、Neighbor Search を
行ってみました。
（粒子の表示は、前回同様です。）
プログラム
NeighborSearch-ThreeJS-WebGPU.html
<!DOCTYPE html>
<html>
<head>
<title>Neighbor Search ThreeJS</title>
<meta charset="utf-8">
</head>
<body>
<script type="module">
import * as THREE from '../build/three.module.js';
import WebGPURenderer from './jsm/renderers/webgpu/WebGPURenderer.js';
import WebGPU from './jsm/renderers/webgpu/WebGPU.js';
import {
Matrix3Uniform,
Matrix4Uniform,
Vector2Uniform
} from './jsm/renderers/webgpu/WebGPUUniform.js';
import WebGPUStorageBuffer from './jsm/renderers/webgpu/WebGPUStorageBuffer.js';
import WebGPUUniformsGroup from './jsm/renderers/webgpu/WebGPUUniformsGroup.js';
let camera, scene, renderer;
let pointer;
let stepPQUniform;
const computeParams_initialize_buckets = [];
const computeParams_sort_buckets = [];
const computeParams_work = [];
const computeParams_bucketRef = [];
const computeParams_findNeighbors = [];
const computeParams_clearNeighbors = [];
const computeParams_integrate = [];
const powerN = 4;
const particleNum = 2 ** powerN
init().then( animate ).catch( error );
async function init() {
if ( WebGPU.isAvailable() === false ) {
document.body.appendChild( WebGPU.getErrorMessage() );
throw 'No WebGPU support';
}
camera = new THREE.PerspectiveCamera( 70, window.innerWidth / window.innerHeight, 0.1, 1000 );
camera.position.z = 1;
camera.position.x = 0.5;
camera.position.y = 0.5;
// particle moving area
// x, y: 0 ~ 1
scene = new THREE.Scene();
scene.background = new THREE.Color( 0x000000 );
const viewRadius = 0.2;
const bucketSize = 2.0 * viewRadius;
const bucketNum = Math.ceil(1.0 / bucketSize);
const bucketIndexNum = bucketNum * bucketNum; // 2d
console.log("powerN: ", powerN);
const particleSize = 3;
// find neighbors of particle with this particleIndex
const particleIndex = 0;
const particleArray = new Float32Array( particleNum * particleSize );
const velocityArray = new Float32Array( particleNum * particleSize );
const colorArray = new Float32Array( particleNum * 3 );
const bucketsArray = new Uint32Array( particleNum * 2) ;
const workBucketsArray = new Uint32Array( particleNum * 2) ;
const bucketRefArray = new Uint32Array( bucketIndexNum * 2 );
const neighborsArray = new Uint32Array( particleNum );
// test bucketIndex
var buckets = [];
for ( let i = 0; i < particleArray.length / particleSize; i ++ ) {
const r = Math.random() * 0.02 + 0.02;
const rad = Math.random() * Math.PI;
const rad2 = Math.random() * Math.PI * 2.0;
particleArray[ i * particleSize + 0 ] = Math.random();
particleArray[ i * particleSize + 1 ] = Math.random();
particleArray[ i * particleSize + 2 ] = 0.0;
velocityArray[ i * particleSize + 0 ] = r * Math.sin( rad ) * Math.cos( rad2 ) * 0.06;
velocityArray[ i * particleSize + 1 ] = r * Math.sin( rad ) * Math.sin( rad2 ) * 0.06;
velocityArray[ i * particleSize + 2 ] = 0.0;
if (i == 0) {
colorArray[ i * 3 + 0 ] = 1.0;
colorArray[ i * 3 + 1 ] = 1.0;
colorArray[ i * 3 + 2 ] = 0.0;
} else {
colorArray[ i * 3 + 0 ] = 1.0;
colorArray[ i * 3 + 1 ] = 0.0;
colorArray[ i * 3 + 2 ] = 0.0;
}
// test bucketIndex ( cpu buckets )
var bucketX = Math.floor(particleArray[i * particleSize + 0] / bucketSize);
var bucketY = Math.floor(particleArray[i * particleSize + 1] / bucketSize);
var bucketIndex = bucketX + bucketY * bucketNum;
buckets.push([bucketIndex, i]);
}
console.log("neighborsArray: ", neighborsArray);
// test buckets (cpu)
// initial buckets data
var buckets_initial = Array.from(buckets);
console.log("initial buckets (cpu): ", buckets_initial);
// sort (cpu)
for (let p = 0; p < powerN; p++) {
for(let q = 0; q <=p; q++) {
const d = 1 << (p - q);
for (let i = 0; i < buckets.length; i++) {
const up = ((i >> p) & 2) === 0;
if ((i & d) == 0 && (buckets[i][0] > buckets[i | d][0]) === up) {
const tmp = buckets[i];
buckets[i] = buckets[i | d];
buckets[i | d] = tmp;
}
}
}
}
console.log("buckets soerted (cpu): ", buckets);
// initialize bucketRef data
for (let i = 0 ; i < bucketIndexNum ; i++){
bucketRefArray[i * 2 + 0] = 65535;
bucketRefArray[i * 2 + 1] = 65535;
}
console.log("bucketRef initialized: ", bucketRefArray);
// cpu bucketRef
var cpu_bucketRef = new Array(bucketIndexNum);
for(let i = 0; i < bucketIndexNum; i++) {
cpu_bucketRef[i] = new Array(2).fill(65535);
}
console.log("cpu bucketRef: ", cpu_bucketRef);
for (let i = 0; i <= particleNum; i++) {
// dispatch PARTICLE_NUM + 1 ( 0 to PARTICLE_NUM)
var refIndex_start;
var refIndex_end;
if (i == 0) {
refIndex_start = buckets[i][0];
cpu_bucketRef[refIndex_start][0] = 0;
} else if (i < particleNum) {
refIndex_start = buckets[i][0];
refIndex_end = buckets[i - 1][0];
if (refIndex_start != refIndex_end) {
cpu_bucketRef[refIndex_start][0] = i;
cpu_bucketRef[refIndex_end][1] = i - 1;
}
} else {
refIndex_end = buckets[(i - 1)][0];
cpu_bucketRef[refIndex_end][1] = i - 1;
}
}
const particleAttribute = new THREE.InstancedBufferAttribute( particleArray, particleSize );
const velocityAttribute = new THREE.BufferAttribute( velocityArray, particleSize );
const colorAttribute = new THREE.InstancedBufferAttribute( colorArray, 3 );
const bucketsAttribute = new THREE.BufferAttribute( bucketsArray, 2 );
const workBucketsAttribute = new THREE.BufferAttribute( workBucketsArray, 2 );
const bucketRefAttribute = new THREE.BufferAttribute( bucketRefArray, 2 );
const neighborsAttribute = new THREE.InstancedBufferAttribute( neighborsArray, 1 );
const particleBuffer = new WebGPUStorageBuffer( 'particle', particleAttribute );
const velocityBuffer = new WebGPUStorageBuffer( 'velocity', velocityAttribute );
const bucketsBuffer = new WebGPUStorageBuffer( 'buckets', bucketsAttribute );
const workBucketsBuffer = new WebGPUStorageBuffer( 'workData', workBucketsAttribute );
const bucketRefBuffer = new WebGPUStorageBuffer( 'bucketRef', bucketRefAttribute );
const neighborsBuffer = new WebGPUStorageBuffer( 'neighbors', neighborsAttribute );
// compute shader uniform
stepPQUniform = new Vector2Uniform( 'pq' );
const stepPQGroup = new WebGPUUniformsGroup( 'stepPQ' );
stepPQGroup.addUniform( stepPQUniform );
const computeBindings_initialize_buckets = [
particleBuffer,
bucketsBuffer
];
const computeBindings_sort_buckets = [
bucketsBuffer,
workBucketsBuffer,
stepPQGroup
];
const computeBindings_work = [
bucketsBuffer,
workBucketsBuffer
];
const computeBindings_bucketRef = [
bucketsBuffer,
bucketRefBuffer
];
const computeBindings_findNeighbors = [
particleBuffer,
bucketsBuffer,
bucketRefBuffer,
neighborsBuffer
];
const computeBindings_clearNeighbors = [
neighborsBuffer
];
const computeBindings_integrate = [
particleBuffer,
velocityBuffer,
bucketsBuffer,
workBucketsBuffer,
bucketRefBuffer,
neighborsBuffer
];
const computeShader_initialize_buckets = `
#version 450
// constants
#define PARTICLE_NUM ${particleNum}
#define PARTICLE_SIZE ${particleSize}
#define BUCKET_NUM ${bucketNum}
#define BUCKET_SIZE ${bucketSize}
layout(set = 0, binding = 0) buffer Particle {
float particle[ PARTICLE_NUM * PARTICLE_SIZE ];
} particle;
layout(set = 0, binding = 1) buffer Buckets {
uint buckets[ PARTICLE_NUM * 2 ];
};
void main()
{
uint i = gl_GlobalInvocationID.x;
if ( i >= PARTICLE_NUM ) { return; }
// buckets index
vec3 position = vec3(
particle.particle[ i * PARTICLE_SIZE + 0 ],
particle.particle[ i * PARTICLE_SIZE + 1 ],
particle.particle[ i * PARTICLE_SIZE + 2 ]
);
uint bucketX = uint( floor(position.x / BUCKET_SIZE) );
uint bucketY = uint( floor(position.y / BUCKET_SIZE) );
uint bucketIndex = bucketX + bucketY * BUCKET_NUM;
buckets[ i * 2 + 0 ] = bucketIndex;
buckets[ i * 2 + 1 ] = i;
}
`;
const computeShader_sort_buckets = `
#version 450
// constants
#define PARTICLE_NUM ${particleNum}
layout(set = 0, binding = 0) buffer Buckets {
uint buckets[ PARTICLE_NUM * 2 ];
};
layout(set = 0, binding = 1) buffer WorkData {
uint numbers[ PARTICLE_NUM * 2 ];
} workData;
layout(set = 0, binding = 2) uniform StepUniform {
vec2 pq;
} stepPQ;
void main()
{
uint index = gl_GlobalInvocationID.x;
if ( index >= PARTICLE_NUM ) { return; }
uint p = uint(stepPQ.pq.x);
uint q = uint(stepPQ.pq.y);
uint d = 1u << (p - q);
bool up = ((index >> p) & 2u) == 0u;
if ((index & d) == 0 && (buckets[index * 2 + 0] > buckets[(index + d) * 2 + 0]) == up) {
workData.numbers[index * 2 + 0] = buckets[(index + d) * 2 + 0];
workData.numbers[index * 2 + 1] = buckets[(index + d) * 2 + 1];
} else if ((index & d) == d && (buckets[(index - d) * 2 + 0] > buckets[index * 2 + 0]) == up) {
workData.numbers[index * 2 + 0] = buckets[(index - d) * 2 + 0];
workData.numbers[index * 2 + 1] = buckets[(index - d) * 2 + 1];
} else {
workData.numbers[index * 2 + 0] = buckets[index * 2 + 0];
workData.numbers[index * 2 + 1] = buckets[index * 2 + 1];
}
}
`;
const computeShader_work = `
#version 450
#define PARTICLE_NUM ${particleNum}
layout(set = 0, binding = 0) buffer Buckets {
uint buckets[ PARTICLE_NUM * 2 ];
};
layout(set = 0, binding = 1) buffer WorkData {
uint numbers[ PARTICLE_NUM * 2 ];
} workData;
void main() {
uint index = gl_GlobalInvocationID.x;
buckets[index * 2 + 0] = workData.numbers[index * 2 + 0];
buckets[index * 2 + 1] = workData.numbers[index * 2 + 1];
}
`;
const computeShader_bucketRef = `
#version 450
#define PARTICLE_NUM ${particleNum}
#define BUCKETIDX_NUM ${bucketIndexNum}
layout(set = 0, binding = 0) buffer BucketsData {
uint buckets[ PARTICLE_NUM * 2 ];
};
layout(set = 0, binding = 1) buffer BucketRef {
uint bucketRef[ BUCKETIDX_NUM * 2 ];
};
void main() {
// dispatch PARTICLE_NUM + 1 ( 0 to PARTICLE_NUM)
uint index = gl_GlobalInvocationID.x;
if (index == 0) {
uint refIndex_start = buckets[index * 2 + 0];
bucketRef[refIndex_start * 2 + 0] = 0;
} else if (index < PARTICLE_NUM) {
uint refIndex_start = buckets[index * 2 + 0];
uint refIndex_end = buckets[(index - 1) * 2 + 0];
if (refIndex_start != refIndex_end) {
bucketRef[refIndex_start * 2 + 0] = index;
bucketRef[refIndex_end * 2 + 1] = index - 1;
}
} else {
uint refIndex_end = buckets[(index - 1) * 2 + 0];
bucketRef[refIndex_end * 2 + 1] = index - 1;
}
}
`;
const computeShader_findNeighbors = `
#version 450
#define PARTICLE_NUM ${particleNum}
#define PARTICLE_SIZE ${particleSize}
#define BUCKET_NUM ${bucketNum}
#define BUCKET_SIZE ${bucketSize}
#define BUCKETIDX_NUM ${bucketIndexNum}
#define PARTICLE_INDEX ${particleIndex}
layout(set = 0, binding = 0) buffer Particle {
float particle[ PARTICLE_NUM * PARTICLE_SIZE ];
} particle;
layout(set = 0, binding = 1) buffer BucketsData {
uint buckets[ PARTICLE_NUM * 2 ];
};
layout(set = 0, binding = 2) buffer BucketRef {
uint bucketRef[ BUCKETIDX_NUM * 2 ];
};
layout(set = 0, binding = 3) buffer Neighbors {
uint neighbors[ PARTICLE_NUM ];
};
void findNeighbors(uint index, ivec2 bucketPosition, uint neighborFlag) {
if (bucketPosition.x < 0 || bucketPosition.x >= BUCKET_NUM || bucketPosition.y < 0 || bucketPosition.y >= BUCKET_NUM) {
return;
}
uint bucketIndex = bucketPosition.x + BUCKET_NUM * bucketPosition.y;
uint cellStart = bucketRef[bucketIndex * 2 + 0];
uint cellEnd = bucketRef[bucketIndex * 2 + 1];
if (cellStart == 65535 || cellEnd == 65535) {
return;
}
for (uint i = cellStart; i <= cellEnd; i++) {
if (buckets[i * 2 + 1] == index) {
neighbors[index] = neighborFlag;
}
}
}
void main() {
uint index = gl_GlobalInvocationID.x;
uint pIndex = PARTICLE_INDEX;
vec2 position = vec2(
particle.particle[ pIndex * PARTICLE_SIZE + 0 ],
particle.particle[ pIndex * PARTICLE_SIZE + 1 ]
);
vec2 bucketPosition = position / BUCKET_SIZE;
int xOffset = fract(bucketPosition.x) < 0.5 ? -1: 1;
int yOffset = fract(bucketPosition.y) < 0.5 ? -1: 1;
ivec2 bucketPosition00 = ivec2(bucketPosition);
ivec2 bucketPosition10 = bucketPosition00 + ivec2(xOffset, 0);
ivec2 bucketPosition01 = bucketPosition00 + ivec2(0, yOffset);
ivec2 bucketPosition11 = bucketPosition00 + ivec2(xOffset, yOffset);
findNeighbors(index, bucketPosition00, 0);
findNeighbors(index, bucketPosition10, 100);
findNeighbors(index, bucketPosition01, 100);
findNeighbors(index, bucketPosition11, 100);
}
`;
const computeShader_clearNeighbors = `
#version 450
#define PARTICLE_NUM ${particleNum}
layout(set = 0, binding = 0) buffer Neighbors {
uint neighbors[ PARTICLE_NUM ];
};
void main() {
uint index = gl_GlobalInvocationID.x;
neighbors[index] = 65535;
}
`;
const computeShader_integrate = `
#version 450
#define PARTICLE_NUM ${particleNum}
#define PARTICLE_SIZE ${particleSize}
#define BUCKETIDX_NUM ${bucketIndexNum}
// Limitation for now: the order should be the same as bindings order
layout(set = 0, binding = 0) buffer Particle {
float particle[ PARTICLE_NUM * PARTICLE_SIZE ];
} particle;
layout(set = 0, binding = 1) buffer Velocity {
float velocity[ PARTICLE_NUM * PARTICLE_SIZE ];
} velocity;
layout(set = 0, binding = 2) buffer Buckets {
uint buckets[ PARTICLE_NUM * 2 ];
};
layout(set = 0, binding = 3) buffer WorkData {
uint numbers[ PARTICLE_NUM * 2 ];
} workData;
layout(set = 0, binding = 4) buffer BucketRef {
uint bucketRef[ BUCKETIDX_NUM * 2 ];
};
layout(set = 0, binding = 5) buffer Neighbors {
uint neighbors[ PARTICLE_NUM ];
};
void main() {
uint index = gl_GlobalInvocationID.x;
if ( index >= PARTICLE_NUM ) { return; }
vec3 position = vec3(
particle.particle[ index * PARTICLE_SIZE + 0 ] + velocity.velocity[ index * PARTICLE_SIZE + 0 ],
particle.particle[ index * PARTICLE_SIZE + 1 ] + velocity.velocity[ index * PARTICLE_SIZE + 1 ],
particle.particle[ index * PARTICLE_SIZE + 2 ] + velocity.velocity[ index * PARTICLE_SIZE + 2 ]
);
if ( position.x <= 0.001 || position.x >= 0.999 ) {
velocity.velocity[ index * PARTICLE_SIZE + 0 ] = - velocity.velocity[ index * PARTICLE_SIZE + 0 ];
}
if ( position.y <= 0.001 || position.y >= 0.999 ) {
velocity.velocity[ index * PARTICLE_SIZE + 1 ] = - velocity.velocity[ index * PARTICLE_SIZE + 1 ];
}
if ( position.z <= 0.001 || position.z >= 0.999 ) {
velocity.velocity[ index * PARTICLE_SIZE + 2 ] = - velocity.velocity[ index * PARTICLE_SIZE + 2 ];
}
particle.particle[ index * PARTICLE_SIZE + 0 ] = position.x;
particle.particle[ index * PARTICLE_SIZE + 1 ] = position.y;
particle.particle[ index * PARTICLE_SIZE + 2 ] = position.z;
}
`;
computeParams_initialize_buckets.push( {
num: particleNum,
shader: computeShader_initialize_buckets,
bindings: computeBindings_initialize_buckets
});
computeParams_sort_buckets.push( {
num: particleNum,
shader: computeShader_sort_buckets,
bindings: computeBindings_sort_buckets
});
computeParams_work.push( {
num: particleNum,
shader: computeShader_work,
bindings: computeBindings_work
});
computeParams_bucketRef.push( {
num: particleNum + 1,
shader: computeShader_bucketRef,
bindings: computeBindings_bucketRef
});
computeParams_findNeighbors.push( {
num: particleNum,
shader: computeShader_findNeighbors,
bindings: computeBindings_findNeighbors
});
computeParams_clearNeighbors.push( {
num: particleNum,
shader: computeShader_clearNeighbors,
bindings: computeBindings_clearNeighbors
});
computeParams_integrate.push( {
num: particleNum,
shader: computeShader_integrate,
bindings: computeBindings_integrate
});
const boxSize = 0.02;
const boxGeometry = new THREE.BoxBufferGeometry( boxSize, boxSize, boxSize );
const geometry = new THREE.InstancedBufferGeometry()
.setAttribute( 'position', boxGeometry.getAttribute( 'position' ) )
.setAttribute( 'instancePosition', particleAttribute )
.setAttribute( 'instanceColor', colorAttribute )
.setAttribute( 'neighbor', neighborsAttribute );
geometry.setIndex( boxGeometry.getIndex() );
geometry.instanceCount = particleNum;
const material = new THREE.RawShaderMaterial( {
vertexShader:
`#version 450
#define PARTICLE_NUM ${particleNum}
#define PARTICLE_SIZE ${particleSize}
layout(location = 0) in vec3 position;
layout(location = 1) in vec3 instancePosition;
layout(location = 2) in vec3 instanceColor;
layout(location = 3) in int neighbor; // need int (not uint)
layout(location = 0) out vec3 vColor;
layout(set = 0, binding = 0) uniform ModelUniforms {
mat4 modelMatrix;
mat4 modelViewMatrix;
mat3 normalMatrix;
} modelUniforms;
layout(set = 0, binding = 1) uniform CameraUniforms {
mat4 projectionMatrix;
mat4 viewMatrix;
} cameraUniforms;
void main(){
uint id = gl_InstanceIndex;
if (id == 0) {
vColor = vec3(1.0, 0.0, 0.0);
} else if (neighbor == 0) {
vColor = vec3(1.0, 1.0, 0.0);
} else if (neighbor == 100) {
vColor = vec3(0.0, 1.0, 0.0);
} else {
vColor = vec3(0.7, 0.7, 0.7);
}
gl_Position = cameraUniforms.projectionMatrix * modelUniforms.modelViewMatrix * vec4( position + instancePosition, 1.0 );
}
`,
fragmentShader:
`#version 450
#define PARTICLE_NUM ${particleNum}
layout(location = 0) in vec3 vColor;
layout(location = 0) out vec4 outColor;
void main() {
outColor = vec4( vColor, 1.0 );
}
`
} );
const bindings = [];
const modelViewUniform = new Matrix4Uniform( 'modelMatrix' );
const modelViewMatrixUniform = new Matrix4Uniform( 'modelViewMatrix' );
const normalMatrixUniform = new Matrix3Uniform( 'normalMatrix' );
const modelGroup = new WebGPUUniformsGroup( 'modelUniforms' );
modelGroup.addUniform( modelViewUniform );
modelGroup.addUniform( modelViewMatrixUniform );
modelGroup.addUniform( normalMatrixUniform );
modelGroup.setOnBeforeUpdate( function ( object/*, camera */ ) {
modelViewUniform.setValue( object.matrixWorld );
modelViewMatrixUniform.setValue( object.modelViewMatrix );
normalMatrixUniform.setValue( object.normalMatrix );
} );
const projectionMatrixUniform = new Matrix4Uniform( 'projectionMatrix' );
const viewMatrixUniform = new Matrix4Uniform( 'viewMatrix' );
const cameraGroup = new WebGPUUniformsGroup( 'cameraUniforms' );
cameraGroup.addUniform( projectionMatrixUniform );
cameraGroup.addUniform( viewMatrixUniform );
cameraGroup.setOnBeforeUpdate( function ( object, camera ) {
projectionMatrixUniform.setValue( camera.projectionMatrix );
viewMatrixUniform.setValue( camera.matrixWorldInverse );
} );
bindings.push( modelGroup );
bindings.push( cameraGroup );
material.bindings = bindings;
const mesh = new THREE.Mesh( geometry, material );
scene.add( mesh );
renderer = new WebGPURenderer();
renderer.setPixelRatio( window.devicePixelRatio );
renderer.setSize( window.innerWidth, window.innerHeight );
document.body.appendChild( renderer.domElement );
window.addEventListener( 'resize', onWindowResize, false );
return renderer.init();
}
function onWindowResize() {
camera.aspect = window.innerWidth / window.innerHeight;
camera.updateProjectionMatrix();
renderer.setSize( window.innerWidth, window.innerHeight );
}
function animate() {
//
requestAnimationFrame( animate );
// clear neighbors
renderer.compute( computeParams_clearNeighbors );
renderer.compute( computeParams_initialize_buckets );
for (let p = 0; p < powerN; p++) {
for (let q = 0; q <= p; q++) {
stepPQUniform.setValue( new THREE.Vector2(p, q) );
renderer.compute( computeParams_sort_buckets );
renderer.compute( computeParams_work );
}
}
renderer.compute( computeParams_bucketRef );
renderer.compute( computeParams_findNeighbors );
renderer.compute( computeParams_integrate );
renderer.render( scene, camera );
}
function error( error ) {
console.error( error );
}
</script>
</body>
</html>
