wip, compute shader improvements, texture re-use

src/mol-gl/compute/histogram-pyramid/reduction.ts

@@ -4,7 +4,7 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { createComputeRenderable } from '../../renderable'
+import { createComputeRenderable, ComputeRenderable } from '../../renderable'
 import { WebGLContext } from '../../webgl/context';
 import { createComputeRenderItem } from '../../webgl/render-item';
 import { Values, TextureSpec, UniformSpec } from '../../renderable/schema';
@@ -22,26 +22,46 @@ const HistopyramidReductionSchema = {
     uSize: UniformSpec('f'),
 }
 
+let HistopyramidReductionRenderable: ComputeRenderable<Values<typeof HistopyramidReductionSchema>>
 function getHistopyramidReductionRenderable(ctx: WebGLContext, initialTexture: Texture) {
-    const values: Values<typeof HistopyramidReductionSchema> = {
-        ...QuadValues,
-        tPreviousLevel: ValueCell.create(initialTexture),
-        uSize: ValueCell.create(0),
+    if (HistopyramidReductionRenderable) {
+        ValueCell.update(HistopyramidReductionRenderable.values.tPreviousLevel, initialTexture)
+        HistopyramidReductionRenderable.update()
+        return HistopyramidReductionRenderable
+    } else {
+        const values: Values<typeof HistopyramidReductionSchema> = {
+            ...QuadValues,
+            tPreviousLevel: ValueCell.create(initialTexture),
+            uSize: ValueCell.create(0),
+        }
+
+        const schema = { ...HistopyramidReductionSchema }
+        const shaderCode = ShaderCode(
+            require('mol-gl/shader/quad.vert').default,
+            require('mol-gl/shader/histogram-pyramid/reduction.frag').default
+        )
+        const renderItem = createComputeRenderItem(ctx, 'triangles', shaderCode, schema, values)
+
+        HistopyramidReductionRenderable = createComputeRenderable(renderItem, values);
+        return HistopyramidReductionRenderable
     }
-
-    const schema = { ...HistopyramidReductionSchema }
-    const shaderCode = ShaderCode(
-        require('mol-gl/shader/quad.vert').default,
-        require('mol-gl/shader/histogram-pyramid/reduction.frag').default
-    )
-    const renderItem = createComputeRenderItem(ctx, 'triangles', shaderCode, schema, values)
-
-    return createComputeRenderable(renderItem, values);
 }
 
 /** name for shared framebuffer used for histogram-pyramid operations */
 const FramebufferName = 'histogram-pyramid-reduction'
 
+const LevelTextures: Texture[] = []
+function getLevelTexture(ctx: WebGLContext, level: number) {
+    let tex = LevelTextures[level]
+    const size = Math.pow(2, level)
+    if (tex === undefined) {
+        tex = createTexture(ctx, 'image-float32', 'rgba', 'float', 'nearest')
+        LevelTextures[level] = tex
+    }
+    tex.define(size, size) // always call to set size AND clear
+    return tex
+}
+
 function setRenderingDefaults(gl: GLRenderingContext) {
     gl.disable(gl.CULL_FACE)
     gl.disable(gl.BLEND)
@@ -51,8 +71,6 @@ function setRenderingDefaults(gl: GLRenderingContext) {
 
 export interface HistogramPyramid {
     pyramidTex: Texture
-    totalTex: Texture
-    initialTex: Texture
     count: number
     height: number
     levels: number
@@ -66,33 +84,24 @@ export function createHistogramPyramid(ctx: WebGLContext, inputTexture: Texture)
 
     // This part set the levels
     const levels = Math.ceil(Math.log(inputTextureMaxDim) / Math.log(2))
+    const maxSize = Math.pow(2, levels)
     // console.log('levels', levels)
 
-    const initialTexture = createTexture(ctx, 'image-float32', 'rgba', 'float', 'nearest')
-    initialTexture.load({ array: new Float32Array(4), width: 1, height: 1 })
-    initialTexture.define(Math.pow(2, levels), Math.pow(2, levels))
+    const initialTexture = getLevelTexture(ctx, levels)
 
     const framebuffer = framebufferCache.get(FramebufferName).value
     inputTexture.attachFramebuffer(framebuffer, 0)
-    initialTexture.define(Math.pow(2, levels), Math.pow(2, levels))
     // TODO need to initialize texSubImage2D to make Firefox happy
     gl.copyTexSubImage2D(gl.TEXTURE_2D, 0, 0, 0, 0, 0, inputTexture.width, inputTexture.height);
 
-    const initialTextureMaxDim = Math.max(initialTexture.width, initialTexture.height)
-
     const pyramidTexture = createTexture(ctx, 'image-float32', 'rgba', 'float', 'nearest')
-    pyramidTexture.define(Math.pow(2, levels), Math.pow(2, levels))
+    pyramidTexture.define(maxSize, maxSize)
 
-    // TODO cache globally for reuse
     const levelTextures: Texture[] = []
-    for (let i = 0; i < levels; ++i) {
-        const tex = createTexture(ctx, 'image-float32', 'rgba', 'float', 'nearest')
-        tex.define(Math.pow(2, i), Math.pow(2, i))
-        levelTextures.push(tex)
-    }
+    for (let i = 0; i < levels; ++i) levelTextures.push(getLevelTexture(ctx, i))
 
     const renderable = getHistopyramidReductionRenderable(ctx, initialTexture)
-    renderable.update()
+    setRenderingDefaults(gl)
 
     let offset = 0;
     for (let i = 0; i < levels; i++) {
@@ -102,15 +111,13 @@ export function createHistogramPyramid(ctx: WebGLContext, inputTexture: Texture)
         const size = Math.pow(2, currLevel)
         // console.log('size', size, 'draw-level', currLevel, 'read-level', levels - i)
         gl.clear(gl.COLOR_BUFFER_BIT)
-
-        ValueCell.update(renderable.values.uSize, Math.pow(2, i + 1) / initialTextureMaxDim)
-        const readTex = i === 0 ? initialTexture : levelTextures[levels - i]
-        // console.log(readTex.width, readTex.height)
-        ValueCell.update(renderable.values.tPreviousLevel, readTex)
-
-        renderable.update()
-        setRenderingDefaults(gl)
         gl.viewport(0, 0, size, size)
+
+        ValueCell.update(renderable.values.uSize, Math.pow(2, i + 1) / maxSize)
+        if (i > 0) {
+            ValueCell.update(renderable.values.tPreviousLevel, levelTextures[levels - i])
+            renderable.update()
+        }
         renderable.render()
 
         pyramidTexture.bind(0)
@@ -138,18 +145,11 @@ export function createHistogramPyramid(ctx: WebGLContext, inputTexture: Texture)
 
     const finalCount = getHistopyramidSum(ctx, levelTextures[0])
     const height = Math.ceil(finalCount / Math.pow(2, levels))
-    // console.log('height', height, 'finalCount', finalCount)
-
-    //
+    const scale = Vec2.create(maxSize / inputTexture.width, maxSize / inputTexture.height)
+    // console.log('height', height, 'finalCount', finalCount, 'scale', scale)
 
-    const scale = Vec2.create(
-        initialTexture.width / inputTexture.width,
-        initialTexture.height / inputTexture.height
-    )
     return {
         pyramidTex: pyramidTexture,
-        totalTex: levelTextures[0],
-        initialTex: initialTexture,
         count: finalCount,
         height,
         levels,

src/mol-gl/compute/histogram-pyramid/sum.ts

@@ -4,7 +4,7 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { createComputeRenderable } from '../../renderable'
+import { createComputeRenderable, ComputeRenderable } from '../../renderable'
 import { WebGLContext } from '../../webgl/context';
 import { createComputeRenderItem } from '../../webgl/render-item';
 import { Values, TextureSpec } from '../../renderable/schema';
@@ -19,20 +19,36 @@ const HistopyramidSumSchema = {
     tTexture: TextureSpec('texture', 'rgba', 'float', 'nearest'),
 }
 
+let HistopyramidSumRenderable: ComputeRenderable<Values<typeof HistopyramidSumSchema>>
 function getHistopyramidSumRenderable(ctx: WebGLContext, texture: Texture) {
-    const values: Values<typeof HistopyramidSumSchema> = {
-        ...QuadValues,
-        tTexture: ValueCell.create(texture),
-    }
+    if (HistopyramidSumRenderable) {
+        ValueCell.update(HistopyramidSumRenderable.values.tTexture, texture)
+        HistopyramidSumRenderable.update()
+        return HistopyramidSumRenderable
+    } else {
+        const values: Values<typeof HistopyramidSumSchema> = {
+            ...QuadValues,
+            tTexture: ValueCell.create(texture),
+        }
+
+        const schema = { ...HistopyramidSumSchema }
+        const shaderCode = ShaderCode(
+            require('mol-gl/shader/quad.vert').default,
+            require('mol-gl/shader/histogram-pyramid/sum.frag').default
+        )
+        const renderItem = createComputeRenderItem(ctx, 'triangles', shaderCode, schema, values)
 
-    const schema = { ...HistopyramidSumSchema }
-    const shaderCode = ShaderCode(
-        require('mol-gl/shader/quad.vert').default,
-        require('mol-gl/shader/histogram-pyramid/sum.frag').default
-    )
-    const renderItem = createComputeRenderItem(ctx, 'triangles', shaderCode, schema, values)
+        HistopyramidSumRenderable = createComputeRenderable(renderItem, values)
+        return HistopyramidSumRenderable
+    }
+}
 
-    return createComputeRenderable(renderItem, values);
+let SumTexture: Texture
+function getSumTexture(ctx: WebGLContext) {
+    if (SumTexture) return SumTexture
+    SumTexture = createTexture(ctx, 'image-uint8', 'rgba', 'ubyte', 'nearest')
+    SumTexture.define(1, 1)
+    return SumTexture
 }
 
 /** name for shared framebuffer used for histogram-pyramid operations */
@@ -42,21 +58,14 @@ const sumArray = new Uint8Array(4)
 export function getHistopyramidSum(ctx: WebGLContext, pyramidTopTexture: Texture) {
     const { gl, framebufferCache } = ctx
 
-    const framebuffer = framebufferCache.get(FramebufferName).value
-    framebuffer.bind()
-    gl.viewport(0, 0, 1, 1)
-
     const renderable = getHistopyramidSumRenderable(ctx, pyramidTopTexture)
-    renderable.update()
-
-    // TODO cache globally for reuse
-    const sumTexture = createTexture(ctx, 'image-uint8', 'rgba', 'ubyte', 'nearest')
-    sumTexture.define(1, 1)
+    const framebuffer = framebufferCache.get(FramebufferName).value
+    const sumTexture = getSumTexture(ctx)
     sumTexture.attachFramebuffer(framebuffer, 0)
 
+    gl.viewport(0, 0, 1, 1)
     renderable.render()
     ctx.readPixels(0, 0, 1, 1, sumArray)
-
     ctx.unbindFramebuffer()
 
     return decodeFloatRGB(sumArray[0], sumArray[1], sumArray[2])

src/mol-gl/compute/marching-cubes/active-voxels.ts

@@ -70,7 +70,6 @@ export function calcActiveVoxels(ctx: WebGLContext, cornerTex: Texture, gridDime
     activeVoxelsTex.define(width, height)
 
     const renderable = getActiveVoxelsRenderable(ctx, cornerTex, gridDimensions, isoValue)
-    renderable.update()
 
     activeVoxelsTex.attachFramebuffer(framebuffer, 0)
     setRenderingDefaults(gl)

src/mol-gl/compute/marching-cubes/isosurface.ts

@@ -27,11 +27,11 @@ const IsosurfaceSchema = {
     tActiveVoxelsPyramid: TextureSpec('texture', 'rgba', 'float', 'nearest'),
     tActiveVoxelsBase: TextureSpec('texture', 'rgba', 'float', 'nearest'),
     tVolumeData: TextureSpec('texture', 'rgba', 'ubyte', 'nearest'),
-    tActiveVoxelsTotal: TextureSpec('texture', 'rgba', 'float', 'nearest'),
     uIsoValue: UniformSpec('f'),
 
     uSize: UniformSpec('f'),
     uLevels: UniformSpec('f'),
+    uCount: UniformSpec('f'),
 
     uGridDim: UniformSpec('v3'),
     uGridTexDim: UniformSpec('v3'),
@@ -40,7 +40,7 @@ const IsosurfaceSchema = {
     uScale: UniformSpec('v2'),
 }
 
-function getIsosurfaceRenderable(ctx: WebGLContext, activeVoxelsPyramid: Texture, activeVoxelsBase: Texture, volumeData: Texture, activeVoxelsTotal: Texture, gridDimensions: Vec3, transform: Mat4, isoValue: number, levels: number, scale: Vec2) {
+function getIsosurfaceRenderable(ctx: WebGLContext, activeVoxelsPyramid: Texture, activeVoxelsBase: Texture, volumeData: Texture, gridDimensions: Vec3, transform: Mat4, isoValue: number, levels: number, scale: Vec2, count: number) {
     // console.log('uSize', Math.pow(2, levels))
     const values: Values<typeof IsosurfaceSchema> = {
         ...QuadValues,
@@ -49,11 +49,11 @@ function getIsosurfaceRenderable(ctx: WebGLContext, activeVoxelsPyramid: Texture
         tActiveVoxelsPyramid: ValueCell.create(activeVoxelsPyramid),
         tActiveVoxelsBase: ValueCell.create(activeVoxelsBase),
         tVolumeData: ValueCell.create(volumeData),
-        tActiveVoxelsTotal: ValueCell.create(activeVoxelsTotal),
         uIsoValue: ValueCell.create(isoValue),
 
         uSize: ValueCell.create(Math.pow(2, levels)),
         uLevels: ValueCell.create(levels),
+        uCount: ValueCell.create(count),
 
         uGridDim: ValueCell.create(gridDimensions),
         uGridTexDim: ValueCell.create(Vec3.create(volumeData.width, volumeData.height, 0)),
@@ -82,7 +82,7 @@ function setRenderingDefaults(gl: GLRenderingContext) {
 
 export function createIsosurfaceBuffers(ctx: WebGLContext, activeVoxelsBase: Texture, volumeData: Texture, histogramPyramid: HistogramPyramid, gridDimensions: Vec3, transform: Mat4, isoValue: number) {
     const { gl, framebufferCache } = ctx
-    const { pyramidTex, totalTex, height, levels, scale, count } = histogramPyramid
+    const { pyramidTex, height, levels, scale, count } = histogramPyramid
 
     const framebuffer = framebufferCache.get(FramebufferName).value
     framebuffer.bind()
@@ -108,8 +108,7 @@ export function createIsosurfaceBuffers(ctx: WebGLContext, activeVoxelsBase: Tex
     // const indexTex = createTexture(ctx, 'image-float32', 'rgba', 'float', 'nearest')
     // indexTex.define(pyramidTex.width, pyramidTex.height)
 
-    const pr = getIsosurfaceRenderable(ctx, pyramidTex, activeVoxelsBase, volumeData, totalTex, gridDimensions, transform, isoValue, levels, scale)
-    pr.update()
+    const renderable = getIsosurfaceRenderable(ctx, pyramidTex, activeVoxelsBase, volumeData, gridDimensions, transform, isoValue, levels, scale, count)
 
     vertexGroupTexture.attachFramebuffer(framebuffer, 0)
     normalTexture.attachFramebuffer(framebuffer, 1)
@@ -134,7 +133,7 @@ export function createIsosurfaceBuffers(ctx: WebGLContext, activeVoxelsBase: Tex
     setRenderingDefaults(gl)
     gl.viewport(0, 0, pyramidTex.width, pyramidTex.height)
     gl.scissor(0, 0, pyramidTex.width, height)
-    pr.render()
+    renderable.render()
     gl.disable(gl.SCISSOR_TEST)
 
     // const vgt = readTexture(ctx, vertexGroupTexture, pyramidTex.width, height)

src/mol-gl/shader/marching-cubes/isosurface.frag

@@ -3,13 +3,13 @@ precision highp sampler2D;
 
 uniform sampler2D tActiveVoxelsPyramid;
 uniform sampler2D tActiveVoxelsBase;
-uniform sampler2D tActiveVoxelsTotal;
 uniform sampler2D tVolumeData;
 uniform sampler2D tTriIndices;
 
 uniform float uIsoValue;
 uniform float uLevels;
 uniform float uSize;
+uniform float uCount;
 
 uniform vec3 uGridDim;
 uniform vec3 uGridTexDim;
@@ -61,7 +61,7 @@ void main(void) {
     float vI = dot(floor(uSize * vCoordinate), vec2(1.0, uSize));
 
     // ignore 1D indices outside of the grid
-    if(vI >= texture2D(tActiveVoxelsTotal, vec2(0.5)).r) discard;
+    if(vI >= uCount) discard;
 
     float offset = uSize - 2.;
     float k = 1. / uSize;

src/tests/browser/marching-cubes.ts

@@ -61,27 +61,29 @@ async function init() {
     }
     const isoValue = Math.exp(-props.smoothness)
 
-    // console.time('gpu gaussian2')
-    // const densityTextureData2 = await computeGaussianDensityTexture2d(position, box, radius, props, webgl).run()
-    // webgl.waitForGpuCommandsCompleteSync()
-    // console.timeEnd('gpu gaussian2')
-
-    // console.time('gpu mc2')
-    // console.time('gpu mc active2')
-    // const activeVoxelsTex2 = calcActiveVoxels(webgl, densityTextureData2.texture, densityTextureData2.gridDimension, isoValue)
-    // webgl.waitForGpuCommandsCompleteSync()
-    // console.timeEnd('gpu mc active2')
-
-    // console.time('gpu mc pyramid2')
-    // const compacted2 = createHistogramPyramid(webgl, activeVoxelsTex2)
-    // webgl.waitForGpuCommandsCompleteSync()
-    // console.timeEnd('gpu mc pyramid2')
-
-    // console.time('gpu mc vert2')
-    // const gv2 = createIsosurfaceBuffers(webgl, activeVoxelsTex2, densityTextureData2.texture, compacted2, densityTextureData2.gridDimension, densityTextureData2.transform, isoValue)
-    // webgl.waitForGpuCommandsCompleteSync()
-    // console.timeEnd('gpu mc vert2')
-    // console.timeEnd('gpu mc2')
+    // if (true) {
+    //     console.time('gpu gaussian2')
+    //     const densityTextureData2 = await computeGaussianDensityTexture2d(position, box, radius, props, webgl).run()
+    //     webgl.waitForGpuCommandsCompleteSync()
+    //     console.timeEnd('gpu gaussian2')
+
+    //     console.time('gpu mc2')
+    //     console.time('gpu mc active2')
+    //     const activeVoxelsTex2 = calcActiveVoxels(webgl, densityTextureData2.texture, densityTextureData2.gridDimension, isoValue)
+    //     webgl.waitForGpuCommandsCompleteSync()
+    //     console.timeEnd('gpu mc active2')
+
+    //     console.time('gpu mc pyramid2')
+    //     const compacted2 = createHistogramPyramid(webgl, activeVoxelsTex2)
+    //     webgl.waitForGpuCommandsCompleteSync()
+    //     console.timeEnd('gpu mc pyramid2')
+
+    //     console.time('gpu mc vert2')
+    //     const gv2 = createIsosurfaceBuffers(webgl, activeVoxelsTex2, densityTextureData2.texture, compacted2, densityTextureData2.gridDimension, densityTextureData2.transform, isoValue)
+    //     webgl.waitForGpuCommandsCompleteSync()
+    //     console.timeEnd('gpu mc vert2')
+    //     console.timeEnd('gpu mc2')
+    // }
 
     console.time('gpu gaussian')
     const densityTextureData = await computeGaussianDensityTexture2d(position, box, radius, props, webgl).run()
@@ -107,9 +109,7 @@ async function init() {
 
     console.log({ ...webgl.stats, programCount: webgl.programCache.count, shaderCount: webgl.shaderCache.count })
 
-    const mcBoundingSphere = Sphere3D.zero()
-    Sphere3D.addVec3(mcBoundingSphere, mcBoundingSphere, densityTextureData.gridDimension)
-    console.log('mcBoundingSphere', mcBoundingSphere, densityTextureData.gridDimension)
+    const mcBoundingSphere = Sphere3D.fromBox3D(Sphere3D(), densityTextureData.bbox)
     const mcIsosurface = TextureMesh.create(gv.vertexCount, 1, gv.vertexGroupTexture, gv.normalTexture, mcBoundingSphere)
     const mcIsoSurfaceProps = { doubleSided: true, flatShaded: false, alpha: 1.0 }
     const mcIsoSurfaceValues = TextureMesh.Utils.createValuesSimple(mcIsosurface, mcIsoSurfaceProps, Color(0x112299), 1)
@@ -138,6 +138,7 @@ async function init() {
     const surface = await computeMarchingCubesMesh(params).run()
     console.timeEnd('cpu mc')
     // console.log('surface', surface)
+    Mesh.transformImmediate(surface, densityData.transform)
     Mesh.computeNormalsImmediate(surface)
     const meshProps = { doubleSided: true, flatShaded: false, alpha: 1.0 }
     const meshValues = Mesh.Utils.createValuesSimple(surface, meshProps, Color(0x995511), 1)