Skip to content

Instantly share code, notes, and snippets.

@puny-d
Forked from CodyJasonBennett/index.ts
Created May 31, 2023 07:59
Show Gist options
  • Save puny-d/b58edcf74eb4696aa2c8dad88e92e77c to your computer and use it in GitHub Desktop.
Save puny-d/b58edcf74eb4696aa2c8dad88e92e77c to your computer and use it in GitHub Desktop.

Revisions

  1. @CodyJasonBennett CodyJasonBennett revised this gist Nov 8, 2022. 1 changed file with 12 additions and 13 deletions.
    25 changes: 12 additions & 13 deletions index.ts
    Original file line number Diff line number Diff line change
    @@ -86,10 +86,15 @@ export interface Compiled {
    */
    export class WebGLCompute {
    readonly gl: WebGL2RenderingContext
    private _fragmentShader: WebGLShader
    private _compiled = new Map<WebGLComputeOptions, Compiled>()

    constructor(gl = document.createElement('canvas').getContext('webgl2')!) {
    this.gl = gl

    this._fragmentShader = this.gl.createShader(this.gl.FRAGMENT_SHADER)!
    this.gl.shaderSource(this._fragmentShader, '#version 300 es\nvoid main(){}')
    this.gl.compileShader(this._fragmentShader)
    }

    /**
    @@ -124,28 +129,21 @@ export class WebGLCompute {
    this.gl.shaderSource(vertexShader, options.compute)
    this.gl.compileShader(vertexShader)

    const fragmentShader = this.gl.createShader(this.gl.FRAGMENT_SHADER)!
    this.gl.shaderSource(fragmentShader, '#version 300 es\nvoid main(){}')
    this.gl.compileShader(fragmentShader)

    this.gl.attachShader(program, vertexShader)
    this.gl.attachShader(program, fragmentShader)
    this.gl.attachShader(program, this._fragmentShader)

    this.gl.transformFeedbackVaryings(program, outputs, this.gl.SEPARATE_ATTRIBS)
    this.gl.linkProgram(program)

    for (const shader of [vertexShader, fragmentShader]) {
    const error = this.gl.getShaderInfoLog(shader)
    if (error) throw `Error compiling shader: ${error}\n${lineNumbers(this.gl.getShaderSource(shader)!)}`
    }
    const shaderError = this.gl.getShaderInfoLog(vertexShader)
    if (shaderError) throw `Error compiling shader: ${shaderError}\n${lineNumbers(options.compute)}`

    const error = this.gl.getProgramInfoLog(program)
    if (error) throw `Error compiling program: ${this.gl.getProgramInfoLog(program)}`
    const programError = this.gl.getProgramInfoLog(program)
    if (programError) throw `Error compiling program: ${this.gl.getProgramInfoLog(program)}`

    this.gl.detachShader(program, vertexShader)
    this.gl.detachShader(program, fragmentShader)
    this.gl.detachShader(program, this._fragmentShader)
    this.gl.deleteShader(vertexShader)
    this.gl.deleteShader(fragmentShader)

    // Init VAO state (input)
    const VAO = this.gl.createVertexArray()!
    @@ -250,6 +248,7 @@ export class WebGLCompute {
    * Disposes the compute pipeline from GPU memory.
    */
    dispose(): void {
    this.gl.deleteShader(this._fragmentShader)
    for (const [, compiled] of this._compiled) {
    this.gl.deleteProgram(compiled.program)
    this.gl.deleteVertexArray(compiled.VAO)
  2. @CodyJasonBennett CodyJasonBennett revised this gist Nov 8, 2022. 1 changed file with 2 additions and 1 deletion.
    3 changes: 2 additions & 1 deletion index.ts
    Original file line number Diff line number Diff line change
    @@ -90,7 +90,6 @@ export class WebGLCompute {

    constructor(gl = document.createElement('canvas').getContext('webgl2')!) {
    this.gl = gl
    this.gl.enable(this.gl.RASTERIZER_DISCARD)
    }

    /**
    @@ -221,6 +220,7 @@ export class WebGLCompute {
    this.gl.useProgram(compiled.program)
    this.gl.bindVertexArray(compiled.VAO)
    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, compiled.transformFeedback)
    this.gl.enable(this.gl.RASTERIZER_DISCARD)

    this.gl.beginTransformFeedback(this.gl.POINTS)
    this.gl.drawArraysInstanced(this.gl.POINTS, 0, compiled.length, options.instances ?? 1)
    @@ -229,6 +229,7 @@ export class WebGLCompute {
    this.gl.useProgram(null)
    this.gl.bindVertexArray(null)
    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, null)
    this.gl.disable(this.gl.RASTERIZER_DISCARD)

    // Read output buffer data
    const output: WebGLComputeResult = {}
  3. @CodyJasonBennett CodyJasonBennett revised this gist Nov 8, 2022. 1 changed file with 127 additions and 62 deletions.
    189 changes: 127 additions & 62 deletions index.ts
    Original file line number Diff line number Diff line change
    @@ -45,12 +45,21 @@ export interface WebGLComputeInput {
    * The size (per vertex) of the data array. Used to allocate data to each vertex.
    */
    size: number
    /**
    * The size (per instance) of the data array. Used to allocate data to each instance.
    */
    divisor?: number
    /**
    * Flags this input for update.
    */
    needsUpdate?: boolean
    }

    /**
    * WebGLCompute constructor parameters. Accepts a list of program inputs and compute shader source.
    */
    export interface WebGLComputeOptions {
    instances?: number
    inputs: Record<string, WebGLComputeInput>
    compute: string
    }
    @@ -60,143 +69,199 @@ export interface WebGLComputeOptions {
    */
    export type WebGLComputeResult = Record<string, Float32Array>

    /**
    * Represents internal compiled state.
    */
    export interface Compiled {
    program: WebGLProgram
    VAO: WebGLVertexArrayObject
    transformFeedback: WebGLTransformFeedback
    buffers: Map<string, WebGLBuffer>
    containers: Map<string, Float32Array>
    length: number
    }

    /**
    * Constructs a WebGL compute program via transform feedback. Can be used to compute and serialize data from the GPU.
    */
    export class WebGLCompute {
    readonly gl: WebGL2RenderingContext
    readonly program: WebGLProgram
    readonly VAO: WebGLVertexArrayObject
    readonly transformFeedback: WebGLTransformFeedback
    readonly buffers = new Map<string, WebGLBuffer>()
    readonly containers = new Map<string, ArrayBufferView>()
    private _length = 0

    constructor(options: WebGLComputeOptions, gl = document.createElement('canvas').getContext('webgl2')!) {
    private _compiled = new Map<WebGLComputeOptions, Compiled>()

    constructor(gl = document.createElement('canvas').getContext('webgl2')!) {
    this.gl = gl
    this.gl.enable(this.gl.RASTERIZER_DISCARD)
    }

    /**
    * Compiles a transform feedback program from compute options.
    */
    compile(options: WebGLComputeOptions): Compiled {
    let compiled = this._compiled.get(options)
    if (compiled) {
    this.gl.bindVertexArray(compiled.VAO)
    for (const [name, buffer] of compiled.buffers) {
    const input = options.inputs[name]
    if (!input.needsUpdate) continue

    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, buffer)
    this.gl.bufferData(this.gl.ARRAY_BUFFER, input.data, this.gl.DYNAMIC_READ)
    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, null)

    input.needsUpdate = false
    }
    this.gl.bindVertexArray(null)

    return compiled
    }

    // Parse outputs from shader source
    const outputs = Array.from(options.compute.matchAll(VARYING_REGEX)).map(([, varying]) => varying)

    // Compile shaders, configure output varyings
    this.program = this.gl.createProgram()!
    const program = this.gl.createProgram()!

    const vertexShader = this.gl.createShader(this.gl.VERTEX_SHADER)!
    this.gl.shaderSource(vertexShader, options.compute)
    this.gl.compileShader(vertexShader)

    const error = this.gl.getShaderInfoLog(vertexShader)
    if (error) throw `${error}\n${lineNumbers(options.compute)}`

    const fragmentShader = this.gl.createShader(this.gl.FRAGMENT_SHADER)!
    this.gl.shaderSource(fragmentShader, '#version 300 es\nprecision highp float;\nvoid main(){}')
    this.gl.shaderSource(fragmentShader, '#version 300 es\nvoid main(){}')
    this.gl.compileShader(fragmentShader)

    this.gl.attachShader(this.program, vertexShader)
    this.gl.attachShader(this.program, fragmentShader)
    this.gl.attachShader(program, vertexShader)
    this.gl.attachShader(program, fragmentShader)

    this.gl.transformFeedbackVaryings(this.program, outputs, this.gl.SEPARATE_ATTRIBS)
    this.gl.linkProgram(this.program)
    this.gl.transformFeedbackVaryings(program, outputs, this.gl.SEPARATE_ATTRIBS)
    this.gl.linkProgram(program)

    this.gl.detachShader(this.program, vertexShader)
    this.gl.detachShader(this.program, fragmentShader)
    for (const shader of [vertexShader, fragmentShader]) {
    const error = this.gl.getShaderInfoLog(shader)
    if (error) throw `Error compiling shader: ${error}\n${lineNumbers(this.gl.getShaderSource(shader)!)}`
    }

    const error = this.gl.getProgramInfoLog(program)
    if (error) throw `Error compiling program: ${this.gl.getProgramInfoLog(program)}`

    this.gl.detachShader(program, vertexShader)
    this.gl.detachShader(program, fragmentShader)
    this.gl.deleteShader(vertexShader)
    this.gl.deleteShader(fragmentShader)

    if (!this.gl.getProgramParameter(this.program, this.gl.LINK_STATUS)) {
    throw this.gl.getProgramInfoLog(this.program)
    }

    // Init VAO state (input)
    this.VAO = this.gl.createVertexArray()!
    this.gl.bindVertexArray(this.VAO)
    const VAO = this.gl.createVertexArray()!
    this.gl.bindVertexArray(VAO)

    let length = 0

    const buffers = new Map<string, WebGLBuffer>()
    for (const name in options.inputs) {
    const { data, size } = options.inputs[name]
    const input = options.inputs[name]

    const buffer = this.gl.createBuffer()!
    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, buffer)
    this.gl.bufferData(this.gl.ARRAY_BUFFER, data, this.gl.STATIC_READ)
    this.gl.bufferData(this.gl.ARRAY_BUFFER, input.data, this.gl.STATIC_READ)

    const location = this.gl.getAttribLocation(program, name)

    const location = this.gl.getAttribLocation(this.program, name)
    this.gl.enableVertexAttribArray(location)
    const slots = Math.min(4, Math.max(1, Math.floor(input.size / 3)))
    for (let i = 0; i < slots; i++) {
    this.gl.enableVertexAttribArray(location + i)

    const dataType = getDataType(data)!
    if (dataType === this.gl.INT || dataType === this.gl.UNSIGNED_INT) {
    this.gl.vertexAttribIPointer(location, size, dataType, 0, 0)
    } else {
    this.gl.vertexAttribPointer(location, size, dataType, false, 0, 0)
    if (input.data instanceof Float32Array) {
    this.gl.vertexAttribPointer(location, input.size, this.gl.FLOAT, false, 0, 0)
    } else {
    const dataType = getDataType(input.data)!
    this.gl.vertexAttribIPointer(location, input.size, dataType, 0, 0)
    }

    if (input.divisor) this.gl.vertexAttribDivisor(location + i, input.divisor)
    }

    this.buffers.set(name, buffer)
    this._length = Math.max(this._length, (data as unknown as ArrayLike<number>).length / size)
    }
    buffers.set(name, buffer)
    length = Math.max(length, (input.data as unknown as ArrayLike<number>).length / input.size)

    input.needsUpdate = false
    }
    this.gl.bindVertexArray(null)

    // Init feedback state (output)
    this.transformFeedback = this.gl.createTransformFeedback()!
    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, this.transformFeedback)
    const transformFeedback = this.gl.createTransformFeedback()!
    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, transformFeedback)

    for (const name of outputs) {
    const data = new Float32Array(this._length)
    this.containers.set(name, data)
    const containers = new Map<string, Float32Array>()
    for (let i = 0; i < outputs.length; i++) {
    const output = outputs[i]
    const data = new Float32Array(length)
    containers.set(output, data)

    const buffer = this.gl.createBuffer()!
    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, buffer)
    this.gl.bufferData(this.gl.ARRAY_BUFFER, data, this.gl.STATIC_COPY)
    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, null)
    buffers.set(output, buffer)

    this.gl.bindBufferBase(this.gl.TRANSFORM_FEEDBACK_BUFFER, this.containers.size - 1, buffer)
    this.buffers.set(name, buffer)
    this.gl.bindBufferBase(this.gl.TRANSFORM_FEEDBACK_BUFFER, i, buffer)
    }

    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, null)

    compiled = { program, VAO, transformFeedback, buffers, containers, length }
    this._compiled.set(options, compiled)

    return compiled
    }

    /**
    * Runs and reads from the compute program.
    */
    compute(): WebGLComputeResult {
    compute(options: WebGLComputeOptions): WebGLComputeResult {
    const compiled = this.compile(options)

    // Run compute
    this.gl.useProgram(this.program)
    this.gl.bindVertexArray(this.VAO)
    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, this.transformFeedback)
    this.gl.enable(this.gl.RASTERIZER_DISCARD)
    this.gl.useProgram(compiled.program)
    this.gl.bindVertexArray(compiled.VAO)
    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, compiled.transformFeedback)

    this.gl.beginTransformFeedback(this.gl.POINTS)
    this.gl.drawArrays(this.gl.POINTS, 0, this._length)
    this.gl.drawArraysInstanced(this.gl.POINTS, 0, compiled.length, options.instances ?? 1)
    this.gl.endTransformFeedback()

    this.gl.useProgram(null)
    this.gl.bindVertexArray(null)
    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, null)
    this.gl.disable(this.gl.RASTERIZER_DISCARD)

    // Read output buffer data
    return Array.from(this.containers).reduce((acc, [name, data]) => {
    const buffer = this.buffers.get(name)!
    const output: WebGLComputeResult = {}
    for (const [name, data] of compiled.containers) {
    const buffer = compiled.buffers.get(name)!

    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, buffer)
    this.gl.getBufferSubData(this.gl.ARRAY_BUFFER, 0, data)
    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, null)

    return { ...acc, [name]: data }
    }, {})
    output[name] = data
    }

    return output
    }

    /**
    * Disposes the compute pipeline from GPU memory.
    */
    dispose(): void {
    this.gl.deleteProgram(this.program)
    this.gl.deleteVertexArray(this.VAO)
    this.gl.deleteTransformFeedback(this.transformFeedback)
    this.buffers.forEach((buffer) => this.gl.deleteBuffer(buffer))
    for (const [, compiled] of this._compiled) {
    this.gl.deleteProgram(compiled.program)
    this.gl.deleteVertexArray(compiled.VAO)
    this.gl.deleteTransformFeedback(compiled.transformFeedback)
    compiled.buffers.forEach((buffer) => this.gl.deleteBuffer(buffer))
    }
    this._compiled.clear()
    }
    }

    const compute = new WebGLCompute({
    const renderer = new WebGLCompute()
    const result = renderer.compute({
    instances: 1,
    inputs: {
    source: {
    data: new Float32Array([0, 1, 2, 3, 4]),
    @@ -214,4 +279,4 @@ const compute = new WebGLCompute({
    })

    // { result: Float32Array(5) [0, 2, 4, 6, 8] }
    console.log(compute.compute())
    console.log(result)
  4. @CodyJasonBennett CodyJasonBennett revised this gist Aug 8, 2022. 1 changed file with 1 addition and 1 deletion.
    2 changes: 1 addition & 1 deletion index.ts
    Original file line number Diff line number Diff line change
    @@ -44,7 +44,7 @@ export interface WebGLComputeInput {
    /**
    * The size (per vertex) of the data array. Used to allocate data to each vertex.
    */
    size: 1 | 2 | 3 | 4
    size: number
    }

    /**
  5. @CodyJasonBennett CodyJasonBennett created this gist Jul 1, 2022.
    217 changes: 217 additions & 0 deletions index.ts
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,217 @@
    /**
    * Matches against GLSL shader outputs.
    */
    const VARYING_REGEX = /[^\w](?:varying|out)\s+\w+\s+(\w+)\s*;/g

    /**
    * Adds line numbers to a string with an optional starting offset.
    */
    const lineNumbers = (source: string, offset = 0): string => source.replace(/^/gm, () => `${offset++}:`)

    /**
    * Gets the appropriate WebGL data type for a data view.
    */
    const getDataType = (data: ArrayBufferView): number | null => {
    switch (data.constructor) {
    case Float32Array:
    return 5126 // FLOAT
    case Int8Array:
    return 5120 // BYTE
    case Int16Array:
    return 5122 // SHORT
    case Int32Array:
    return 5124 // INT
    case Uint8Array:
    case Uint8ClampedArray:
    return 5121 // UNSIGNED_BYTE
    case Uint16Array:
    return 5123 // UNSIGNED_SHORT
    case Uint32Array:
    return 5125 // UNSIGNED_INT
    default:
    return null
    }
    }

    /**
    * Represents compute input data.
    */
    export interface WebGLComputeInput {
    /**
    * Input data view.
    */
    data: ArrayBufferView
    /**
    * The size (per vertex) of the data array. Used to allocate data to each vertex.
    */
    size: 1 | 2 | 3 | 4
    }

    /**
    * WebGLCompute constructor parameters. Accepts a list of program inputs and compute shader source.
    */
    export interface WebGLComputeOptions {
    inputs: Record<string, WebGLComputeInput>
    compute: string
    }

    /**
    * Represents a compute result.
    */
    export type WebGLComputeResult = Record<string, Float32Array>

    /**
    * Constructs a WebGL compute program via transform feedback. Can be used to compute and serialize data from the GPU.
    */
    export class WebGLCompute {
    readonly gl: WebGL2RenderingContext
    readonly program: WebGLProgram
    readonly VAO: WebGLVertexArrayObject
    readonly transformFeedback: WebGLTransformFeedback
    readonly buffers = new Map<string, WebGLBuffer>()
    readonly containers = new Map<string, ArrayBufferView>()
    private _length = 0

    constructor(options: WebGLComputeOptions, gl = document.createElement('canvas').getContext('webgl2')!) {
    this.gl = gl

    // Parse outputs from shader source
    const outputs = Array.from(options.compute.matchAll(VARYING_REGEX)).map(([, varying]) => varying)

    // Compile shaders, configure output varyings
    this.program = this.gl.createProgram()!

    const vertexShader = this.gl.createShader(this.gl.VERTEX_SHADER)!
    this.gl.shaderSource(vertexShader, options.compute)
    this.gl.compileShader(vertexShader)

    const error = this.gl.getShaderInfoLog(vertexShader)
    if (error) throw `${error}\n${lineNumbers(options.compute)}`

    const fragmentShader = this.gl.createShader(this.gl.FRAGMENT_SHADER)!
    this.gl.shaderSource(fragmentShader, '#version 300 es\nprecision highp float;\nvoid main(){}')
    this.gl.compileShader(fragmentShader)

    this.gl.attachShader(this.program, vertexShader)
    this.gl.attachShader(this.program, fragmentShader)

    this.gl.transformFeedbackVaryings(this.program, outputs, this.gl.SEPARATE_ATTRIBS)
    this.gl.linkProgram(this.program)

    this.gl.detachShader(this.program, vertexShader)
    this.gl.detachShader(this.program, fragmentShader)
    this.gl.deleteShader(vertexShader)
    this.gl.deleteShader(fragmentShader)

    if (!this.gl.getProgramParameter(this.program, this.gl.LINK_STATUS)) {
    throw this.gl.getProgramInfoLog(this.program)
    }

    // Init VAO state (input)
    this.VAO = this.gl.createVertexArray()!
    this.gl.bindVertexArray(this.VAO)

    for (const name in options.inputs) {
    const { data, size } = options.inputs[name]

    const buffer = this.gl.createBuffer()!
    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, buffer)
    this.gl.bufferData(this.gl.ARRAY_BUFFER, data, this.gl.STATIC_READ)

    const location = this.gl.getAttribLocation(this.program, name)
    this.gl.enableVertexAttribArray(location)

    const dataType = getDataType(data)!
    if (dataType === this.gl.INT || dataType === this.gl.UNSIGNED_INT) {
    this.gl.vertexAttribIPointer(location, size, dataType, 0, 0)
    } else {
    this.gl.vertexAttribPointer(location, size, dataType, false, 0, 0)
    }

    this.buffers.set(name, buffer)
    this._length = Math.max(this._length, (data as unknown as ArrayLike<number>).length / size)
    }

    this.gl.bindVertexArray(null)

    // Init feedback state (output)
    this.transformFeedback = this.gl.createTransformFeedback()!
    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, this.transformFeedback)

    for (const name of outputs) {
    const data = new Float32Array(this._length)
    this.containers.set(name, data)

    const buffer = this.gl.createBuffer()!
    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, buffer)
    this.gl.bufferData(this.gl.ARRAY_BUFFER, data, this.gl.STATIC_COPY)
    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, null)

    this.gl.bindBufferBase(this.gl.TRANSFORM_FEEDBACK_BUFFER, this.containers.size - 1, buffer)
    this.buffers.set(name, buffer)
    }

    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, null)
    }

    /**
    * Runs and reads from the compute program.
    */
    compute(): WebGLComputeResult {
    // Run compute
    this.gl.useProgram(this.program)
    this.gl.bindVertexArray(this.VAO)
    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, this.transformFeedback)
    this.gl.enable(this.gl.RASTERIZER_DISCARD)

    this.gl.beginTransformFeedback(this.gl.POINTS)
    this.gl.drawArrays(this.gl.POINTS, 0, this._length)
    this.gl.endTransformFeedback()

    this.gl.useProgram(null)
    this.gl.bindVertexArray(null)
    this.gl.bindTransformFeedback(this.gl.TRANSFORM_FEEDBACK, null)
    this.gl.disable(this.gl.RASTERIZER_DISCARD)

    // Read output buffer data
    return Array.from(this.containers).reduce((acc, [name, data]) => {
    const buffer = this.buffers.get(name)!

    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, buffer)
    this.gl.getBufferSubData(this.gl.ARRAY_BUFFER, 0, data)
    this.gl.bindBuffer(this.gl.ARRAY_BUFFER, null)

    return { ...acc, [name]: data }
    }, {})
    }

    /**
    * Disposes the compute pipeline from GPU memory.
    */
    dispose(): void {
    this.gl.deleteProgram(this.program)
    this.gl.deleteVertexArray(this.VAO)
    this.gl.deleteTransformFeedback(this.transformFeedback)
    this.buffers.forEach((buffer) => this.gl.deleteBuffer(buffer))
    }
    }

    const compute = new WebGLCompute({
    inputs: {
    source: {
    data: new Float32Array([0, 1, 2, 3, 4]),
    size: 1,
    },
    },
    compute: /* glsl */ `#version 300 es
    in float source;
    out float result;
    void main() {
    result = source + float(gl_VertexID);
    }
    `,
    })

    // { result: Float32Array(5) [0, 2, 4, 6, 8] }
    console.log(compute.compute())