import { AST_TOKEN_TYPES , DefineValuePlaceholder , GLSLContext , IComputeModel , IConfigService , IDENTIFIER , IRendererService , IShaderModuleService , ISystem , KernelBundle , STORAGE_CLASS , Target , createEntity , } from '@antv/g-webgpu-core' ; import * as WebGPUConstants from '@webgpu/types/dist/constants' ; import { inject , injectable } from 'inversify' ; import { isArray , isNil , isNumber , isTypedArray } from 'lodash' ; import { createCanvas } from './utils/canvas' ; @ injectable ( ) export class Kernel { @ inject ( IDENTIFIER . RenderEngine ) private readonly engine : IRendererService ; @ inject ( IDENTIFIER . ConfigService ) private readonly configService : IConfigService ; private entity = createEntity ( ) ; private model : IComputeModel ; private dirty : boolean = true ; private compiledBundle : KernelBundle ; private initPromise : Promise < void > ; public init ( ) { const { canvas , engineOptions } = this . configService . get ( ) ; this . initPromise = this . engine . init ( { canvas : canvas || createCanvas ( ) , swapChainFormat : WebGPUConstants . TextureFormat . BGRA8Unorm , antialiasing : false , ... engineOptions , } ) ; } public setBundle ( bundle ) { this . compiledBundle = JSON . parse ( JSON . stringify ( bundle ) ) ; } public setDispatch ( dispatch ) { if ( this . compiledBundle . context ) { this . compiledBundle . context . dispatch = dispatch ; } return this ; } public setMaxIteration ( maxIteration ) { if ( this . compiledBundle . context ) { this . compiledBundle . context . maxIteration = maxIteration ; } return this ; } public setBinding ( name , data ? ) { if ( typeof name === 'string' ) { const isNumberLikeData = isNumber ( data ) || isTypedArray ( data ) || isArray ( data ) ; if ( this . compiledBundle && this . compiledBundle . context ) { const existedDefine = this . compiledBundle . context . defines . find ( ( b ) => b . name === name ) ; if ( existedDefine ) { existedDefine . value = data as number ; return this ; } const existedBinding = this . compiledBundle . context . uniforms . find ( ( b ) => b . name === name ) ; if ( existedBinding ) { if ( isNumberLikeData ) { existedBinding . data = data ; existedBinding . isReferer = false ; if ( existedBinding . storageClass === STORAGE_CLASS . Uniform ) { if ( this . model ) { this . model . updateUniform ( name , data ) ; } } else { if ( this . model ) { this . model . updateBuffer ( name , data ) ; } } } else { existedBinding . isReferer = true ; existedBinding . data = data as Kernel ; } } } } else { Object . keys ( name ) . forEach ( ( key ) => { this . setBinding ( key , name [ key ] ) ; } ) ; } return this ; } public async execute ( iteration = 1 ) { if ( this . dirty ) { if ( this . compiledBundle . context ) { if ( iteration > 1 ) { this . compiledBundle . context . maxIteration = iteration ; } else { this . compiledBundle . context . maxIteration ++ ; } } await this . compile ( ) ; this . dirty = false ; } this . engine . beginFrame ( ) ; this . engine . clear ( { } ) ; if ( this . compiledBundle . context ) { this . compiledBundle . context . uniforms . filter ( ( { isReferer } ) => isReferer ) . forEach ( ( { data , name } ) => { this . model . confirmInput ( ( data as Kernel ) . model , name ) ; } ) ; } for ( let i = 0 ; i < iteration ; i ++ ) { this . model . run ( ) ; } this . engine . endFrame ( ) ; return this ; } /**
     * read output from GPUBuffer
     */ public async getOutput ( ) { return this . model . readData ( ) ; } private async compile ( ) { await this . initPromise ; const context = { ... this . compiledBundle . context ! , } ; const target = this . engine . supportWebGPU ? this . engine . useWGSL ? Target . WGSL : Target . GLSL450 : Target . GLSL100 ; let shader = this . compiledBundle . shaders [ target ] ; context . defines . filter ( ( define ) => define . runtime ) . forEach ( ( define ) => { const valuePlaceHolder = `${ DefineValuePlaceholder }${ define . name }` ; shader = shader . replace ( valuePlaceHolder , `${ define . value }` ) ; } ) ; context . shader = shader ; context . uniforms . forEach ( ( uniform ) => { if ( ! uniform . data ) { if ( uniform . storageClass === STORAGE_CLASS . StorageBuffer ) { let sizePerElement = 1 ; if ( uniform . type === AST_TOKEN_TYPES . FloatArray ) { sizePerElement = 1 ; } else if ( uniform . type === AST_TOKEN_TYPES . Vector4FloatArray ) { sizePerElement = 4 ; } uniform . data = new Float32Array ( context . output . length ! * sizePerElement ) . fill ( 0 ) ; } } } ) ; this . compiledBundle . context = context ; this . model = await this . engine . createComputeModel ( this . compiledBundle . context ) ; } }