license import './flags_webgl' ; import * as tf from '@tensorflow/tfjs-core' ; import { backend_util , BackendValues , buffer , DataId , DataStorage , DataToGPUWebGLOption , DataType , DataValues , engine , env , GPUData , kernel_impls , KernelBackend , MemoryInfo , NumericDataType , Rank , RecursiveArray , scalar , ShapeMap , Tensor , Tensor2D , TensorBuffer , TensorInfo , tidy , TimingInfo , TypedArray , util } from '@tensorflow/tfjs-core' ; import { getWebGLContext } from './canvas_util' ; import { DecodeMatrixProgram } from './decode_matrix_gpu' ; import { DecodeMatrixPackedProgram } from './decode_matrix_packed_gpu' ; import { EncodeFloatProgram } from './encode_float_gpu' ; import { EncodeFloatPackedProgram } from './encode_float_packed_gpu' ; import { EncodeMatrixProgram } from './encode_matrix_gpu' ; import { EncodeMatrixPackedProgram } from './encode_matrix_packed_gpu' ; import { GPGPUContext } from './gpgpu_context' ; import * as gpgpu_math from './gpgpu_math' ; import { GPGPUBinary , GPGPUProgram , TensorData } from './gpgpu_math' ; import { simpleAbsImplCPU } from './kernel_utils/shared' ; import { PackProgram } from './pack_gpu' ; import { ReshapePackedProgram } from './reshape_packed_gpu' ; import * as tex_util from './tex_util' ; import { Texture , TextureData , TextureUsage } from './tex_util' ; import { TextureManager } from './texture_manager' ; import * as unary_op from './unaryop_gpu' ; import { UnaryOpProgram } from './unaryop_gpu' ; import { UnaryOpPackedProgram } from './unaryop_packed_gpu' ; import { UnpackProgram } from './unpack_gpu' ; import * as webgl_util from './webgl_util' ; const whereImpl = kernel_impls . whereImpl ; export const EPSILON_FLOAT32 = 1e-7 ; export const EPSILON_FLOAT16 = 1e-4 ; type KernelInfo = { name : string ; query : Promise < number > ; } ; export type TimerNode = RecursiveArray < KernelInfo > | KernelInfo ; export interface CPUTimerQuery { startMs : number ; endMs ? : number ; } export interface WebGLMemoryInfo extends MemoryInfo { numBytesInGPU : number ; numBytesInGPUAllocated : number ; numBytesInGPUFree : number ; unreliable : boolean ; } export interface WebGLTimingInfo extends TimingInfo { uploadWaitMs : number ; downloadWaitMs : number ; } const binaryCaches = { } ; export function getBinaryCache ( webGLVersion ) { if ( webGLVersion in binaryCaches ) { return binaryCaches [ webGLVersion ] ; } binaryCaches [ webGLVersion ] = { } ; return binaryCaches [ webGLVersion ] ; } const CPU_HANDOFF_SIZE_THRESHOLD = env ( ) . getNumber ( 'CPU_HANDOFF_SIZE_THRESHOLD' ) ; const BEFORE_PAGING_CONSTANT = 600 ; function numMBBeforeWarning ( ) { if ( env ( ) . global . screen == null ) { return 1024 ; } return ( env ( ) . global . screen . height * env ( ) . global . screen . width * window . devicePixelRatio ) * BEFORE_PAGING_CONSTANT / 1024 / 1024 ; } export class MathBackendWebGL extends KernelBackend { texData : DataStorage < TextureData > ; gpgpu : GPGPUContext ; private static nextDataId = 0 ; private nextDataId ( ) { return MathBackendWebGL . nextDataId ++ ; } private pendingRead = new WeakMap < DataId , Array < ( arr : TypedArray ) => void > > ( ) ; private pendingDisposal = new WeakSet < DataId > ( ) ; dataRefCount = new WeakMap < DataId , number > ( ) ; private numBytesInGPU = 0 ; private canvas : HTMLCanvasElement | OffscreenCanvas ; private programTimersStack : TimerNode [ ] ; private activeTimers : TimerNode [ ] ; private uploadWaitMs = 0 ; private downloadWaitMs = 0 ; private lastGlFlushTime = 0 ; private floatPrecisionValue : 32 | 16 ; private textureManager : TextureManager ; private binaryCache : { [ key ] : GPGPUBinary ; } ; private gpgpuCreatedLocally : boolean ; private numMBBeforeWarning : number ; private warnedAboutMemory = false ; constructor ( gpuResource ? ) { super ( ) ; if ( ! env ( ) . getBool ( 'HAS_WEBGL' ) ) { throw new Error ( 'WebGL is not supported on this device' ) ; } let newGPGPU ; if ( gpuResource != null ) { if ( gpuResource instanceof GPGPUContext ) { newGPGPU = gpuResource ; } else { const gl = getWebGLContext ( env ( ) . getNumber ( 'WEBGL_VERSION' ) , gpuResource ) ; newGPGPU = new GPGPUContext ( gl ) ; } this . binaryCache = { } ; this . gpgpuCreatedLocally = false ; } else { const gl = getWebGLContext ( env ( ) . getNumber ( 'WEBGL_VERSION' ) ) ; newGPGPU = new GPGPUContext ( gl ) ; this . binaryCache = getBinaryCache ( env ( ) . getNumber ( 'WEBGL_VERSION' ) ) ; this . gpgpuCreatedLocally = true ; } this . gpgpu = newGPGPU ; this . canvas = this . gpgpu . gl . canvas ; this . textureManager = new TextureManager ( this . gpgpu ) ; this . numMBBeforeWarning = numMBBeforeWarning ( ) ; this . texData = new DataStorage ( this , engine ( ) ) ; } numDataIds ( ) { return this . texData . numDataIds ( ) - this . pendingDeletes ; } write ( values , shape , dtype ) { if ( env ( ) . getBool ( 'WEBGL_CHECK_NUMERICAL_PROBLEMS' ) || env ( ) . getBool ( 'DEBUG' ) ) { this . checkNumericalProblems ( values ) ; } if ( dtype === 'complex64' && values != null ) { throw new Error ( `Cannot write to a complex64 dtype. ` + `Please use tf.complex(real, imag).` ) ; } const dataId = { id : this . nextDataId ( ) } ; this . texData . set ( dataId , { shape , dtype , values , usage : TextureUsage . UPLOAD , refCount : 1 } ) ; return dataId ; } /** Return refCount of a `TensorData`. */ refCount ( dataId ) { if ( this . texData . has ( dataId ) ) { const tensorData = this . texData . get ( dataId ) ; return tensorData . refCount ; } return 0 ; } /** Increase refCount of a `TextureData`. */ incRef ( dataId ) { const texData = this . texData . get ( dataId ) ; texData . refCount ++ ; } /** Decrease refCount of a `TextureData`. */ decRef ( dataId ) { if ( this . texData . has ( dataId ) ) { const texData = this . texData . get ( dataId ) ; texData . refCount -- ; } } move ( dataId , values , shape , dtype , refCount ) { if ( env ( ) . getBool ( 'DEBUG' ) ) { this . checkNumericalProblems ( values ) ; } if ( dtype === 'complex64' ) { throw new Error ( `Cannot write to a complex64 dtype. ` + `Please use tf.complex(real, imag).` ) ; } this . texData . set ( dataId , { shape , dtype , values , usage : TextureUsage . UPLOAD , refCount } ) ; } disposeIntermediateTensorInfo ( tensorInfo ) { this . disposeData ( tensorInfo . dataId ) ; } readSync ( dataId ) { const texData = this . texData . get ( dataId ) ; const { values , dtype , complexTensorInfos , slice , shape , isPacked } = texData ; if ( slice != null ) { let program ; if ( isPacked ) { program = new UnaryOpPackedProgram ( shape , unary_op . CLONE ) ; } else { program = new UnaryOpProgram ( shape , unary_op . CLONE ) ; } const res = this . runWebGLProgram ( program , [ { dataId , shape , dtype } ] , dtype ) ; const data = this . readSync ( res . dataId ) ; this . disposeIntermediateTensorInfo ( res ) ; return data ; } if ( values != null ) { return this . convertAndCacheOnCPU ( dataId ) ; } if ( dtype === 'string' ) { return values ; } const shouldTimeProgram = this . activeTimers != null ; let start ; if ( shouldTimeProgram ) { start = util . now ( ) ; } let result ; if ( dtype === 'complex64' ) { const realValues = this . readSync ( complexTensorInfos . real . dataId ) as Float32Array ; const imagValues = this . readSync ( complexTensorInfos . imag . dataId ) as Float32Array ; result = backend_util . mergeRealAndImagArrays ( realValues , imagValues ) ; } else { result = this . getValuesFromTexture ( dataId ) ; } if ( shouldTimeProgram ) { this . downloadWaitMs += util . now ( ) - start ; } return this . convertAndCacheOnCPU ( dataId , result ) ; } async read ( dataId ) { if ( this . pendingRead . has ( dataId ) ) { const subscribers = this . pendingRead . get ( dataId ) ; return new Promise < TypedArray > ( resolve => subscribers . push ( resolve ) ) ; } const texData = this . texData . get ( dataId ) ; const { values , shape , slice , dtype , complexTensorInfos , isPacked } = texData ; if ( slice != null ) { let program ; if ( isPacked ) { program = new UnaryOpPackedProgram ( shape , unary_op . CLONE ) ; } else { program = new UnaryOpProgram ( shape , unary_op . CLONE ) ; } const res = this . runWebGLProgram ( program , [ { dataId , shape , dtype } ] , dtype ) ; const data = this . read ( res . dataId ) ; this . disposeIntermediateTensorInfo ( res ) ; return data ; } if ( values != null ) { return this . convertAndCacheOnCPU ( dataId ) ; } if ( env ( ) . getBool ( 'DEBUG' ) ) { if ( ! env ( ) . getBool ( 'WEBGL_DOWNLOAD_FLOAT_ENABLED' ) && env ( ) . getNumber ( 'WEBGL_VERSION' ) === 2 ) { throw new Error ( `tensor.data() with WEBGL_DOWNLOAD_FLOAT_ENABLED=false and ` + `WEBGL_VERSION=2 not yet supported.` ) ; } } let buffer = null ; let tmpDownloadTarget ; if ( dtype !== 'complex64' && env ( ) . get ( 'WEBGL_BUFFER_SUPPORTED' ) ) { tmpDownloadTarget = this . decode ( dataId ) ; const tmpData = this . texData . get ( tmpDownloadTarget . dataId ) ; buffer = this . gpgpu . createBufferFromTexture ( tmpData . texture . texture , ... tex_util . getDenseTexShape ( shape ) ) ; } this . pendingRead . set ( dataId , [ ] ) ; if ( dtype !== 'complex64' ) { await this . gpgpu . createAndWaitForFence ( ) ; } let vals ; if ( dtype === 'complex64' ) { const ps = await Promise . all ( [ this . read ( complexTensorInfos . real . dataId ) , this . read ( complexTensorInfos . imag . dataId ) ] ) ; const realValues = ps [ 0 ] ; const imagValues = ps [ 1 ] ; vals = backend_util . mergeRealAndImagArrays ( realValues as Float32Array , imagValues as Float32Array ) ; } else if ( buffer == null ) { vals = this . getValuesFromTexture ( dataId ) ; } else { const size = util . sizeFromShape ( shape ) ; vals = this . gpgpu . downloadFloat32MatrixFromBuffer ( buffer , size ) ; } if ( tmpDownloadTarget != null ) { this . disposeIntermediateTensorInfo ( tmpDownloadTarget ) ; } if ( buffer != null ) { const gl = this . gpgpu . gl ; webgl_util . callAndCheck ( gl , ( ) => gl . deleteBuffer ( buffer ) ) ; } const dTypeVals = this . convertAndCacheOnCPU ( dataId , vals ) ; const subscribers = this . pendingRead . get ( dataId ) ; this . pendingRead . delete ( dataId ) ; subscribers . forEach ( resolve => resolve ( dTypeVals ) ) ; if ( this . pendingDisposal . has ( dataId ) ) { this . pendingDisposal . delete ( dataId ) ; if ( this . disposeData ( dataId ) ) { engine ( ) . removeDataId ( dataId , this ) ; } this . pendingDeletes -- ; } return dTypeVals ; } param dataId param options readToGPU ( dataId , options = { } ) { const texData = this . texData . get ( dataId ) ; const { values , shape , slice , dtype , isPacked , texture } = texData ; if ( dtype === 'complex64' ) { throw new Error ( 'Does not support reading texture for complex64 dtype.' ) ; } if ( slice != null ) { let program ; if ( isPacked ) { program = new UnaryOpPackedProgram ( shape , unary_op . CLONE ) ; } else { program = new UnaryOpProgram ( shape , unary_op . CLONE ) ; } const res = this . runWebGLProgram ( program , [ { dataId , shape , dtype } ] , dtype ) ; const gpuResouorce = this . readToGPU ( res , options ) ; this . disposeIntermediateTensorInfo ( res ) ; return gpuResouorce ; } if ( texture == null ) { if ( values != null ) { throw new Error ( 'Data is not on GPU but on CPU.' ) ; } else { throw new Error ( 'There is no data on GPU or CPU.' ) ; } } const tmpTarget = this . decode ( dataId , options . customTexShape ) ; const tensorRef = engine ( ) . makeTensorFromDataId ( tmpTarget . dataId , tmpTarget . shape , tmpTarget . dtype ) ; const tmpData = this . texData . get ( tmpTarget . dataId ) ; return { tensorRef , ... tmpData . texture } ; } bufferSync < R extends Rank > ( t ) { const data = this . readSync ( t . dataId ) ; let decodedData = data as DataValues ; if ( t . dtype === 'string' ) { try { decodedData = ( data as Uint8Array [ ] ) . map ( d => util . decodeString ( d ) ) ; } catch { throw new Error ( 'Failed to decode encoded string bytes into utf-8' ) ; } } return buffer ( t . shape as ShapeMap [ R ] , t . dtype , decodedData ) as TensorBuffer < R > ; } private checkNumericalProblems ( values ) { if ( values == null ) { return ; } for ( let i = 0 ; i < values . length ; i ++ ) { const num = values [ i ] as number ; if ( ! webgl_util . canBeRepresented ( num ) ) { if ( env ( ) . getBool ( 'WEBGL_RENDER_FLOAT32_CAPABLE' ) ) { throw Error ( `The value ${ num } cannot be represented with your ` + `current settings. Consider enabling float32 rendering: ` + `'tf.env().set('WEBGL_RENDER_FLOAT32_ENABLED', true);'` ) ; } throw Error ( `The value ${ num } cannot be represented on this device.` ) ; } } } private getValuesFromTexture ( dataId ) { const { shape , dtype , isPacked } = this . texData . get ( dataId ) ; const size = util . sizeFromShape ( shape ) ; if ( env ( ) . getBool ( 'WEBGL_DOWNLOAD_FLOAT_ENABLED' ) ) { const tmpTarget = this . decode ( dataId ) ; const tmpData = this . texData . get ( tmpTarget . dataId ) ; const vals = this . gpgpu . downloadMatrixFromPackedTexture ( tmpData . texture . texture , ... tex_util . getDenseTexShape ( shape ) ) . subarray ( 0 , size ) ; this . disposeIntermediateTensorInfo ( tmpTarget ) ; return vals ; } const shouldUsePackedProgram = env ( ) . getBool ( 'WEBGL_PACK' ) && isPacked === true ; const outputShape = shouldUsePackedProgram ? webgl_util . getShapeAs3D ( shape ) : shape ; const program = shouldUsePackedProgram ? new EncodeFloatPackedProgram ( outputShape as [ number , number , number ] ) : new EncodeFloatProgram ( outputShape ) ; const output = this . runWebGLProgram ( program , [ { shape : outputShape , dtype , dataId } ] , 'float32' ) ; const tmpData = this . texData . get ( output . dataId ) ; const vals = this . gpgpu . downloadByteEncodedFloatMatrixFromOutputTexture ( tmpData . texture . texture , tmpData . texShape [ 0 ] , tmpData . texShape [ 1 ] ) . subarray ( 0 , size ) ; this . disposeIntermediateTensorInfo ( output ) ; return vals ; } timerAvailable ( ) { return env ( ) . getNumber ( 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE' ) > 0 ; } time ( f ) { const oldActiveTimers = this . activeTimers ; const newActiveTimers = [ ] ; let outerMostTime = false ; if ( this . programTimersStack == null ) { this . programTimersStack = newActiveTimers ; outerMostTime = true ; } else { this . activeTimers . push ( newActiveTimers ) ; } this . activeTimers = newActiveTimers ; f ( ) ; const flattenedActiveTimerQueries = util . flatten ( this . activeTimers . map ( ( d ) => d . query ) ) . filter ( d => d != null ) ; const flattenedActiveTimerNames = util . flatten ( this . activeTimers . map ( ( d ) => d . name ) ) . filter ( d => d != null ) ; this . activeTimers = oldActiveTimers ; if ( outerMostTime ) { this . programTimersStack = null ; } const res = { uploadWaitMs : this . uploadWaitMs , downloadWaitMs : this . downloadWaitMs , kernelMs : null , wallMs : null } ; return ( async ( ) => { if ( env ( ) . getNumber ( 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE' ) > 0 ) { const kernelMs = await Promise . all ( flattenedActiveTimerQueries ) ; res [ 'kernelMs' ] = util . sum ( kernelMs ) ; res [ 'getExtraProfileInfo' ] = ( ) => kernelMs . map ( ( d , i ) => ( { name : flattenedActiveTimerNames [ i ] , ms : d } ) ) . map ( d => `${ d . name }: ${ d . ms }` ) . join ( ', ' ) ; } else { res [ 'kernelMs' ] = { error : 'WebGL query timers are not supported in this environment.' } ; } this . uploadWaitMs = 0 ; this . downloadWaitMs = 0 ; return res ; } ) ( ) ; } memory ( ) { return { unreliable : false , numBytesInGPU : this . numBytesInGPU , numBytesInGPUAllocated : this . textureManager . numBytesAllocated , numBytesInGPUFree : this . textureManager . numBytesFree } as WebGLMemoryInfo ; } private startTimer ( ) { if ( env ( ) . getNumber ( 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE' ) > 0 ) { return this . gpgpu . beginQuery ( ) ; } return { startMs : util . now ( ) , endMs : null } ; } private endTimer ( query ) { if ( env ( ) . getNumber ( 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE' ) > 0 ) { this . gpgpu . endQuery ( ) ; return query ; } ( query as CPUTimerQuery ) . endMs = util . now ( ) ; return query ; } private async getQueryTime ( query ) { if ( env ( ) . getNumber ( 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE' ) > 0 ) { return this . gpgpu . waitForQueryAndGetTime ( query as WebGLQuery ) ; } const timerQuery = query as CPUTimerQuery ; return timerQuery . endMs - timerQuery . startMs ; } private pendingDeletes = 0 ; param dataId oaram disposeData ( dataId , force = false ) { if ( this . pendingDisposal . has ( dataId ) ) { return false ; } if ( ! this . texData . has ( dataId ) ) { return true ; } if ( force ) { this . texData . get ( dataId ) . refCount = 0 ; } else { this . texData . get ( dataId ) . refCount -- ; } if ( ! force && this . texData . get ( dataId ) . refCount > 0 ) { return false ; } if ( this . pendingRead . has ( dataId ) ) { this . pendingDisposal . add ( dataId ) ; this . pendingDeletes ++ ; return false ; } this . releaseGPUData ( dataId ) ; const { complexTensorInfos } = this . texData . get ( dataId ) ; if ( complexTensorInfos != null ) { this . disposeData ( complexTensorInfos . real . dataId , force ) ; this . disposeData ( complexTensorInfos . imag . dataId , force ) ; } this . texData . delete ( dataId ) ; return true ; } private releaseGPUData ( dataId ) { const { texture , dtype , texShape , usage , isPacked , slice } = this . texData . get ( dataId ) ; const key = slice && slice . origDataId || dataId ; const refCount = this . dataRefCount . get ( key ) ; if ( refCount > 1 ) { this . dataRefCount . set ( key , refCount - 1 ) ; } else { this . dataRefCount . delete ( key ) ; if ( texture != null ) { this . numBytesInGPU -= this . computeBytes ( texShape , dtype ) ; this . textureManager . releaseTexture ( texture , texShape , usage , isPacked ) ; } } const texData = this . texData . get ( dataId ) ; texData . texture = null ; texData . texShape = null ; texData . isPacked = false ; texData . slice = null ; } getTexture ( dataId ) { this . uploadToGPU ( dataId ) ; return this . texData . get ( dataId ) . texture . texture ; } /**
     * Returns internal information for the specific data bucket. Used in unit
     * tests.
     */ getDataInfo ( dataId ) { return this . texData . get ( dataId ) ; } shouldExecuteOnCPU ( inputs , sizeThreshold = CPU_HANDOFF_SIZE_THRESHOLD ) { return env ( ) . getBool ( 'WEBGL_CPU_FORWARD' ) && inputs . every ( input => this . texData . get ( input . dataId ) . texture == null && util . sizeFromShape ( input . shape ) < sizeThreshold ) ; } getGPGPUContext ( ) { return this . gpgpu ; } where ( condition ) { backend_util . warn ( 'tf.where() in webgl locks the UI thread. ' + 'Call tf.whereAsync() instead' ) ; const condVals = condition . dataSync ( ) ; return whereImpl ( condition . shape , condVals ) ; } private packedUnaryOp ( x , op , dtype ) { const program = new UnaryOpPackedProgram ( x . shape , op ) ; const outInfo = this . compileAndRun ( program , [ x ] , dtype ) ; return engine ( ) . makeTensorFromDataId ( outInfo . dataId , outInfo . shape , outInfo . dtype ) ; } abs < T extends Tensor > ( x ) { if ( this . shouldExecuteOnCPU ( [ x ] ) && x . dtype !== 'complex64' ) { const outValues = simpleAbsImplCPU ( this . texData . get ( x . dataId ) . values as TypedArray ) ; return this . makeOutput ( x . shape , x . dtype , outValues ) ; } if ( env ( ) . getBool ( 'WEBGL_PACK_UNARY_OPERATIONS' ) ) { return this . packedUnaryOp ( x , unary_op . ABS , x . dtype ) as T ; } const program = new UnaryOpProgram ( x . shape , unary_op . ABS ) ; const outInfo = this . compileAndRun ( program , [ x ] ) ; return engine ( ) . makeTensorFromDataId ( outInfo . dataId , outInfo . shape , outInfo . dtype ) as T ; } makeTensorInfo ( shape , dtype , values ? ) { let dataId ; if ( dtype === 'string' && values != null && values . length > 0 && util . isString ( values [ 0 ] ) ) { const encodedValues = ( values as { } as string [ ] ) . map ( d => util . encodeString ( d ) ) ; dataId = this . write ( encodedValues , shape , dtype ) ; } else { dataId = this . write ( values as TypedArray , shape , dtype ) ; } this . texData . get ( dataId ) . usage = null ; return { dataId , shape , dtype } ; } private makeOutput < T extends Tensor > ( shape , dtype , values ? ) { const { dataId } = this . makeTensorInfo ( shape , dtype , values ) ; return engine ( ) . makeTensorFromDataId ( dataId , shape , dtype , this ) as T ; } unpackTensor ( input ) { const program = new UnpackProgram ( input . shape ) ; return this . runWebGLProgram ( program , [ input ] , input . dtype ) ; } packTensor ( input ) { const program = new PackProgram ( input . shape ) ; const preventEagerUnpackingOutput = true ; return this . runWebGLProgram ( program , [ input ] , input . dtype , null , preventEagerUnpackingOutput ) ; } private packedReshape ( input , afterShape ) { const input3DShape = [ webgl_util . getBatchDim ( input . shape ) , ... webgl_util . getRowsCols ( input . shape ) ] as [ number , number , number ] ; const input3D = { dtype : input . dtype , shape : input3DShape , dataId : input . dataId } ; const afterShapeAs3D = [ webgl_util . getBatchDim ( afterShape ) , ... webgl_util . getRowsCols ( afterShape ) ] as [ number , number , number ] ; const program = new ReshapePackedProgram ( afterShapeAs3D , input3DShape ) ; const preventEagerUnpackingOfOutput = true ; const customValues = [ input3DShape ] ; const output = this . runWebGLProgram ( program , [ input3D ] , input . dtype , customValues , preventEagerUnpackingOfOutput ) ; return { dataId : output . dataId , shape : afterShape , dtype : output . dtype } ; } private decode ( dataId , customTexShape ? ) { const texData = this . texData . get ( dataId ) ; const { isPacked , shape , dtype } = texData ; if ( customTexShape != null ) { const size = util . sizeFromShape ( shape ) ; const texSize = customTexShape [ 0 ] * customTexShape [ 1 ] * 4 ; util . assert ( size <= texSize , ( ) => 'customTexShape is too small. ' + 'Row * Column * 4 should be equal or larger than the ' + 'size of the tensor data.' ) ; } const shapeAs3D = webgl_util . getShapeAs3D ( shape ) as [ number , number , number ] ; let program ; if ( isPacked ) { program = new DecodeMatrixPackedProgram ( shapeAs3D ) ; } else { program = new DecodeMatrixProgram ( shapeAs3D ) ; } const preventEagerUnpackingOfOutput = true ; const customValues = [ customTexShape != null ? customTexShape : tex_util . getDenseTexShape ( shapeAs3D ) ] ; const out = this . runWebGLProgram ( program , [ { shape : shapeAs3D , dtype , dataId } ] , dtype , customValues , preventEagerUnpackingOfOutput , customTexShape ) ; return { dtype , shape , dataId : out . dataId } ; } runWebGLProgram ( program , inputs , outputDtype , customUniformValues ? , preventEagerUnpackingOfOutput = false , customTexShape ? ) { const output = this . makeTensorInfo ( program . outputShape , outputDtype ) ; const outData = this . texData . get ( output . dataId ) ; if ( program . packedOutput ) { outData . isPacked = true ; } if ( program . outPackingScheme === tex_util . PackingScheme . DENSE ) { const texelShape = customTexShape != null ? customTexShape : tex_util . getDenseTexShape ( program . outputShape ) ; outData . texShape = texelShape . map ( d => d * 2 ) as [ number , number ] ; } if ( program . outTexUsage != null ) { outData . usage = program . outTexUsage ; } if ( util . sizeFromShape ( output . shape ) === 0 ) { outData . values = util . getTypedArrayFromDType ( output . dtype as 'float32' , 0 ) ; return output ; } const dataToDispose = [ ] ; const inputsData = inputs . map ( input => { if ( input . dtype === 'complex64' ) { throw new Error ( `GPGPUProgram does not support complex64 input. For complex64 ` + `dtypes, please separate the program into real and imaginary ` + `parts.` ) ; } let texData = this . texData . get ( input . dataId ) ; if ( texData . texture == null ) { if ( ! program . packedInputs && util . sizeFromShape ( input . shape ) <= env ( ) . getNumber ( 'WEBGL_SIZE_UPLOAD_UNIFORM' ) ) { return { shape : input . shape , texData : null , isUniform : true , uniformValues : texData . values as TypedArray } ; } if ( program . packedInputs ) { texData . isPacked = true ; texData . shape = input . shape ; } } this . uploadToGPU ( input . dataId ) ; if ( ! ! texData . isPacked !== ! ! program . packedInputs ) { input = texData . isPacked ? this . unpackTensor ( input ) : this . packTensor ( input ) ; dataToDispose . push ( input ) ; texData = this . texData . get ( input . dataId ) ; } else if ( texData . isPacked && ! webgl_util . isReshapeFree ( texData . shape , input . shape ) ) { const savedInput = input ; const targetShape = input . shape ; input . shape = texData . shape ; input = this . packedReshape ( input as Tensor , targetShape ) ; dataToDispose . push ( input ) ; texData = this . texData . get ( input . dataId ) ; savedInput . shape = targetShape ; } return { shape : input . shape , texData , isUniform : false } ; } ) ; this . uploadToGPU ( output . dataId ) ; const outputData = { shape : output . shape , texData : outData , isUniform : false } ; const key = gpgpu_math . makeShaderKey ( program , inputsData , outputData ) ; const binary = this . getAndSaveBinary ( key , ( ) => { return gpgpu_math . compileProgram ( this . gpgpu , program , inputsData , outputData ) ; } ) ; const shouldTimeProgram = this . activeTimers != null ; let query ; if ( shouldTimeProgram ) { query = this . startTimer ( ) ; } gpgpu_math . runProgram ( this . gpgpu , binary , inputsData , outputData , customUniformValues ) ; dataToDispose . forEach ( info => this . disposeIntermediateTensorInfo ( info ) ) ; if ( shouldTimeProgram ) { query = this . endTimer ( query ) ; this . activeTimers . push ( { name : program . constructor . name , query : this . getQueryTime ( query ) } ) ; } const glFlushThreshold = env ( ) . get ( 'WEBGL_FLUSH_THRESHOLD' ) ; if ( glFlushThreshold > 0 ) { const time = util . now ( ) ; if ( ( time - this . lastGlFlushTime ) > glFlushThreshold ) { this . gpgpu . gl . flush ( ) ; this . lastGlFlushTime = time ; } } if ( ! env ( ) . getBool ( 'WEBGL_LAZILY_UNPACK' ) && outData . isPacked && preventEagerUnpackingOfOutput === false ) { const unpacked = this . unpackTensor ( output ) ; this . disposeIntermediateTensorInfo ( output ) ; return unpacked ; } return output ; } compileAndRun ( program , inputs , outputDtype ? , customUniformValues ? , preventEagerUnpackingOfOutput = false ) { outputDtype = outputDtype || inputs [ 0 ] . dtype ; const outInfo = this . runWebGLProgram ( program , inputs , outputDtype , customUniformValues , preventEagerUnpackingOfOutput ) ; return outInfo ; } private getAndSaveBinary ( key , getBinary ) { if ( ! ( key in this . binaryCache ) ) { this . binaryCache [ key ] = getBinary ( ) ; } return this . binaryCache [ key ] ; } getTextureManager ( ) { return this . textureManager ; } private disposed = false ; dispose ( ) { if ( this . disposed ) { return ; } if ( ! env ( ) . getBool ( 'IS_TEST' ) ) { const allKeys = Object . keys ( this . binaryCache ) ; allKeys . forEach ( key => { this . gpgpu . deleteProgram ( this . binaryCache [ key ] . webGLProgram ) ; delete this . binaryCache [ key ] ; } ) ; } this . textureManager . dispose ( ) ; if ( this . canvas != null && ( typeof ( HTMLCanvasElement ) !== 'undefined' && this . canvas instanceof HTMLCanvasElement ) ) { this . canvas . remove ( ) ; } else { this . canvas = null ; } if ( this . gpgpuCreatedLocally ) { this . gpgpu . program = null ; this . gpgpu . dispose ( ) ; } this . disposed = true ; } floatPrecision ( ) { if ( this . floatPrecisionValue == null ) { this . floatPrecisionValue = tidy ( ( ) => { if ( ! env ( ) . get ( 'WEBGL_RENDER_FLOAT32_ENABLED' ) ) { const debugFlag = env ( ) . getBool ( 'DEBUG' ) ; env ( ) . set ( 'DEBUG' , false ) ; const underflowCheckValue = this . abs ( scalar ( 1e-8 ) ) . dataSync ( ) [ 0 ] ; env ( ) . set ( 'DEBUG' , debugFlag ) ; if ( underflowCheckValue > 0 ) { return 32 ; } } return 16 ; } ) ; } return this . floatPrecisionValue ; } /** Returns the smallest representable number.  */ epsilon ( ) { return this . floatPrecision ( ) === 32 ? EPSILON_FLOAT32 : EPSILON_FLOAT16 ; } uploadToGPU ( dataId ) { const texData = this . texData . get ( dataId ) ; const { shape , dtype , values , texture , usage , isPacked } = texData ; if ( texture != null ) { return ; } const shouldTimeProgram = this . activeTimers != null ; let start ; if ( shouldTimeProgram ) { start = util . now ( ) ; } let texShape = texData . texShape ; if ( texShape == null ) { texShape = webgl_util . getTextureShapeFromLogicalShape ( shape , isPacked ) ; texData . texShape = texShape ; } if ( values != null ) { const shapeAs3D = webgl_util . getShapeAs3D ( shape ) ; let program ; let width = texShape [ 1 ] , height = texShape [ 0 ] ; const isByteArray = values instanceof Uint8Array || values instanceof Uint8ClampedArray ; if ( isPacked || ! isByteArray ) { [ width , height ] = tex_util . getPackedMatrixTextureShapeWidthHeight ( texShape [ 0 ] , texShape [ 1 ] ) ; } if ( isPacked ) { program = new EncodeMatrixPackedProgram ( shapeAs3D , isByteArray ) ; } else { program = new EncodeMatrixProgram ( shapeAs3D , isByteArray ) ; } const tempDenseInputTexShape = isByteArray ? [ height , width ] : texShape ; const tempDenseInputHandle = this . makeTensorInfo ( tempDenseInputTexShape , dtype ) ; const tempDenseInputTexData = this . texData . get ( tempDenseInputHandle . dataId ) ; if ( isByteArray ) { tempDenseInputTexData . usage = TextureUsage . PIXELS ; } else { tempDenseInputTexData . usage = TextureUsage . UPLOAD ; } tempDenseInputTexData . texShape = tempDenseInputTexShape ; this . gpgpu . uploadDenseMatrixToTexture ( this . getTexture ( tempDenseInputHandle . dataId ) , width , height , values as TypedArray ) ; const customValues = [ [ height , width ] ] ; const preventEagerUnpacking = true ; const encodedOutputTarget = this . runWebGLProgram ( program , [ tempDenseInputHandle ] , dtype , customValues , preventEagerUnpacking ) ; const outputTexData = this . texData . get ( encodedOutputTarget . dataId ) ; texData . texture = outputTexData . texture ; texData . texShape = outputTexData . texShape ; texData . isPacked = outputTexData . isPacked ; texData . usage = outputTexData . usage ; this . disposeIntermediateTensorInfo ( tempDenseInputHandle ) ; this . texData . delete ( encodedOutputTarget . dataId ) ; texData . values = null ; if ( shouldTimeProgram ) { this . uploadWaitMs += util . now ( ) - start ; } } else { const newTexture = this . acquireTexture ( texShape , usage , dtype , isPacked ) ; texData . texture = newTexture ; } } private convertAndCacheOnCPU ( dataId , float32Values ? ) { const texData = this . texData . get ( dataId ) ; const { dtype } = texData ; this . releaseGPUData ( dataId ) ; if ( float32Values != null ) { texData . values = float32ToTypedArray ( float32Values , dtype as 'float32' ) ; } return texData . values as TypedArray ; } private acquireTexture ( texShape , texType , dtype , isPacked ) { this . numBytesInGPU += this . computeBytes ( texShape , dtype ) ; if ( ! this . warnedAboutMemory && this . numBytesInGPU > this . numMBBeforeWarning * 1024 * 1024 ) { const mb = ( this . numBytesInGPU / 1024 / 1024 ) . toFixed ( 2 ) ; this . warnedAboutMemory = true ; console . warn ( `High memory usage in GPU: ${ mb } MB, ` + `most likely due to a memory leak` ) ; } return this . textureManager . acquireTexture ( texShape , texType , isPacked ) ; } private computeBytes ( shape , dtype ) { return shape [ 0 ] * shape [ 1 ] * util . bytesPerElement ( dtype ) ; } } function float32ToTypedArray < D extends NumericDataType > ( a , dtype ) { if ( dtype === 'float32' || dtype === 'complex64' ) { return a as tf . DataTypeMap [ D ] ; } else if ( dtype === 'int32' || dtype === 'bool' ) { const result = ( dtype === 'int32' ) ? new Int32Array ( a . length ) : new Uint8Array ( a . length ) ; for ( let i = 0 ; i < result . length ; ++ i ) { result [ i ] = Math . round ( a [ i ] ) ; } return result as tf . DataTypeMap [ D ] ; } else { throw new Error ( `Unknown dtype ${ dtype }` ) ; } }