From 6013b09bd43d2c926a7da6bc6372c58b983f0500 Mon Sep 17 00:00:00 2001 From: rishab Date: Thu, 2 Jan 2025 13:43:32 +0530 Subject: [PATCH 1/2] added compute-shader --- src/core/compute-shader.js | 268 +++++++++++++++++++++++++++++++++++++ src/webgl/index.js | 2 + 2 files changed, 270 insertions(+) create mode 100644 src/core/compute-shader.js diff --git a/src/core/compute-shader.js b/src/core/compute-shader.js new file mode 100644 index 0000000000..16b8bd8590 --- /dev/null +++ b/src/core/compute-shader.js @@ -0,0 +1,268 @@ +class ComputeShader { + constructor(p5Instance, config) { + this.p5 = p5Instance; + this.gl = this.p5._renderer.GL; + + if (!this.gl) { + throw new Error('ComputeShader requires WEBGL mode'); + } + + this.particleCount = config.particleCount || 1000; + this.particleStruct = config.particleStruct; + this.computeFunction = config.computeFunction; + + this._initCapabilities(); + this._initShaders(); + this._initFramebuffers(); + } + + _initCapabilities() { + const gl = this.gl; + + if (gl.getExtension('OES_texture_float') && + (gl.getExtension('WEBGL_color_buffer_float') || gl.getExtension('EXT_color_buffer_float'))) { + this.textureType = gl.FLOAT; + console.log('Using FLOAT textures'); + } else if (gl.getExtension('OES_texture_half_float') && + (gl.getExtension('EXT_color_buffer_half_float') || gl.getExtension('WEBGL_color_buffer_float'))) { + this.textureType = gl.getExtension('OES_texture_half_float').HALF_FLOAT_OES; + console.log('Using HALF_FLOAT textures'); + } else { + console.warn('Float textures not supported. Falling back to UNSIGNED_BYTE.'); + this.textureType = gl.UNSIGNED_BYTE; + } + } + + _initShaders() { + const vertexShader = ` + attribute vec2 aPosition; + varying vec2 vTexCoord; + + void main() { + vTexCoord = aPosition * 0.5 + 0.5; + gl_Position = vec4(aPosition, 0.0, 1.0); + } + `; + + const fragmentShader = this._generateFragmentShader(); + + this.shader = this.p5.createShader(vertexShader, fragmentShader); + } + + _generateFragmentShader() { + const structFields = Object.entries(this.particleStruct) + .map(([name, type]) => ` ${type} ${name};`) + .join('\n'); + + const floatsPerParticle = Object.values(this.particleStruct).reduce((sum, type) => sum + (type === 'float' ? 1 : parseInt(type.slice(3))), 0); + const pixelsPerParticle = Math.ceil(floatsPerParticle / 4); + + return ` + #ifdef GL_ES + precision highp float; + #endif + + uniform sampler2D uState; + uniform vec2 uResolution; + varying vec2 vTexCoord; + + struct Particle { +${structFields} + }; + + ${this._generateReadParticleFunction(floatsPerParticle, pixelsPerParticle)} + ${this._generateWriteParticleFunction(floatsPerParticle, pixelsPerParticle)} + ${this.computeFunction} + + void main() { + ivec2 pixelCoord = ivec2(gl_FragCoord.xy); + int particleIndex = int(pixelCoord.y) * int(uResolution.x) + int(pixelCoord.x); + int pixelIndex = particleIndex / ${pixelsPerParticle}; + + if (float(pixelIndex) >= ${this.particleCount}.0) { + gl_FragColor = vec4(0.0); + return; + } + + Particle particle = readParticle(pixelIndex); + particle = compute(particle); + writeParticle(particle, particleIndex); + } + `; + } + + _generateReadParticleFunction(floatsPerParticle, pixelsPerParticle) { + let readFunction = ` + Particle readParticle(int index) { + Particle p; + int baseIndex = index * ${pixelsPerParticle}; + `; + + let componentIndex = 0; + let pixelOffset = 0; + + for (const [name, type] of Object.entries(this.particleStruct)) { + const components = type === 'float' ? 1 : parseInt(type.slice(3)); + for (let i = 0; i < components; i++) { + readFunction += ` p.${name}${components > 1 ? `[${i}]` : ''} = texture2D(uState, vec2((float(baseIndex + ${pixelOffset}) + 0.5) / uResolution.x, 0.5)).${['r', 'g', 'b', 'a'][componentIndex]};\n`; + componentIndex++; + if (componentIndex === 4) { + componentIndex = 0; + pixelOffset++; + } + } + } + + readFunction += ` + return p; + } + `; + + return readFunction; + } + + _generateWriteParticleFunction(floatsPerParticle, pixelsPerParticle) { + let writeFunction = ` + void writeParticle(Particle p, int particleIndex) { + int pixelIndex = int(mod(float(particleIndex), float(${pixelsPerParticle}))); + vec4 color = vec4(0.0, 0.0, 0.0, 1.0); + `; + + let componentIndex = 0; + let pixelOffset = 0; + + for (const [name, type] of Object.entries(this.particleStruct)) { + const components = type === 'float' ? 1 : parseInt(type.slice(3)); + for (let i = 0; i < components; i++) { + writeFunction += ` if (pixelIndex == ${pixelOffset}) color.${['r', 'g', 'b', 'a'][componentIndex]} = p.${name}${components > 1 ? `[${i}]` : ''};\n`; + componentIndex++; + if (componentIndex === 4) { + componentIndex = 0; + pixelOffset++; + } + } + } + + writeFunction += ` + gl_FragColor = color; + } + `; + + return writeFunction; + } + + _initFramebuffers() { + const floatsPerParticle = Object.values(this.particleStruct).reduce((sum, type) => sum + (type === 'float' ? 1 : parseInt(type.slice(3))), 0); + const pixelsPerParticle = Math.ceil(floatsPerParticle / 4); + this.textureWidth = Math.ceil(Math.sqrt(this.particleCount * pixelsPerParticle)); + this.textureHeight = Math.ceil((this.particleCount * pixelsPerParticle) / this.textureWidth); + + const options = { + format: this.p5.RGBA, + type: this.textureType === this.gl.FLOAT ? this.p5.FLOAT : + this.textureType === this.gl.HALF_FLOAT ? this.p5.HALF_FLOAT : + this.p5.UNSIGNED_BYTE, + width: this.textureWidth, + height: this.textureHeight + }; + + this.inputFramebuffer = this.p5.createFramebuffer(options); + this.outputFramebuffer = this.p5.createFramebuffer(options); + } + + compute() { + this.p5.push(); + this.p5.noStroke(); + + this.outputFramebuffer.begin(); + this.p5.shader(this.shader); + + this.shader.setUniform('uState', this.inputFramebuffer.color); + this.shader.setUniform('uResolution', [this.textureWidth, this.textureHeight]); + + this.p5.quad(-1, -1, 1, -1, 1, 1, -1, 1); + + this.outputFramebuffer.end(); + + // Swap input and output framebuffers + [this.inputFramebuffer, this.outputFramebuffer] = [this.outputFramebuffer, this.inputFramebuffer]; + + this.p5.pop(); + } + + setParticles(particles) { + const floatsPerParticle = Object.values(this.particleStruct).reduce((sum, type) => sum + (type === 'float' ? 1 : parseInt(type.slice(3))), 0); + const pixelsPerParticle = Math.ceil(floatsPerParticle / 4); + const data = new Float32Array(this.textureWidth * this.textureHeight * 4); + + let index = 0; + for (let i = 0; i < particles.length; i++) { + const p = particles[i]; + for (const [name, type] of Object.entries(this.particleStruct)) { + if (type === 'float') { + data[index++] = p[name]; + } else { + const components = parseInt(type.slice(3)); + for (let j = 0; j < components; j++) { + data[index++] = p[name][j]; + } + } + } + // Pad with zeros if necessary + while (index % (pixelsPerParticle * 4) !== 0) { + data[index++] = 0; + } + } + + this.inputFramebuffer.begin(); + this.p5.background(0); + this.inputFramebuffer.loadPixels(); + this.inputFramebuffer.pixels.set(data); + this.inputFramebuffer.updatePixels(); + this.inputFramebuffer.end(); + } + + getParticles() { + const floatsPerParticle = Object.values(this.particleStruct).reduce((sum, type) => sum + (type === 'float' ? 1 : parseInt(type.slice(3))), 0); + const pixelsPerParticle = Math.ceil(floatsPerParticle / 4); + + this.inputFramebuffer.loadPixels(); + const data = this.inputFramebuffer.pixels; + + const particles = []; + let index = 0; + for (let i = 0; i < this.particleCount; i++) { + const particle = {}; + for (const [name, type] of Object.entries(this.particleStruct)) { + if (type === 'float') { + particle[name] = data[index++]; + } else { + const components = parseInt(type.slice(3)); + particle[name] = []; + for (let j = 0; j < components; j++) { + particle[name].push(data[index++]); + } + } + } + // Skip padding + index = (i + 1) * pixelsPerParticle * 4; + particles.push(particle); + } + + return particles; + } +} + +function computeShaderAdditions(p5, fn) { + p5.ComputeShader = ComputeShader; + + fn.createComputeShader = function(config) { + if (!this._renderer || !this._renderer.GL) { + throw new Error('ComputeShader requires WEBGL mode. Use createCanvas(w, h, WEBGL)'); + } + return new ComputeShader(this, config); + }; +} + +export { ComputeShader }; +export default computeShaderAdditions; \ No newline at end of file diff --git a/src/webgl/index.js b/src/webgl/index.js index c2515fce5a..c7017597d6 100644 --- a/src/webgl/index.js +++ b/src/webgl/index.js @@ -14,6 +14,7 @@ import shader from './p5.Shader'; import camera from './p5.Camera'; import texture from './p5.Texture'; import rendererGL from './p5.RendererGL'; +import computeShaderAdditions from '../core/compute-shader'; export default function(p5){ rendererGL(p5, p5.prototype); @@ -32,4 +33,5 @@ export default function(p5){ dataArray(p5, p5.prototype); shader(p5, p5.prototype); texture(p5, p5.prototype); + computeShaderAdditions(p5, p5.prototype); } From 7a10a52394f22e336fbdfdb7aeba4b8c073ff3aa Mon Sep 17 00:00:00 2001 From: rishab Date: Sun, 5 Jan 2025 22:37:35 +0530 Subject: [PATCH 2/2] added tests --- src/core/compute-shader.js | 15 ++-- test/unit/core/compute-shader.js | 125 +++++++++++++++++++++++++++++++ test/unit/spec.js | 3 +- 3 files changed, 137 insertions(+), 6 deletions(-) create mode 100644 test/unit/core/compute-shader.js diff --git a/src/core/compute-shader.js b/src/core/compute-shader.js index 16b8bd8590..2fd9eb12a3 100644 --- a/src/core/compute-shader.js +++ b/src/core/compute-shader.js @@ -8,8 +8,8 @@ class ComputeShader { } this.particleCount = config.particleCount || 1000; - this.particleStruct = config.particleStruct; - this.computeFunction = config.computeFunction; + this.particleStruct = config.particleStruct || {}; + this.computeFunction = config.computeFunction || ''; this._initCapabilities(); this._initShaders(); @@ -70,6 +70,11 @@ class ComputeShader { ${structFields} }; + // Custom random function for the shader + float rand(vec2 co) { + return fract(sin(dot(co.xy ,vec2(12.9898,78.233))) * 43758.5453); + } + ${this._generateReadParticleFunction(floatsPerParticle, pixelsPerParticle)} ${this._generateWriteParticleFunction(floatsPerParticle, pixelsPerParticle)} ${this.computeFunction} @@ -154,8 +159,8 @@ ${structFields} _initFramebuffers() { const floatsPerParticle = Object.values(this.particleStruct).reduce((sum, type) => sum + (type === 'float' ? 1 : parseInt(type.slice(3))), 0); const pixelsPerParticle = Math.ceil(floatsPerParticle / 4); - this.textureWidth = Math.ceil(Math.sqrt(this.particleCount * pixelsPerParticle)); - this.textureHeight = Math.ceil((this.particleCount * pixelsPerParticle) / this.textureWidth); + this.textureWidth = this.particleCount * pixelsPerParticle; + this.textureHeight = 1; const options = { format: this.p5.RGBA, @@ -251,7 +256,7 @@ ${structFields} return particles; } -} +} function computeShaderAdditions(p5, fn) { p5.ComputeShader = ComputeShader; diff --git a/test/unit/core/compute-shader.js b/test/unit/core/compute-shader.js new file mode 100644 index 0000000000..c164469531 --- /dev/null +++ b/test/unit/core/compute-shader.js @@ -0,0 +1,125 @@ +import p5 from '../../../src/app.js'; +import { ComputeShader } from '../../../src/core/compute-shader.js'; + +suite('compute_shader', function() { + let myp5; + + beforeAll(function() { + myp5 = new p5(function(p) { + p.setup = function() { + p.createCanvas(100, 100, p.WEBGL); + }; + }); + }); + + afterAll(function() { + myp5.remove(); + }); + + test('ComputeShader initialization', function() { + const computeShader = new ComputeShader(myp5, { + particleCount: 100, + particleStruct: { + position: 'vec3', + velocity: 'vec2', + age: 'float' + }, + computeFunction: ` + Particle compute(Particle p) { + p.position += vec3(p.velocity, 0.0); + p.age += 0.01; + return p; + } + ` + }); + + assert(computeShader instanceof ComputeShader, 'ComputeShader was not created successfully'); + assert(computeShader.particleCount === 100, 'Particle count was not set correctly'); + assert(Object.keys(computeShader.particleStruct).length === 3, 'Particle struct does not have the correct number of properties'); + assert(computeShader.particleStruct.position === 'vec3', 'Position type is incorrect'); + assert(computeShader.particleStruct.velocity === 'vec2', 'Velocity type is incorrect'); + assert(computeShader.particleStruct.age === 'float', 'Age type is incorrect'); + }); + + test('ComputeShader texture size calculation', function() { + const computeShader = new ComputeShader(myp5, { + particleCount: 1000, + particleStruct: { + position: 'vec3', + velocity: 'vec3', + color: 'vec3', + size: 'float' + }, + computeFunction: ` + Particle compute(Particle p) { + return p; + } + ` + }); + + const expectedPixelsPerParticle = 3; // (3 + 3 + 3 + 1) components / 4 components per pixel, rounded up + const expectedTextureWidth = 1000 * expectedPixelsPerParticle; + + assert(computeShader.textureWidth === expectedTextureWidth, `Texture width should be ${expectedTextureWidth}`); + assert(computeShader.textureHeight === 1, 'Texture height should be 1'); + }); + + test('ComputeShader setParticles and getParticles', function() { + const computeShader = new ComputeShader(myp5, { + particleCount: 2, + particleStruct: { + position: 'vec3', + velocity: 'vec2', + age: 'float' + }, + computeFunction: ` + Particle compute(Particle p) { + return p; + } + ` + }); + + const initialParticles = [ + { position: [0, 0, 0], velocity: [1, 1], age: 0 }, + { position: [1, 1, 1], velocity: [-1, -1], age: 1 } + ]; + + computeShader.setParticles(initialParticles); + const retrievedParticles = computeShader.getParticles(); + + assert(retrievedParticles.length === 2, 'Retrieved particles count is incorrect'); + assert.deepEqual(retrievedParticles[0], initialParticles[0], 'First particle data does not match'); + assert.deepEqual(retrievedParticles[1], initialParticles[1], 'Second particle data does not match'); + }); + + test('ComputeShader compute function', function() { + const computeShader = new ComputeShader(myp5, { + particleCount: 1, + particleStruct: { + position: 'vec3', + velocity: 'vec2', + age: 'float' + }, + computeFunction: ` + Particle compute(Particle p) { + p.position += vec3(p.velocity, 0.0); + p.age += 1.0; + return p; + } + ` + }); + + const initialParticle = [ + { position: [0, 0, 0], velocity: [0.1, 0.2], age: 0 } + ]; + + computeShader.setParticles(initialParticle); + computeShader.compute(); + const updatedParticle = computeShader.getParticles()[0]; + + assert.closeTo(updatedParticle.position[0], 0.1, 0.001, 'X position not updated correctly'); + assert.closeTo(updatedParticle.position[1], 0.2, 0.001, 'Y position not updated correctly'); + assert.closeTo(updatedParticle.position[2], 0, 0.001, 'Z position should not change'); + assert.closeTo(updatedParticle.age, 1, 0.001, 'Age not updated correctly'); + }); +}); \ No newline at end of file diff --git a/test/unit/spec.js b/test/unit/spec.js index 8df31317f2..61c9ddc856 100644 --- a/test/unit/spec.js +++ b/test/unit/spec.js @@ -17,7 +17,8 @@ var spec = { 'structure', 'transform', 'version', - 'vertex' + 'vertex', + 'compute_shader', ], data: ['p5.TypedDict', 'local_storage'], dom: ['dom'],