import { AST_TOKEN_TYPES , createEntity , DataType , GLSLContext , IComputeModel , STORAGE_CLASS , } from '@antv/g-webgpu-core' ; import { isTypedArray } from 'lodash' ; import regl from 'regl' ; import quadVert from './shaders/quad.vert.glsl' ; interface DataTextureDescriptor { id : number ; data : number | number [ ] | Float32Array | Uint8Array | Uint16Array | Uint32Array | Int8Array | Int16Array | Int32Array | undefined ; textureWidth : number ; texture : regl . Texture2D ; texelCount : number ; originalDataLength : number ; elementsPerTexel : number ; typedArrayConstructor ? : Function ; isOutput : boolean ; } let textureId = 0 ; const debug = false ; /**
 * adaptor for regl.DrawCommand
 */ export default class ReglComputeModel implements IComputeModel { private entity = createEntity ( ) ; private texFBO : regl . Framebuffer2D ; private computeCommand : regl . DrawCommand ; private textureCache : { [ textureName ] : DataTextureDescriptor ; } = { } ; private outputTextureName : string ; private swapOutputTextureName : string ; private compiledPingpong : boolean ; private dynamicPingpong : boolean ; constructor ( private reGl , private context ) { const uniforms = { } ; this . context . uniforms . forEach ( ( uniform ) => { const { name , type , data , isReferer , storageClass } = uniform ; if ( storageClass === STORAGE_CLASS . StorageBuffer ) { if ( ! isReferer ) { this . textureCache [ name ] = this . calcDataTexture ( name , type , data ! ) ; const { textureWidth : width , isOutput } = this . textureCache [ name ] ; uniforms [ `${ name }Size` ] = [ width , width ] ; if ( isOutput ) { this . outputTextureName = name ; if ( this . context . needPingpong ) { this . outputTextureName = `${ name }Output` ; this . textureCache [ this . outputTextureName ] = this . calcDataTexture ( name , type , data ! ) ; } } } else { this . textureCache [ name ] = { data : undefined , } ; uniforms [ `${ name }Size` ] = ( ) => data . compiledBundle . context . output . textureSize ; } uniforms [ name ] = ( ) => { if ( debug ) { console . log ( `[${ this . entity }]: ${ name } ${ this . textureCache [ name ] . id }` ) ; } return this . textureCache [ name ] . texture ; } ; } else if ( storageClass === STORAGE_CLASS . Uniform ) { if ( data && ( Array . isArray ( data ) || isTypedArray ( data ) ) && ( data as ArrayLike < number > ) . length > 16 ) { throw new Error ( `invalid data type ${ type }` ) ; } uniforms [ name ] = ( ) => uniform . data ; } } ) ; const { textureWidth , texelCount } = this . getOuputDataTexture ( ) ; uniforms . u_OutputTextureSize = [ textureWidth , textureWidth ] ; uniforms . u_OutputTexelCount = texelCount ; this . context . output . textureSize = [ textureWidth ! , textureWidth ! ] ; const drawParams = { attributes : { a_Position : [ [ - 1 , 1 , 0 ] , [ - 1 , - 1 , 0 ] , [ 1 , 1 , 0 ] , [ 1 , - 1 , 0 ] , ] , a_TexCoord : [ [ 0 , 1 ] , [ 0 , 0 ] , [ 1 , 1 ] , [ 1 , 0 ] , ] , } , frag : `#ifdef GL_FRAGMENT_PRECISION_HIGH
  precision highp float;
#else
  precision mediump float;
#endif
${ this . context . shader }` , uniforms , vert : quadVert , primitive : 'triangle strip' , count : 4 , } ; this . computeCommand = this . reGl ( drawParams ) ; } public run ( ) { if ( this . context . maxIteration > 1 && this . context . needPingpong ) { this . compiledPingpong = true ; } if ( this . compiledPingpong || this . dynamicPingpong ) { this . swap ( ) ; } this . texFBO = this . reGl . framebuffer ( { color : this . getOuputDataTexture ( ) . texture , } ) ; this . texFBO . use ( ( ) => { this . computeCommand ( ) ; } ) ; if ( debug ) { console . log ( `[${ this . entity }]: output ${ this . getOuputDataTexture ( ) . id }` ) ; } } public async readData ( ) { let pixels ; this . reGl ( { framebuffer : this . texFBO , } ) ( ( ) => { pixels = this . reGl . read ( ) ; } ) ; if ( pixels ) { const { originalDataLength , elementsPerTexel , typedArrayConstructor = Float32Array , } = this . getOuputDataTexture ( ) ; let formattedPixels = [ ] ; if ( elementsPerTexel !== 4 ) { for ( let i = 0 ; i < pixels . length ; i += 4 ) { if ( elementsPerTexel === 1 ) { formattedPixels . push ( pixels [ i ] ) ; } else if ( elementsPerTexel === 2 ) { formattedPixels . push ( pixels [ i ] , pixels [ i + 1 ] ) ; } else { formattedPixels . push ( pixels [ i ] , pixels [ i + 1 ] , pixels [ i + 2 ] ) ; } } } else { formattedPixels = pixels ; } return new typedArrayConstructor ( formattedPixels . slice ( 0 , originalDataLength ) ) ; } return new Float32Array ( ) ; } public confirmInput ( model , inputName ) { let inputModel ; if ( this . entity === ( model as ReglComputeModel ) . entity ) { this . dynamicPingpong = true ; inputModel = this ; } else { inputModel = model as ReglComputeModel ; } this . textureCache [ inputName ] . id = inputModel . getOuputDataTexture ( ) . id ; this . textureCache [ inputName ] . texture = inputModel . getOuputDataTexture ( ) . texture ; if ( debug ) { console . log ( `[${ this . entity }]: confirm input ${ inputName } from model ${ inputModel . entity }, ${ ( inputModel as ReglComputeModel ) . getOuputDataTexture ( ) . id }` ) ; } } public updateUniform ( ) { } public updateBuffer ( bufferName , data , offset = 0 ) { const buffer = this . context . uniforms . find ( ( { name } ) => name === bufferName ) ; if ( buffer ) { const { texture , data : paddingData } = this . calcDataTexture ( bufferName , buffer . type , data ) ; this . textureCache [ bufferName ] . data = paddingData ; this . textureCache [ bufferName ] . texture = texture ; } } public destroy ( ) { } private swap ( ) { if ( ! this . swapOutputTextureName ) { this . createSwapOutputDataTexture ( ) ; } if ( this . compiledPingpong ) { const outputTextureUniformName = this . context . output . name ; this . textureCache [ outputTextureUniformName ] . id = this . getOuputDataTexture ( ) . id ; this . textureCache [ outputTextureUniformName ] . texture = this . getOuputDataTexture ( ) . texture ; } const tmp = this . outputTextureName ; this . outputTextureName = this . swapOutputTextureName ; this . swapOutputTextureName = tmp ; if ( debug ) { console . log ( `[${ this . entity }]: after swap, output ${ this . getOuputDataTexture ( ) . id }` ) ; } } private getOuputDataTexture ( ) { return this . textureCache [ this . outputTextureName ] ; } private createSwapOutputDataTexture ( ) { const texture = this . cloneDataTexture ( this . getOuputDataTexture ( ) ) ; this . swapOutputTextureName = `${ this . entity }-swap` ; this . textureCache [ this . swapOutputTextureName ] = texture ; } private cloneDataTexture ( texture ) { const { data , textureWidth } = texture ; return { ... texture , id : textureId ++ , texture : this . reGl . texture ( { width : textureWidth , height : textureWidth , data , type : 'float' , } ) , } ; } private calcDataTexture ( name , type , data ) { let elementsPerTexel = 1 ; if ( type === AST_TOKEN_TYPES . Vector4FloatArray ) { elementsPerTexel = 4 ; } const paddingData = [ ] ; for ( let i = 0 ; i < ( data as number [ ] ) . length ; i += elementsPerTexel ) { if ( elementsPerTexel === 1 ) { paddingData . push ( ( data as number [ ] ) [ i ] , 0 , 0 , 0 ) ; } else if ( elementsPerTexel === 2 ) { paddingData . push ( ( data as number [ ] ) [ i ] , ( data as number [ ] ) [ i + 1 ] , 0 , 0 ) ; } else if ( elementsPerTexel === 3 ) { paddingData . push ( ( data as number [ ] ) [ i ] , ( data as number [ ] ) [ i + 1 ] , ( data as number [ ] ) [ i + 2 ] , 0 ) ; } else if ( elementsPerTexel === 4 ) { paddingData . push ( ( data as number [ ] ) [ i ] , ( data as number [ ] ) [ i + 1 ] , ( data as number [ ] ) [ i + 2 ] , ( data as number [ ] ) [ i + 3 ] ) ; } } const originalDataLength = ( data as ArrayLike < number > ) . length ; const texelCount = Math . ceil ( originalDataLength / elementsPerTexel ) ; const width = Math . ceil ( Math . sqrt ( texelCount ) ) ; const paddingTexelCount = width * width ; if ( texelCount < paddingTexelCount ) { paddingData . push ( ... new Array ( ( paddingTexelCount - texelCount ) * 4 ) . fill ( 0 ) ) ; } const texture = this . reGl . texture ( { width , height : width , data : paddingData , type : 'float' , } ) ; return { id : textureId ++ , data : paddingData , originalDataLength , typedArrayConstructor : isTypedArray ( data ) ? data ! . constructor : undefined , textureWidth : width , texture , texelCount , elementsPerTexel , isOutput : name === this . context . output . name , } ; } }