wip, gpu gaussian surface

src/mol-gl/webgl/texture.ts

@@ -10,6 +10,7 @@ import { ValueCell } from 'mol-util';
 import { RenderableSchema } from '../renderable/schema';
 import { idFactory } from 'mol-util/id-factory';
 import { Framebuffer } from './framebuffer';
+import { isWebGL2 } from './compat';
 
 const getNextTextureId = idFactory()
 
@@ -22,7 +23,8 @@ export type TextureKindValue = {
 export type TextureKind = keyof TextureKindValue
 export type TextureType = 'ubyte' | 'float'
 export type TextureFormat = 'alpha' | 'rgb' | 'rgba'
-export type TextureAttachment = 'depth' | 'stencil' | 'color0'
+/** Numbers are shortcuts for color attachment */
+export type TextureAttachment = 'depth' | 'stencil' | 'color0' | 'color1' | 'color2' | 'color3' | 'color4' | 'color5' | 'color6' | 'color7' | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7
 export type TextureFilter = 'nearest' | 'linear'
 
 export function getTarget(ctx: Context, kind: TextureKind): number {
@@ -89,8 +91,20 @@ export function getAttachment(ctx: Context, attachment: TextureAttachment): numb
     switch (attachment) {
         case 'depth': return gl.DEPTH_ATTACHMENT
         case 'stencil': return gl.STENCIL_ATTACHMENT
-        case 'color0': return gl.COLOR_ATTACHMENT0
+        case 'color0': case 0: return gl.COLOR_ATTACHMENT0
     }
+    if (isWebGL2(gl)) {
+        switch (attachment) {
+            case 'color1': case 1: return gl.COLOR_ATTACHMENT1
+            case 'color2': case 2: return gl.COLOR_ATTACHMENT2
+            case 'color3': case 3: return gl.COLOR_ATTACHMENT3
+            case 'color4': case 4: return gl.COLOR_ATTACHMENT4
+            case 'color5': case 5: return gl.COLOR_ATTACHMENT5
+            case 'color6': case 6: return gl.COLOR_ATTACHMENT6
+            case 'color7': case 7: return gl.COLOR_ATTACHMENT7
+        }
+    }
+    throw new Error('unknown texture attachment')
 }
 
 export interface Texture {
@@ -100,10 +114,13 @@ export interface Texture {
     readonly internalFormat: number
     readonly type: number
 
+    define: (x: number, y: number, z: number) => void
     load: (image: TextureImage<any>) => void
     bind: (id: TextureId) => void
     unbind: (id: TextureId) => void
-    attachFramebuffer: (framebuffer: Framebuffer, attachment: TextureAttachment) => void
+    /** Use `layer` to attach a z-slice of a 3D texture */
+    attachFramebuffer: (framebuffer: Framebuffer, attachment: TextureAttachment, layer?: number) => void
+    detachFramebuffer: (framebuffer: Framebuffer, attachment: TextureAttachment) => void
     destroy: () => void
 }
 
@@ -126,6 +143,14 @@ export function createTexture(ctx: Context, kind: TextureKind, _format: TextureF
     const internalFormat = getInternalFormat(ctx, _format, _type)
     const type = getType(ctx, _type)
 
+    gl.bindTexture(target, texture)
+    gl.texParameteri(target, gl.TEXTURE_MAG_FILTER, filter)
+    gl.texParameteri(target, gl.TEXTURE_MIN_FILTER, filter)
+    // clamp-to-edge needed for non-power-of-two textures
+    gl.texParameteri(target, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
+    gl.texParameteri(target, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
+    gl.bindTexture(target, null)
+
     let destroyed = false
     ctx.textureCount += 1
 
@@ -136,6 +161,17 @@ export function createTexture(ctx: Context, kind: TextureKind, _format: TextureF
         internalFormat,
         type,
 
+        define: (width: number, height: number, depth?: number) => {
+            gl.bindTexture(target, texture)
+            if (target === gl.TEXTURE_2D) {
+                // TODO remove cast when webgl2 types are fixed
+                (gl as WebGLRenderingContext).texImage2D(target, 0, internalFormat, width, height, 0, format, type, null)
+            } else if (target === (gl as WebGL2RenderingContext).TEXTURE_3D && depth !== undefined) {
+                (gl as WebGL2RenderingContext).texImage3D(target, 0, internalFormat, width, height, depth, 0, format, type, null)
+            } else {
+                throw new Error('unknown texture target')
+            }
+        },
         load: (data: TextureImage<any> | TextureVolume<any>) => {
             gl.bindTexture(target, texture)
             // unpack alignment of 1 since we use textures only for data
@@ -151,11 +187,6 @@ export function createTexture(ctx: Context, kind: TextureKind, _format: TextureF
             } else {
                 throw new Error('unknown texture target')
             }
-            gl.texParameteri(target, gl.TEXTURE_MAG_FILTER, filter)
-            gl.texParameteri(target, gl.TEXTURE_MIN_FILTER, filter)
-            // clamp-to-edge needed for non-power-of-two textures
-            gl.texParameteri(target, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
-            gl.texParameteri(target, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
             gl.bindTexture(target, null)
         },
         bind: (id: TextureId) => {
@@ -166,10 +197,22 @@ export function createTexture(ctx: Context, kind: TextureKind, _format: TextureF
             gl.activeTexture(gl.TEXTURE0 + id)
             gl.bindTexture(target, null)
         },
-        attachFramebuffer: (framebuffer: Framebuffer, attachment: TextureAttachment) => {
-            if (target !== gl.TEXTURE_2D) throw new Error('framebuffer texture must be 2d')
+        attachFramebuffer: (framebuffer: Framebuffer, attachment: TextureAttachment, layer?: number) => {
+            framebuffer.bind()
+            if (target === (gl as WebGL2RenderingContext).TEXTURE_3D) {
+                if (layer === undefined) throw new Error('need `layer` to attach 3D texture');
+                (gl as WebGL2RenderingContext).framebufferTextureLayer(gl.FRAMEBUFFER, getAttachment(ctx, attachment), texture, 0, layer)
+            } else {
+                gl.framebufferTexture2D(gl.FRAMEBUFFER, getAttachment(ctx, attachment), gl.TEXTURE_2D, texture, 0)
+            }
+        },
+        detachFramebuffer: (framebuffer: Framebuffer, attachment: TextureAttachment) => {
             framebuffer.bind()
-            gl.framebufferTexture2D(gl.FRAMEBUFFER, getAttachment(ctx, attachment), gl.TEXTURE_2D, texture, 0)
+            if (target === (gl as WebGL2RenderingContext).TEXTURE_3D) {
+                (gl as WebGL2RenderingContext).framebufferTextureLayer(gl.FRAMEBUFFER, getAttachment(ctx, attachment), null, 0, 0)
+            } else {
+                gl.framebufferTexture2D(gl.FRAMEBUFFER, getAttachment(ctx, attachment), gl.TEXTURE_2D, null, 0)
+            }
         },
         destroy: () => {
             if (destroyed) return

src/mol-math/geometry/gaussian-density/gpu.ts

@@ -17,6 +17,9 @@ import { RenderableState } from 'mol-gl/renderable'
 import { createRenderable, createGaussianDensityRenderObject } from 'mol-gl/render-object'
 import { createRenderTarget } from 'mol-gl/webgl/render-target'
 import { Context, createContext, getGLContext } from 'mol-gl/webgl/context';
+import { createFramebuffer } from 'mol-gl/webgl/framebuffer';
+import { createTexture, Texture, TextureAttachment } from 'mol-gl/webgl/texture';
+import { GLRenderingContext } from 'mol-gl/webgl/compat';
 
 let webglContext: Context
 function getWebGLContext() {
@@ -39,7 +42,22 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
 
     if (webgl.maxDrawBuffers > 0) {
         console.log('GaussianDensityMultiDrawBuffer')
-        return GaussianDensityMultiDrawBuffer(ctx, webgl, position, box, radius, props)
+        const { texture, scale, bbox, dim } = await GaussianDensityMultiDrawBuffer(ctx, webgl, position, box, radius, props)
+
+        console.time('gpu gaussian density 3d texture read')
+        const field = fieldFromTexture3d(webgl, texture, dim)
+        console.timeEnd('gpu gaussian density 3d texture read')
+
+        const idData = field.space.create()
+        const idField = Tensor.create(field.space, idData)
+
+        const transform = Mat4.identity()
+        Mat4.fromScaling(transform, scale)
+        Mat4.setTranslation(transform, bbox.min)
+
+        const renderTarget = createRenderTarget(webgl, dim[0], dim[1])
+
+        return { field, idField, transform, renderTarget, bbox, gridDimension: dim }
     } else {
         console.log('GaussianDensitySingleDrawBuffer')
         return GaussianDensitySingleDrawBuffer(ctx, webgl, position, box, radius, props)
@@ -77,7 +95,7 @@ async function prepareGaussianDensityData(ctx: RuntimeContext, position: Positio
     const expandedBox = Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad));
     const extent = Vec3.sub(Vec3.zero(), expandedBox.max, expandedBox.min)
 
-    const delta = getDelta(Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad)), resolution)
+    const delta = getDelta(expandedBox, resolution)
     const dim = Vec3.zero()
     Vec3.ceil(dim, Vec3.mul(dim, extent, delta))
     console.log('grid dim gpu', dim)
@@ -116,6 +134,20 @@ function getGaussianDensityRenderObject(webgl: Context, drawCount: number, posit
     return renderObject
 }
 
+//
+
+function setRenderingDefaults(gl: GLRenderingContext) {
+    gl.disable(gl.CULL_FACE)
+    gl.frontFace(gl.CCW)
+    gl.cullFace(gl.BACK)
+
+    gl.blendFunc(gl.SRC_ALPHA, gl.ONE_MINUS_SRC_ALPHA)
+    gl.blendEquation(gl.FUNC_ADD)
+    gl.enable(gl.BLEND)
+}
+
+//
+
 async function GaussianDensitySingleDrawBuffer(ctx: RuntimeContext, webgl: Context, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps): Promise<DensityData> {
     const { readSlices, smoothness } = props
 
@@ -161,14 +193,7 @@ async function GaussianDensitySingleDrawBuffer(ctx: RuntimeContext, webgl: Conte
 
     program.use()
     renderTarget.bind()
-
-    gl.disable(gl.CULL_FACE)
-    gl.frontFace(gl.CCW)
-    gl.cullFace(gl.BACK)
-
-    gl.blendFunc(gl.SRC_ALPHA, gl.ONE_MINUS_SRC_ALPHA)
-    gl.blendEquation(gl.FUNC_ADD)
-    gl.enable(gl.BLEND)
+    setRenderingDefaults(gl)
 
     const slice = new Uint8Array(dim[0] * dim[1] * 4)
 
@@ -236,38 +261,25 @@ async function GaussianDensitySingleDrawBuffer(ctx: RuntimeContext, webgl: Conte
     return { field, idField, transform, renderTarget, bbox: expandedBox, gridDimension: dim }
 }
 
-async function GaussianDensityMultiDrawBuffer(ctx: RuntimeContext, webgl: Context, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps): Promise<DensityData> {
+async function GaussianDensityMultiDrawBuffer(ctx: RuntimeContext, webgl: Context, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps) {
     const { smoothness } = props
 
     const { drawCount, positions, radii, delta, expandedBox, dim } = await prepareGaussianDensityData(ctx, position, box, radius, props)
+    const [ dx, dy, dz ] = dim
     const renderObject = getGaussianDensityRenderObject(webgl, drawCount, positions, radii, expandedBox, dim, smoothness)
     const renderable = createRenderable(webgl, renderObject)
     const drawBuffers = Math.min(8, webgl.maxDrawBuffers)
 
     //
 
-    const [ dx, dy, dz ] = dim
-
-    const space = Tensor.Space(dim, [2, 1, 0], Float32Array)
-    const data = space.create()
-    const field = Tensor.create(space, data)
-
-    const idData = space.create()
-    const idField = Tensor.create(space, idData)
-
-    //
-
     const gl = webgl.gl as WebGL2RenderingContext
     const { uCurrentSlice } = renderObject.values
 
-    const fb = gl.createFramebuffer()
-    gl.bindFramebuffer(gl.FRAMEBUFFER, fb)
+    const framebuffer = createFramebuffer(webgl)
+    framebuffer.bind()
 
-    const tex = gl.createTexture()
-    gl.bindTexture(gl.TEXTURE_3D, tex)
-    gl.texParameteri(gl.TEXTURE_3D, gl.TEXTURE_MIN_FILTER, gl.LINEAR)
-    gl.texParameteri(gl.TEXTURE_3D, gl.TEXTURE_MAG_FILTER, gl.LINEAR)
-    gl.texImage3D(gl.TEXTURE_3D, 0, gl.RGBA8, dx, dy, dz, 0, gl.RGBA, gl.UNSIGNED_BYTE, null)
+    const texture = createTexture(webgl, 'volume-uint8', 'rgba', 'ubyte', 'linear')
+    texture.define(dx, dy, dz)
 
     if (drawBuffers === 1) {
         gl.drawBuffers([
@@ -285,16 +297,7 @@ async function GaussianDensityMultiDrawBuffer(ctx: RuntimeContext, webgl: Contex
     }
 
     gl.viewport(0, 0, dx, dy)
-
-    gl.disable(gl.CULL_FACE)
-    gl.frontFace(gl.CCW)
-    gl.cullFace(gl.BACK)
-
-    gl.blendFunc(gl.SRC_ALPHA, gl.ONE_MINUS_SRC_ALPHA)
-    gl.blendEquation(gl.FUNC_ADD)
-    gl.enable(gl.BLEND)
-
-    const slice = new Uint8Array(dx * dy * 4)
+    setRenderingDefaults(gl)
 
     //
 
@@ -303,39 +306,58 @@ async function GaussianDensityMultiDrawBuffer(ctx: RuntimeContext, webgl: Contex
     const programMulti = renderable.getProgram('draw')
     programMulti.use()
 
-    console.time('gpu gaussian density 3d texture slices multi')
+    console.time('gpu gaussian density 3d texture multi')
     for (let i = 0; i < dzMulti; i += drawBuffers) {
         ValueCell.update(uCurrentSlice, i)
         for (let k = 0; k < drawBuffers; ++k) {
-            gl.framebufferTextureLayer(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0 + k, tex, 0, i + k)
+            texture.attachFramebuffer(framebuffer, k as TextureAttachment, i + k)
         }
         renderable.render('draw')
     }
-    console.timeEnd('gpu gaussian density 3d texture slices multi')
+    console.timeEnd('gpu gaussian density 3d texture multi')
 
     ValueCell.updateIfChanged(renderable.values.dDrawBuffers, 1)
     renderable.update()
     const programSingle = renderable.getProgram('draw')
     programSingle.use()
 
-    console.time('gpu gaussian density 3d texture slices single')
+    console.time('gpu gaussian density 3d texture single')
     for (let i = dzMulti; i < dz; ++i) {
         ValueCell.update(uCurrentSlice, i)
-        gl.framebufferTextureLayer(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, tex, 0, i)
+        texture.attachFramebuffer(framebuffer, 0, i)
         renderable.render('draw')
     }
-    console.timeEnd('gpu gaussian density 3d texture slices single')
+    console.timeEnd('gpu gaussian density 3d texture single')
 
-    console.time('gpu gaussian density 3d texture slices read')
-    // Must unset framebufferTextureLayer attachments before reading
+    // must detach framebuffer attachments before reading is possible
     for (let k = 0; k < drawBuffers; ++k) {
-        gl.framebufferTextureLayer(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0 + k, null, 0, 0)
+        texture.detachFramebuffer(framebuffer, k as TextureAttachment)
     }
+    framebuffer.destroy() // clean up
+
+    // throw new Error('foo')
+
+    return { texture, scale: Vec3.inverse(Vec3.zero(), delta), bbox: expandedBox, dim }
+}
+
+//
+
+function fieldFromTexture3d(ctx: Context, texture: Texture, dim: Vec3) {
+    const { gl } = ctx
+    const [ dx, dy, dz ] = dim
+
+    const space = Tensor.Space(dim, [2, 1, 0], Float32Array)
+    const data = space.create()
+    const field = Tensor.create(space, data)
+
+    const slice = new Uint8Array(dx * dy * 4)
+
+    const framebuffer = createFramebuffer(ctx)
+    framebuffer.bind()
 
     let j = 0
     for (let i = 0; i < dz; ++i) {
-        gl.framebufferTextureLayer(gl.READ_FRAMEBUFFER, gl.COLOR_ATTACHMENT0, tex, 0, i)
-        gl.readBuffer(gl.COLOR_ATTACHMENT0)
+        texture.attachFramebuffer(framebuffer, 0, i)
         gl.readPixels(0, 0, dx, dy, gl.RGBA, gl.UNSIGNED_BYTE, slice)
         for (let iy = 0; iy < dim[1]; ++iy) {
             for (let ix = 0; ix < dim[0]; ++ix) {
@@ -344,22 +366,8 @@ async function GaussianDensityMultiDrawBuffer(ctx: RuntimeContext, webgl: Contex
             }
         }
     }
-    console.timeEnd('gpu gaussian density 3d texture slices read')
-
-    // clean up
-    gl.bindFramebuffer(gl.FRAMEBUFFER, null)
-
-    //
 
-    const transform = Mat4.identity()
-    Mat4.fromScaling(transform, Vec3.inverse(Vec3.zero(), delta))
-    Mat4.setTranslation(transform, expandedBox.min)
-
-    // throw new Error('foo')
-
-    const renderTarget = createRenderTarget(webgl, dx, dy)
-
-    return { field, idField, transform, renderTarget, bbox: expandedBox, gridDimension: dim }
-}
+    framebuffer.destroy()
 
-// const wait = (ms: number) => new Promise(r => setTimeout(r, ms))
+    return field
+}