Skip to content
Snippets Groups Projects
gpu.ts 11.61 KiB
/**
 * Copyright (c) 2017-2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
 *
 * @author Alexander Rose <alexander.rose@weirdbyte.de>
 * @author Michael Krone <michael.krone@uni-tuebingen.de>
 */

import { RuntimeContext } from 'mol-task'
import { PositionData, DensityData } from '../common'
import { Box3D } from '../../geometry'
import { GaussianDensityProps, getDelta } from '../gaussian-density'
import { OrderedSet } from 'mol-data/int'
import { Vec3, Tensor, Mat4 } from '../../linear-algebra'
import { GaussianDensityValues } from 'mol-gl/renderable/gaussian-density'
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, 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> {
    // 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)

    const positions = new Float32Array(n * 3)
    const radii = new Float32Array(n)

    let maxRadius = 0

    for (let i = 0; i < n; ++i) {
        const j = OrderedSet.getAt(indices, i);

        positions[i * 3] = x[j]
        positions[i * 3 + 1] = y[j]
        positions[i * 3 + 2] = z[j]
        const r = radius(j) + radiusOffset
        if (maxRadius < r) maxRadius = r
        radii[i] = r

        if (i % 10000 === 0 && ctx.shouldUpdate) {
            await ctx.update({ message: 'preparing density data', current: i, max: n })
        }
    }

    
    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)

    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 gpu', dim)

    return { drawCount: n, positions, radii, delta, expandedBox, dim }
}

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 = {
        drawCount: ValueCell.create(drawCount),
        instanceCount: ValueCell.create(1),

        aRadius: ValueCell.create(radii),
        aPosition: ValueCell.create(positions),

        uCurrentSlice: ValueCell.create(0),
        uCurrentX: ValueCell.create(0),
        uCurrentY: ValueCell.create(0),
        uBboxMin: ValueCell.create(box.min),
        uBboxMax: ValueCell.create(box.max),
        uBboxSize: ValueCell.create(extent),
        uGridDim: ValueCell.create(dimensions),
        uAlpha: ValueCell.create(smoothness),

        dDrawBuffers: ValueCell.create(0),
    }
    const state: RenderableState = {
        visible: true,
        depthMask: false
    }

    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)

    //

    // TODO fallback to lower resolution when texture size is not large enough
    const maxTexSize = webgl.maxTextureSize
    let fboTexDimX = 0
    let fboTexDimY = dim[1]
    let fboTexRows = 1
    let fboTexCols = dim[0]
    if (maxTexSize < dim[0] * dim[2]) {
        fboTexCols =  Math.floor(maxTexSize / dim[0])
        fboTexRows = Math.ceil(dim[2] / fboTexCols)
        fboTexDimX = fboTexCols * dim[0]
        fboTexDimY *= fboTexRows
    } else {
        fboTexDimX = dim[0] * dim[2]
    }

    console.log('dim', dim, 'cols', fboTexCols, 'rows', fboTexRows)

    //

    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
    const { uCurrentSlice, uCurrentX, uCurrentY } = renderObject.values

    const program = renderable.getProgram('draw')
    const renderTarget = createRenderTarget(webgl, fboTexDimX, fboTexDimY)

    program.use()
    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(dim[0] * dim[1] * 4)

    console.time('gpu gaussian density slices')
    let currCol = 0
    let currY = 0
    let currX = 0
    let j = 0
    for (let i = 0; i < dim[2]; ++i) {
        if (currCol >= fboTexCols) {
            currCol -= fboTexCols
            currY += dim[1]
            currX = 0
        }
        gl.viewport(currX, currY, dim[0], dim[1])
        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)
            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
                }
            }
        }
        ++currCol
        currX += dim[0]
    }
    console.timeEnd('gpu gaussian density slices')

    //

    if (!readSlices) {
        console.time('gpu gaussian density full')
        renderTarget.getBuffer()
        const { array } = renderTarget.image
        let idx = 0
        let tmpCol = 0
        let tmpRow = 0
        for (let iz = 0; iz < dim[2]; ++iz) {
            if (tmpCol >= fboTexCols ) {
                tmpCol = 0
                tmpRow += dim[1]
            }
            for (let iy = 0; iy < dim[1]; ++iy) {
                for (let ix = 0; ix < dim[0]; ++ix) {
                    data[idx] = array[4 * (tmpCol * dim[0] + (iy + tmpRow) * fboTexDimX + ix)] / 255
                    idx++
                }
            }
            tmpCol++
        }
        console.timeEnd('gpu gaussian density full')
    }

    //

    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 }
}

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 }
}