import { gl , IModel , IModelDrawOptions , IModelInitializationOptions , isSafari , IUniform , } from '@antv/g-webgpu-core' ; import * as WebGPUConstants from '@webgpu/types/dist/constants' ; import isNil from 'lodash/isNil' ; import { WebGPUEngine } from '.' ; import { extractUniforms } from '../utils/uniform' ; import { getColorStateDescriptors , getCullMode , getDepthStencilStateDescriptor , primitiveMap , } from './constants' ; import WebGPUAttribute from './WebGPUAttribute' ; import WebGPUBuffer from './WebGPUBuffer' ; import WebGPUElements from './WebGPUElements' ; import WebGPUFramebuffer from './WebGPUFramebuffer' ; import WebGPUTexture2D from './WebGPUTexture2D' ; function concatenate ( resultConstructor , ... arrays ) { let totalLength = 0 ; for ( const arr of arrays ) { totalLength += arr . length ; } const result = new resultConstructor ( totalLength ) ; let offset = 0 ; for ( const arr of arrays ) { result . set ( arr , offset ) ; offset += arr . length ; } return result ; } export default class WebGPUModel implements IModel { private pipelineLayout : GPUPipelineLayout ; private renderPipeline : GPURenderPipeline ; private uniformsBindGroupLayout : GPUBindGroupLayout ; private uniformBindGroup : GPUBindGroup ; private uniformBuffer : WebGPUBuffer ; private uniforms : { [ key ] : IUniform ; } = { } ; /**
     * 用于后续渲染时动态更新
     */ private uniformGPUBufferLayout : Array < { name : string ; offset : number ; } > = [ ] ; /**
     * vertex
     */ private attributeCache : { [ attributeName ] : WebGPUAttribute ; } = { } ; /**
     * indices's buffer
     */ private indexBuffer : WebGPUBuffer ; private indexCount : number ; constructor ( private engine , private options ) { } public async init ( ) { const { vs , fs , attributes , uniforms , primitive , count , elements , depth , blend , stencil , cull , instances , } = this . options ; const { vertexStage , fragmentStage , } = await this . compilePipelineStageDescriptor ( vs , fs , null ) ; if ( uniforms ) { this . buildUniformBindGroup ( uniforms ) ; } if ( elements ) { this . indexBuffer = ( elements as WebGPUElements ) . get ( ) as WebGPUBuffer ; this . indexCount = ( elements as WebGPUElements ) . indexCount ; } const vertexState = { vertexBuffers : Object . keys ( attributes ) . map ( ( attributeName , i ) => { const attribute = attributes [ attributeName ] as WebGPUAttribute ; const { arrayStride , stepMode , attributes : ats } = attribute . get ( ) ; this . attributeCache [ attributeName ] = attribute ; return { arrayStride , stepMode , attributes : ats , } ; } ) , } ; const descriptor = { sampleCount : this . engine . mainPassSampleCount , primitiveTopology : primitiveMap [ primitive || gl . TRIANGLES ] , rasterizationState : { ... this . getDefaultRasterizationStateDescriptor ( ) , cullMode : getCullMode ( { cull } ) , } , depthStencilState : getDepthStencilStateDescriptor ( { depth , stencil , } ) , colorStates : getColorStateDescriptors ( { blend } , this . engine . options . swapChainFormat ! ) , layout : this . pipelineLayout , vertexStage , fragmentStage , vertexState , } ; this . renderPipeline = this . engine . device . createRenderPipeline ( descriptor ) ; } public addUniforms ( uniforms ) { this . uniforms = { ... this . uniforms , ... extractUniforms ( uniforms ) , } ; } public draw ( options ) { const renderPass = this . engine . getCurrentRenderPass ( ) ; const uniforms = { ... this . uniforms , ... extractUniforms ( options . uniforms || { } ) , } ; const bindGroupBindings = [ ] ; Object . keys ( uniforms ) . forEach ( ( uniformName ) => { const type = typeof uniforms [ uniformName ] ; if ( type === 'boolean' || type === 'number' || Array . isArray ( uniforms [ uniformName ] ) || uniforms [ uniformName ] . BYTES_PER_ELEMENT ) { const offset = this . uniformGPUBufferLayout . find ( ( { name } ) => name === uniformName ) ?. offset ; if ( ! isNil ( offset ) ) { this . uniformBuffer . subData ( { data : uniforms [ uniformName ] , offset , } ) ; } } else { let offset = this . uniformGPUBufferLayout . find ( ( { name } ) => name === uniformName ) ?. offset ; if ( ! isNil ( offset ) ) { const textureOrFramebuffer = ( uniforms [ uniformName ] as WebGPUTexture2D | WebGPUFramebuffer ) . get ( ) ; const { texture , sampler } = textureOrFramebuffer . color || textureOrFramebuffer ; if ( sampler ) { bindGroupBindings . push ( { binding : offset , resource : sampler , } ) ; offset ++ ; } bindGroupBindings . push ( { binding : offset , resource : texture . createView ( ) , } ) ; } } } ) ; if ( this . uniformBuffer ) { bindGroupBindings [ 0 ] = { binding : 0 , resource : { buffer : this . uniformBuffer . get ( ) , } , } ; } this . uniformBindGroup = this . engine . device . createBindGroup ( { layout : this . uniformsBindGroupLayout , entries : bindGroupBindings , } ) ; if ( this . renderPipeline ) { renderPass . setPipeline ( this . renderPipeline ) ; } renderPass . setBindGroup ( 0 , this . uniformBindGroup ) ; if ( this . indexBuffer ) { renderPass . setIndexBuffer ( this . indexBuffer . get ( ) , WebGPUConstants . IndexFormat . Uint32 , 0 ) ; } Object . keys ( this . attributeCache ) . forEach ( ( attributeName , i ) => { renderPass . setVertexBuffer ( 0 + i , this . attributeCache [ attributeName ] . get ( ) . buffer , 0 ) ; } ) ; if ( this . indexBuffer ) { renderPass . drawIndexed ( this . indexCount , this . options . instances || 1 , 0 , 0 , 0 ) ; } else { renderPass . draw ( this . options . count || 0 , this . options . instances || 0 , 0 , 0 ) ; } } public destroy ( ) { throw new Error ( 'Method not implemented.' ) ; } private async compilePipelineStageDescriptor ( vertexCode , fragmentCode , defines ) { const shaderVersion = '#version 450\n' ; let vertexShader = vertexCode ; let fragmentShader = fragmentCode ; if ( ! this . engine . options . useWGSL ) { vertexShader = await this . compileShaderToSpirV ( vertexCode , 'vertex' , shaderVersion ) ; fragmentShader = await this . compileShaderToSpirV ( fragmentCode , 'fragment' , shaderVersion ) ; } return this . createPipelineStageDescriptor ( vertexShader , fragmentShader ) ; } private compileShaderToSpirV ( source , type , shaderVersion ) { return this . compileRawShaderToSpirV ( shaderVersion + source , type ) ; } private compileRawShaderToSpirV ( source , type ) { return this . engine . glslang . compileGLSL ( source , type ) ; } private createPipelineStageDescriptor ( vertexShader , fragmentShader ) { return { vertexStage : { module : this . engine . device . createShaderModule ( { code : vertexShader , isWHLSL : isSafari , } ) , entryPoint : 'main' , } , fragmentStage : { module : this . engine . device . createShaderModule ( { code : fragmentShader , isWHLSL : isSafari , } ) , entryPoint : 'main' , } , } ; } see https private getDefaultRasterizationStateDescriptor ( ) { return { frontFace : WebGPUConstants . FrontFace . CCW , cullMode : WebGPUConstants . CullMode . None , depthBias : 0 , depthBiasSlopeScale : 0 , depthBiasClamp : 0 , } ; } private buildUniformBindGroup ( uniforms ) { let offset = 0 ; const mergedUniformData = concatenate ( Float32Array , ... Object . keys ( uniforms ) . map ( ( uniformName ) => { if ( uniforms [ uniformName ] ) { this . uniformGPUBufferLayout . push ( { name : uniformName , offset , } ) ; offset += ( uniforms [ uniformName ] . length || 1 ) * 4 ; return uniforms [ uniformName ] ; } else { return [ ] ; } } ) ) ; const entries = [ ] ; let hasUniform = false ; if ( mergedUniformData . length ) { hasUniform = true ; entries . push ( { binding : 0 , visibility : WebGPUConstants . ShaderStage . Fragment | WebGPUConstants . ShaderStage . Vertex , type : WebGPUConstants . BindingType . UniformBuffer , } ) ; } Object . keys ( uniforms ) . filter ( ( uniformName ) => uniforms [ uniformName ] === null ) . forEach ( ( uniformName , i ) => { this . uniformGPUBufferLayout . push ( { name : uniformName , offset : i * 2 + ( hasUniform ? 1 : 0 ) , } ) ; entries . push ( { binding : i * 2 + ( hasUniform ? 1 : 0 ) , visibility : WebGPUConstants . ShaderStage . Fragment , type : WebGPUConstants . BindingType . Sampler , } , { binding : i * 2 + ( hasUniform ? 1 : 0 ) + 1 , visibility : WebGPUConstants . ShaderStage . Fragment , type : WebGPUConstants . BindingType . SampledTexture , } ) ; } ) ; this . uniformsBindGroupLayout = this . engine . device . createBindGroupLayout ( { entries , } ) ; this . pipelineLayout = this . engine . device . createPipelineLayout ( { bindGroupLayouts : [ this . uniformsBindGroupLayout ] , } ) ; if ( hasUniform ) { this . uniformBuffer = new WebGPUBuffer ( this . engine , { data : mergedUniformData instanceof Array ? new Float32Array ( mergedUniformData ) : mergedUniformData , usage : WebGPUConstants . BufferUsage . Uniform | WebGPUConstants . BufferUsage . CopyDst , } ) ; } } }