wip, mrt

src/mol-gl/renderable/gaussian-density.ts

@@ -11,8 +11,6 @@ import { AttributeSpec, Values, UniformSpec, ValueSpec, DefineSpec } from './sch
 import { GaussianDensityShaderCode } from '../shader-code';
 
 export const GaussianDensitySchema = {
-    dWebGL2: DefineSpec('boolean'),
-
     drawCount: ValueSpec('number'),
     instanceCount: ValueSpec('number'),
 
@@ -27,6 +25,8 @@ export const GaussianDensitySchema = {
     uBboxSize: UniformSpec('v3'),
     uGridDim: UniformSpec('v3'),
     uAlpha: UniformSpec('f'),
+
+    dDrawBuffers: DefineSpec('number'),
 }
 export type GaussianDensitySchema = typeof GaussianDensitySchema
 export type GaussianDensityValues = Values<GaussianDensitySchema>

src/mol-gl/shader-code.ts

@@ -59,13 +59,13 @@ function getDefinesCode (defines: ShaderDefines) {
     for (const name in defines) {
         const define = defines[name]
         const v = define.ref.value
-        if (v) {
+        if (v !== undefined) {
             if (typeof v === 'string') {
                 lines.push(`#define ${name}_${v}`)
             } else if (typeof v === 'number') {
                 lines.push(`#define ${name} ${v}`)
             } else if (typeof v === 'boolean') {
-                lines.push(`#define ${name}`)
+                if (v) lines.push(`#define ${name}`)
             } else {
                 throw new Error('unknown define type')
             }
@@ -85,7 +85,7 @@ const glsl300VertPrefix = `#version 300 es
 
 const glsl300FragPrefix = `#version 300 es
 #define varying in
-out highp vec4 out_FragColor;
+layout(location = 0) out highp vec4 out_FragColor;
 #define gl_FragColor out_FragColor
 #define gl_FragDepthEXT gl_FragDepth
 #define texture2D texture

src/mol-gl/shader/gaussian-density.frag

@@ -19,16 +19,33 @@ uniform float uCurrentX;
 uniform float uCurrentY;
 uniform float uAlpha;
 
-void main() {
-    vec3 tmpVec = gl_FragCoord.xyz;
-    tmpVec.x = tmpVec.x - uCurrentX;
-    tmpVec.y = tmpVec.y - uCurrentY;
-    vec3 fragPos = vec3(
-        (tmpVec.x - 0.5) / uGridDim.x,
-        (tmpVec.y - 0.5) / uGridDim.y,
-        (uCurrentSlice) / uGridDim.z
-    );
+// #if dDrawBuffers == 8
+    // layout(location = 0) out vec4 out0;
+    layout(location = 1) out vec4 out1;
+    layout(location = 2) out vec4 out2;
+    layout(location = 3) out vec4 out3;
+    layout(location = 4) out vec4 out4;
+    layout(location = 5) out vec4 out5;
+    layout(location = 6) out vec4 out6;
+    layout(location = 7) out vec4 out7;
+// #endif
+
+float calcDensity(float x, float y, float z) {
+    vec3 fragPos = vec3(x, y, z) / uGridDim;
     float dist = length(fragPos * uBboxSize - position * uBboxSize);
     float density = exp(-uAlpha * ((dist * dist) / (radius * radius)));
-    gl_FragColor = vec4(1, 1, 1, density);
+    return density;
+}
+
+void main() {
+    vec2 tmpVec = gl_FragCoord.xy - vec2(uCurrentX, uCurrentY) - 0.5;
+    out_FragColor = vec4(1, 1, 1, calcDensity(tmpVec.x, tmpVec.y, uCurrentSlice + 0.0));
+    out1 = vec4(1, 1, 1, calcDensity(tmpVec.x, tmpVec.y, uCurrentSlice + 1.0));
+    out2 = vec4(1, 1, 1, calcDensity(tmpVec.x, tmpVec.y, uCurrentSlice + 2.0));
+    out3 = vec4(1, 1, 1, calcDensity(tmpVec.x, tmpVec.y, uCurrentSlice + 3.0));
+    out4 = vec4(1, 1, 1, calcDensity(tmpVec.x, tmpVec.y, uCurrentSlice + 4.0));
+    out5 = vec4(1, 1, 1, calcDensity(tmpVec.x, tmpVec.y, uCurrentSlice + 5.0));
+    out6 = vec4(1, 1, 1, calcDensity(tmpVec.x, tmpVec.y, uCurrentSlice + 6.0));
+    out7 = vec4(1, 1, 1, calcDensity(tmpVec.x, tmpVec.y, uCurrentSlice + 7.0));
+    // gl_FragColor = vec4(1, 1, 1, calcDensity(tmpVec.x, tmpVec.y, uCurrentSlice));
 }

src/mol-gl/webgl/context.ts

@@ -96,6 +96,7 @@ export interface Context {
     instancedDrawCount: number
 
     readonly maxTextureSize: number
+    readonly maxDrawBuffers: number
 
     unbindFramebuffer: () => void
     readPixels: (x: number, y: number, width: number, height: number, buffer: Uint8Array) => void
@@ -134,7 +135,8 @@ export function createContext(gl: GLRenderingContext): Context {
     const programCache = createProgramCache()
 
     const parameters = {
-        maxTextureSize: gl.getParameter(gl.MAX_TEXTURE_SIZE)
+        maxTextureSize: gl.getParameter(gl.MAX_TEXTURE_SIZE),
+        maxDrawBuffers: isWebGL2(gl) ? gl.getParameter(gl.MAX_DRAW_BUFFERS) : 0,
     }
 
     return {
@@ -164,6 +166,7 @@ export function createContext(gl: GLRenderingContext): Context {
         instancedDrawCount: 0,
 
         get maxTextureSize () { return parameters.maxTextureSize },
+        get maxDrawBuffers () { return parameters.maxDrawBuffers },
 
         unbindFramebuffer: () => unbindFramebuffer(gl),
         readPixels: (x: number, y: number, width: number, height: number, buffer: Uint8Array) => {

src/mol-gl/webgl/render-target.ts

@@ -69,7 +69,7 @@ export function createRenderTarget (ctx: Context, _width: number, _height: numbe
 
         bind: () => {
             framebuffer.bind()
-            gl.viewport(0, 0, _width, _height);
+            gl.viewport(0, 0, _width, _height)
         },
         setSize: (width: number, height: number) => {
             _width = width

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

@@ -16,10 +16,38 @@ import { ValueCell } from 'mol-util'
 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 } from 'mol-gl/webgl/context';
+import { Context, createContext, getGLContext } from 'mol-gl/webgl/context';
+
+let webglContext: Context
+function getWebGLContext() {
+    if (webglContext) return webglContext
+    const canvas = document.createElement('canvas')
+    const gl = getGLContext(canvas, {
+        alpha: true,
+        antialias: false,
+        depth: false,
+        preserveDrawingBuffer: true
+    })
+    if (!gl) throw new Error('Could not create a WebGL rendering context')
+    webglContext = createContext(gl)
+    return webglContext
+}
 
 export async function GaussianDensityGPU(ctx: RuntimeContext, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps): Promise<DensityData> {
-    const { resolution, radiusOffset, smoothness, readSlices } = props
+    // TODO allow passing a context via props
+    const webgl = getWebGLContext()
+
+    if (webgl.maxDrawBuffers > 0) {
+        console.log('GaussianDensityMultiDrawBuffer')
+        return GaussianDensityMultiDrawBuffer(ctx, webgl, position, box, radius, props)
+    } else {
+        console.log('GaussianDensitySingleDrawBuffer')
+        return GaussianDensitySingleDrawBuffer(ctx, webgl, position, box, radius, props)
+    }
+}
+
+async function prepareGaussianDensityData(ctx: RuntimeContext, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps) {
+    const { resolution, radiusOffset } = props
 
     const { indices, x, y, z } = position
     const n = OrderedSet.size(indices)
@@ -44,6 +72,7 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
         }
     }
 
+    
     const pad = maxRadius * 2 + resolution
     const expandedBox = Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad));
     const extent = Vec3.sub(Vec3.zero(), expandedBox.max, expandedBox.min)
@@ -51,22 +80,16 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     const delta = getDelta(Box3D.expand(Box3D.empty(), box, Vec3.create(pad, pad, pad)), resolution)
     const dim = Vec3.zero()
     Vec3.ceil(dim, Vec3.mul(dim, extent, delta))
-    console.log('grid dim', dim)
-
-    const _r2 = maxRadius * 2
-    const _radius2 = Vec3.create(_r2, _r2, _r2)
-    Vec3.mul(_radius2, _radius2, delta)
-    const updateChunk = Math.ceil(10000 / (_radius2[0] * _radius2[1] * _radius2[2]))
+    console.log('grid dim gpu', dim)
 
-    //
+    return { drawCount: n, positions, radii, delta, expandedBox, dim }
+}
 
-    // TODO do in OffscreenCanvas (https://www.chromestatus.com/feature/5681560598609920)?
-    const webgl = getWebGLContext()
+function getGaussianDensityRenderObject(drawCount: number, positions: Float32Array, radii: Float32Array, box: Box3D, dimensions: Vec3, smoothness: number) {
+    const extent = Vec3.sub(Vec3.zero(), box.max, box.min)
 
     const values: GaussianDensityValues = {
-        dWebGL2: ValueCell.create(webgl.isWebGL2),
-
-        drawCount: ValueCell.create(n),
+        drawCount: ValueCell.create(drawCount),
         instanceCount: ValueCell.create(1),
 
         aRadius: ValueCell.create(radii),
@@ -75,11 +98,13 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
         uCurrentSlice: ValueCell.create(0),
         uCurrentX: ValueCell.create(0),
         uCurrentY: ValueCell.create(0),
-        uBboxMin: ValueCell.create(expandedBox.min),
-        uBboxMax: ValueCell.create(expandedBox.max),
+        uBboxMin: ValueCell.create(box.min),
+        uBboxMax: ValueCell.create(box.max),
         uBboxSize: ValueCell.create(extent),
-        uGridDim: ValueCell.create(dim),
+        uGridDim: ValueCell.create(dimensions),
         uAlpha: ValueCell.create(smoothness),
+
+        dDrawBuffers: ValueCell.create(0),
     }
     const state: RenderableState = {
         visible: true,
@@ -87,6 +112,15 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     }
 
     const renderObject = createGaussianDensityRenderObject(values, state)
+
+    return renderObject
+}
+
+async function GaussianDensitySingleDrawBuffer(ctx: RuntimeContext, webgl: Context, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps): Promise<DensityData> {
+    const { readSlices, smoothness } = props
+
+    const { drawCount, positions, radii, delta, expandedBox, dim } = await prepareGaussianDensityData(ctx, position, box, radius, props)
+    const renderObject = getGaussianDensityRenderObject(drawCount, positions, radii, expandedBox, dim, smoothness)
     const renderable = createRenderable(webgl, renderObject)
 
     //
@@ -120,6 +154,7 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     //
 
     const { gl } = webgl
+    const { uCurrentSlice, uCurrentX, uCurrentY } = renderObject.values
 
     const program = renderable.getProgram('draw')
     const renderTarget = createRenderTarget(webgl, fboTexDimX, fboTexDimY)
@@ -154,9 +189,9 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
             currX = 0
         }
         gl.viewport(currX, currY, dim[0], dim[1])
-        ValueCell.update(values.uCurrentSlice, i)
-        ValueCell.update(values.uCurrentX, currX)
-        ValueCell.update(values.uCurrentY, currY)
+        ValueCell.update(uCurrentSlice, i)
+        ValueCell.update(uCurrentX, currX)
+        ValueCell.update(uCurrentY, currY)
         renderable.render('draw')
         if (readSlices) {
             renderTarget.readBuffer(currX, currY, dim[0], dim[1], slice)
@@ -169,10 +204,6 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
         }
         ++currCol
         currX += dim[0]
-
-        if (i % updateChunk === 0 && ctx.shouldUpdate) {
-            await ctx.update({ message: 'filling density grid', current: i, max: n })
-        }
     }
     console.timeEnd('gpu gaussian density slices')
 
@@ -210,17 +241,105 @@ export async function GaussianDensityGPU(ctx: RuntimeContext, position: Position
     return { field, idField, transform, renderTarget, bbox: expandedBox, gridDimension: dim }
 }
 
-let webglContext: Context
-function getWebGLContext() {
-    if (webglContext) return webglContext
-    const canvas = document.createElement('canvas')
-    const gl = canvas.getContext('webgl', {
-        alpha: true,
-        antialias: false,
-        depth: false,
-        preserveDrawingBuffer: true
-    })
-    if (!gl) throw new Error('Could not create a WebGL rendering context')
-    webglContext = createContext(gl)
-    return webglContext
+async function GaussianDensityMultiDrawBuffer(ctx: RuntimeContext, webgl: Context, position: PositionData, box: Box3D, radius: (index: number) => number,  props: GaussianDensityProps): Promise<DensityData> {
+    const { smoothness } = props
+
+    const { drawCount, positions, radii, delta, expandedBox, dim } = await prepareGaussianDensityData(ctx, position, box, radius, props)
+    const renderObject = getGaussianDensityRenderObject(drawCount, positions, radii, expandedBox, dim, smoothness)
+    const renderable = createRenderable(webgl, renderObject)
+
+    //
+
+    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)
+
+    //
+
+    console.log('webgl.maxDrawBuffers', webgl.maxDrawBuffers)
+
+    const gl = webgl.gl as WebGL2RenderingContext
+    const { uCurrentSlice } = renderObject.values
+
+    const program = renderable.getProgram('draw')
+    const renderTarget = createRenderTarget(webgl, dx, dy)
+
+    const fb = gl.createFramebuffer()
+    gl.bindFramebuffer(gl.FRAMEBUFFER, fb)
+
+    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)
+
+    gl.drawBuffers([
+        gl.COLOR_ATTACHMENT0,
+        gl.COLOR_ATTACHMENT1,
+        gl.COLOR_ATTACHMENT2,
+        gl.COLOR_ATTACHMENT3,
+        gl.COLOR_ATTACHMENT4,
+        gl.COLOR_ATTACHMENT5,
+        gl.COLOR_ATTACHMENT6,
+        gl.COLOR_ATTACHMENT7,
+    ]);
+
+    program.use()
+    gl.viewport(0, 0, dx, dy)
+
+    // renderTarget.bind()
+
+    gl.disable(gl.CULL_FACE)
+    gl.frontFace(gl.CCW)
+    gl.cullFace(gl.BACK)
+
+    gl.depthMask(true)
+    gl.clearColor(0, 0, 0, 0)
+    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT)
+    gl.depthMask(false)
+
+    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)
+
+    const dzFactor = Math.floor(dz / webgl.maxDrawBuffers)
+    const dzFull = dzFactor * webgl.maxDrawBuffers
+    const dzRest = dz - dzFull
+    console.log(dz, webgl.maxDrawBuffers, dzFull, dzRest)
+
+    console.time('gpu gaussian density 3d texture slices')
+    let j = 0
+    for (let i = 0; i < dz; i += 8) {
+        ValueCell.update(uCurrentSlice, i)
+        for (let k = 0; k < webgl.maxDrawBuffers && k + i < dz; ++k) {
+            gl.framebufferTextureLayer(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0 + k, tex, 0, i + k)
+        }
+        renderable.render('draw')
+        for (let k = 0; k < webgl.maxDrawBuffers && k + i < dz; ++k) {
+            gl.readBuffer(gl.COLOR_ATTACHMENT0 + k);
+            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) {
+                    data[j] = slice[4 * (iy * dim[0] + ix)] / 255
+                    ++j
+                }
+            }
+        }
+    }
+    console.timeEnd('gpu gaussian density 3d texture slices')
+
+    //
+
+    const transform = Mat4.identity()
+    Mat4.fromScaling(transform, Vec3.inverse(Vec3.zero(), delta))
+    Mat4.setTranslation(transform, expandedBox.min)
+
+    return { field, idField, transform, renderTarget, bbox: expandedBox, gridDimension: dim }
 }