license import * as tf from '@tensorflow/tfjs-core' ; import { engine , GPUData , test_util , util } from '@tensorflow/tfjs-core' ; import { describeWithFlags } from '@tensorflow/tfjs-core/dist/jasmine_util' ; const { expectArraysClose , expectArraysEqual } = test_util ; const { decodeString } = util ; import { getBinaryCache , MathBackendWebGL , WebGLMemoryInfo , WebGLTimingInfo } from './backend_webgl' ; import { computeBytes } from './texture_manager' ; import { PhysicalTextureType } from './tex_util' ; import { WEBGL_ENVS } from './backend_webgl_test_registry' ; import { GPGPUContext } from './gpgpu_context' ; function decodeStrings ( bytes ) { return bytes . map ( b => decodeString ( b ) ) ; } const RENDER_FLOAT32_ENVS = { flags : { 'WEBGL_RENDER_FLOAT32_ENABLED' : true } , predicate : WEBGL_ENVS . predicate } ; describeWithFlags ( 'forced f16 render' , RENDER_FLOAT32_ENVS , ( ) => { beforeEach ( ( ) => { tf . env ( ) . set ( 'WEBGL_RENDER_FLOAT32_ENABLED' , false ) ; } ) ; afterAll ( ( ) => tf . env ( ) . reset ( ) ) ; it ( 'should overflow if larger than 66k' , async ( ) => { const a = tf . tensor1d ( [ Math . pow ( 2 , 17 ) ] , 'float32' ) ; const b = tf . relu ( a ) ; const data = await b . data ( ) ; if ( data && data . length === 1 && ! isNaN ( data [ 0 ] ) ) { expect ( data ) . toBeLessThan ( Math . pow ( 2 , 17 ) ) ; } } ) ; it ( 'should error in debug mode' , ( ) => { spyOn ( console , 'warn' ) ; const debug = tf . env ( ) . get ( 'DEBUG' ) ; tf . env ( ) . set ( 'DEBUG' , true ) ; const a = ( ) => tf . tensor1d ( [ 2 , Math . pow ( 2 , 17 ) ] , 'float32' ) ; expect ( a ) . toThrowError ( ) ; tf . env ( ) . set ( 'DEBUG' , debug ) ; } ) ; } ) ; describeWithFlags ( 'lazy packing and unpacking' , WEBGL_ENVS , ( ) => { let webglLazilyUnpackFlagSaved ; let webglCpuForwardFlagSaved ; beforeAll ( ( ) => { webglLazilyUnpackFlagSaved = tf . env ( ) . getBool ( 'WEBGL_LAZILY_UNPACK' ) ; webglCpuForwardFlagSaved = tf . env ( ) . getBool ( 'WEBGL_CPU_FORWARD' ) ; tf . env ( ) . set ( 'WEBGL_LAZILY_UNPACK' , true ) ; tf . env ( ) . set ( 'WEBGL_CPU_FORWARD' , false ) ; } ) ; afterAll ( ( ) => { tf . env ( ) . set ( 'WEBGL_LAZILY_UNPACK' , webglLazilyUnpackFlagSaved ) ; tf . env ( ) . set ( 'WEBGL_CPU_FORWARD' , webglCpuForwardFlagSaved ) ; } ) ; it ( 'should not leak memory when lazily unpacking' , ( ) => { const a = tf . tensor2d ( [ 1 , 2 , 3 , 4 , 5 , 6 ] , [ 2 , 3 ] ) ; const b = tf . tensor2d ( [ 0 , 1 , - 3 , 2 , 2 , 1 ] , [ 3 , 2 ] ) ; const c = tf . matMul ( a , b ) ; const startNumBytes = tf . memory ( ) . numBytes ; const startNumTensors = tf . memory ( ) . numTensors ; const startNumBytesInGPU = ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPU ; const webglPackBinaryOperationsFlagSaved = tf . env ( ) . getBool ( 'WEBGL_PACK_BINARY_OPERATIONS' ) ; tf . env ( ) . set ( 'WEBGL_PACK_BINARY_OPERATIONS' , false ) ; tf . add ( c , 1 ) ; tf . env ( ) . set ( 'WEBGL_PACK_BINARY_OPERATIONS' , webglPackBinaryOperationsFlagSaved ) ; expect ( tf . memory ( ) . numBytes - startNumBytes ) . toEqual ( 16 ) ; expect ( tf . memory ( ) . numTensors - startNumTensors ) . toEqual ( 1 ) ; expect ( ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPU - startNumBytesInGPU ) . toEqual ( 4 * tf . util . bytesPerElement ( a . dtype ) ) ; } ) ; it ( 'should not leak memory when lazily packing' , ( ) => { const a = tf . tensor2d ( [ 1 , 2 , 3 , 4 , 5 , 6 ] , [ 2 , 3 ] ) ; const b = tf . tensor2d ( [ 0 , 1 , - 3 , 2 , 2 , 1 ] , [ 3 , 2 ] ) ; const c = tf . add ( a , 1 ) ; const startNumBytes = tf . memory ( ) . numBytes ; const startNumTensors = tf . memory ( ) . numTensors ; const startNumBytesInGPU = ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPU ; tf . matMul ( b , c ) ; expect ( tf . memory ( ) . numBytes - startNumBytes ) . toEqual ( 36 ) ; expect ( tf . memory ( ) . numTensors - startNumTensors ) . toEqual ( 1 ) ; expect ( ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPU - startNumBytesInGPU ) . toEqual ( 24 * tf . util . bytesPerElement ( a . dtype ) ) ; } ) ; it ( 'should work when the same input must be represented by' + 'different textures' , async ( ) => { const a = tf . tensor1d ( [ 1 , 2 ] ) ; const res = tf . dot ( a , a ) ; expectArraysClose ( await res . data ( ) , [ 5 ] ) ; } ) ; } ) ; describeWithFlags ( 'backendWebGL' , WEBGL_ENVS , ( ) => { let prevBackend ; beforeAll ( ( ) => { prevBackend = tf . getBackend ( ) ; } ) ; afterEach ( ( ) => { tf . setBackend ( prevBackend ) ; tf . removeBackend ( 'test-storage' ) ; } ) ; it ( 'register string tensor with values' , ( ) => { const backend = new MathBackendWebGL ( ) ; tf . registerBackend ( 'test-storage' , ( ) => backend ) ; tf . setBackend ( 'test-storage' ) ; const t = engine ( ) . makeTensor ( [ 'a' , 'b' , 'c' ] , [ 3 ] , 'string' ) ; expectArraysEqual ( decodeStrings ( backend . readSync ( t . dataId ) as Uint8Array [ ] ) , [ 'a' , 'b' , 'c' ] ) ; } ) ; it ( 'register string tensor with values and wrong shape throws error' , ( ) => { const backend = new MathBackendWebGL ( ) ; tf . registerBackend ( 'test-storage' , ( ) => backend ) ; tf . setBackend ( 'test-storage' ) ; expect ( ( ) => tf . tensor ( [ 'a' , 'b' , 'c' ] , [ 4 ] , 'string' ) ) . toThrowError ( ) ; } ) ; it ( 'reading' , ( ) => { const backend = new MathBackendWebGL ( null ) ; tf . registerBackend ( 'test-storage' , ( ) => backend ) ; tf . setBackend ( 'test-storage' ) ; const texManager = backend . getTextureManager ( ) ; const t = engine ( ) . makeTensor ( new Float32Array ( [ 1 , 2 , 3 ] ) , [ 3 ] , 'float32' ) ; expect ( texManager . getNumUsedTextures ( ) ) . toBe ( 0 ) ; backend . getTexture ( t . dataId ) ; expect ( texManager . getNumUsedTextures ( ) ) . toBe ( 1 ) ; expectArraysClose ( backend . readSync ( t . dataId ) as Float32Array , new Float32Array ( [ 1 , 2 , 3 ] ) ) ; expect ( texManager . getNumUsedTextures ( ) ) . toBe ( 0 ) ; backend . getTexture ( t . dataId ) ; expect ( texManager . getNumUsedTextures ( ) ) . toBe ( 1 ) ; backend . disposeData ( t . dataId ) ; expect ( texManager . getNumUsedTextures ( ) ) . toBe ( 0 ) ; } ) ; it ( 'read packed and then use by an unpacked op' , async ( ) => { const backend = new MathBackendWebGL ( null ) ; tf . registerBackend ( 'test-storage' , ( ) => backend ) ; tf . copyRegisteredKernels ( 'webgl' , 'test-storage' ) ; tf . setBackend ( 'test-storage' ) ; const webglPackFlagSaved = tf . env ( ) . getBool ( 'WEBGL_PACK' ) ; tf . env ( ) . set ( 'WEBGL_PACK' , true ) ; const webglSizeUploadUniformSaved = tf . env ( ) . getNumber ( 'WEBGL_SIZE_UPLOAD_UNIFORM' ) ; tf . env ( ) . set ( 'WEBGL_SIZE_UPLOAD_UNIFORM' , 0 ) ; const a = tf . tensor2d ( [ 1 , 2 ] , [ 2 , 1 ] ) ; const b = tf . tensor2d ( [ 1 ] , [ 1 , 1 ] ) ; const c = tf . matMul ( a , b ) ; backend . readSync ( c . dataId ) ; tf . env ( ) . set ( 'WEBGL_PACK' , false ) ; const d = tf . add ( c , 1 ) ; tf . env ( ) . set ( 'WEBGL_PACK' , webglPackFlagSaved ) ; tf . env ( ) . set ( 'WEBGL_SIZE_UPLOAD_UNIFORM' , webglSizeUploadUniformSaved ) ; expectArraysClose ( await d . data ( ) , [ 2 , 3 ] ) ; } ) ; it ( 'delayed storage, overwriting' , ( ) => { const backend = new MathBackendWebGL ( null ) ; tf . registerBackend ( 'test-storage' , ( ) => backend ) ; tf . setBackend ( 'test-storage' ) ; const texManager = backend . getTextureManager ( ) ; const t = engine ( ) . makeTensor ( new Float32Array ( [ 1 , 2 , 3 ] ) , [ 3 ] , 'float32' ) ; backend . getTexture ( t . dataId ) ; expect ( texManager . getNumUsedTextures ( ) ) . toBe ( 1 ) ; expectArraysClose ( backend . readSync ( t . dataId ) as Float32Array , new Float32Array ( [ 1 , 2 , 3 ] ) ) ; backend . getTexture ( t . dataId ) ; expect ( texManager . getNumUsedTextures ( ) ) . toBe ( 1 ) ; expectArraysClose ( backend . readSync ( t . dataId ) as Float32Array , new Float32Array ( [ 1 , 2 , 3 ] ) ) ; expect ( texManager . getNumUsedTextures ( ) ) . toBe ( 0 ) ; } ) ; } ) ; describeWithFlags ( 'Webgl backend disposal' , WEBGL_ENVS , ( ) => { it ( 'register and dispose a backend outside unit test' , ( ) => { tf . ENV . set ( 'IS_TEST' , false ) ; const backend = new MathBackendWebGL ( ) ; tf . registerBackend ( 'test-disposal' , ( ) => backend ) ; tf . copyRegisteredKernels ( 'webgl' , 'test-disposal' ) ; tf . setBackend ( 'test-disposal' ) ; tf . zeros ( [ 1000 ] ) . sqrt ( ) . dataSync ( ) ; tf . backend ( ) . dispose ( ) ; const cache = getBinaryCache ( tf . ENV . getNumber ( 'WEBGL_VERSION' ) ) ; expect ( Object . keys ( cache ) . length ) . toBe ( 0 ) ; tf . removeBackend ( 'test-disposal' ) ; } ) ; it ( 'register and dispose a backend inside unit test' , ( ) => { tf . ENV . set ( 'IS_TEST' , true ) ; const backend = new MathBackendWebGL ( ) ; tf . registerBackend ( 'test-disposal' , ( ) => backend ) ; tf . copyRegisteredKernels ( 'webgl' , 'test-disposal' ) ; tf . setBackend ( 'test-disposal' ) ; tf . zeros ( [ 1000 ] ) . sqrt ( ) . dataSync ( ) ; tf . backend ( ) . dispose ( ) ; const cache = getBinaryCache ( tf . ENV . getNumber ( 'WEBGL_VERSION' ) ) ; expect ( Object . keys ( cache ) . length ) . toBeGreaterThan ( 0 ) ; tf . removeBackend ( 'test-disposal' ) ; } ) ; it ( 'register, dispose and re-register a backend outside unit test' , ( ) => { tf . ENV . set ( 'IS_TEST' , false ) ; tf . registerBackend ( 'test-disposal' , ( ) => new MathBackendWebGL ( ) ) ; tf . copyRegisteredKernels ( 'webgl' , 'test-disposal' ) ; tf . setBackend ( 'test-disposal' ) ; tf . zeros ( [ 1000 ] ) . sqrt ( ) . dataSync ( ) ; tf . backend ( ) . dispose ( ) ; tf . removeBackend ( 'test-disposal' ) ; tf . registerBackend ( 'test-disposal' , ( ) => new MathBackendWebGL ( ) ) ; tf . copyRegisteredKernels ( 'webgl' , 'test-disposal' ) ; tf . setBackend ( 'test-disposal' ) ; tf . zeros ( [ 1000 ] ) . sqrt ( ) . dataSync ( ) ; tf . backend ( ) . dispose ( ) ; const cache = getBinaryCache ( tf . ENV . getNumber ( 'WEBGL_VERSION' ) ) ; expect ( Object . keys ( cache ) . length ) . toBe ( 0 ) ; tf . removeBackend ( 'test-disposal' ) ; } ) ; } ) ; describeWithFlags ( 'Custom window size' , WEBGL_ENVS , ( ) => { const customBackendName = 'custom-webgl' ; beforeAll ( ( ) => { const kernelFunc = tf . getKernel ( 'Square' , 'webgl' ) . kernelFunc ; tf . registerKernel ( { kernelName : 'Square' , backendName : customBackendName , kernelFunc } ) ; } ) ; afterAll ( ( ) => { tf . unregisterKernel ( 'Square' , customBackendName ) ; } ) ; it ( 'Set screen area to be 1x1' , ( ) => { spyOnProperty ( window , 'screen' , 'get' ) . and . returnValue ( { height : 1 , width : 1 } ) ; tf . registerBackend ( customBackendName , ( ) => new MathBackendWebGL ( ) ) ; tf . setBackend ( customBackendName ) ; const a = tf . ones ( [ 100 , 100 ] ) ; expect ( ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPU ) . toBe ( 0 ) ; let numWarnCalls = 0 ; spyOn ( console , 'warn' ) . and . callFake ( ( ) => { numWarnCalls ++ ; } ) ; a . square ( ) ; expect ( numWarnCalls ) . toBe ( 1 ) ; expect ( ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPU ) . toBe ( 100 * 100 * 4 * 2 ) ; a . square ( ) ; expect ( numWarnCalls ) . toBe ( 1 ) ; expect ( ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPU ) . toBe ( 100 * 100 * 4 * 3 ) ; tf . removeBackend ( customBackendName ) ; } ) ; } ) ; const SIZE_UPLOAD_UNIFORM = 4 ; const FLOAT32_WEBGL_ENVS = { flags : { 'WEBGL_RENDER_FLOAT32_ENABLED' : true } , predicate : WEBGL_ENVS . predicate } ; describeWithFlags ( 'upload tensors as uniforms' , FLOAT32_WEBGL_ENVS , ( ) => { let savedUploadUniformValue ; beforeAll ( ( ) => { savedUploadUniformValue = tf . env ( ) . get ( 'WEBGL_SIZE_UPLOAD_UNIFORM' ) as number ; tf . env ( ) . set ( 'WEBGL_SIZE_UPLOAD_UNIFORM' , SIZE_UPLOAD_UNIFORM ) ; } ) ; afterAll ( ( ) => { tf . env ( ) . set ( 'WEBGL_SIZE_UPLOAD_UNIFORM' , savedUploadUniformValue ) ; } ) ; it ( 'small tensor gets uploaded as scalar' , ( ) => { let m = tf . memory ( ) as WebGLMemoryInfo ; expect ( m . numBytesInGPU ) . toBe ( 0 ) ; const a = tf . zeros ( [ SIZE_UPLOAD_UNIFORM - 1 ] ) ; a . square ( ) ; m = tf . memory ( ) as WebGLMemoryInfo ; expect ( m . numBytesInGPU ) . toBe ( a . size * 4 ) ; } ) ; it ( 'large tensor gets uploaded to gpu' , ( ) => { let m = tf . memory ( ) as WebGLMemoryInfo ; expect ( m . numBytesInGPU ) . toBe ( 0 ) ; const a = tf . zeros ( [ SIZE_UPLOAD_UNIFORM + 1 ] ) ; a . square ( ) ; m = tf . memory ( ) as WebGLMemoryInfo ; expect ( m . numBytesInGPU ) . toBe ( a . size * 4 * 2 ) ; } ) ; it ( 'download and re-upload an output of a shader' , async ( ) => { const vals = new Float32Array ( SIZE_UPLOAD_UNIFORM + 1 ) ; vals . fill ( 2 ) ; const a = tf . square ( vals ) ; a . dataSync ( ) ; const res = a . square ( ) ; const expected = new Float32Array ( SIZE_UPLOAD_UNIFORM + 1 ) ; expected . fill ( 16 ) ; expectArraysClose ( await res . data ( ) , expected ) ; } ) ; } ) ; describeWithFlags ( 'indexing for large tensors' , FLOAT32_WEBGL_ENVS , ( ) => { it ( 'properly indexes large tensors' , async ( ) => { const range = 3000 * 3000 ; const aData = new Float32Array ( range ) ; for ( let i = 0 ; i < range ; i ++ ) { aData [ i ] = i / range ; } const a = tf . tensor1d ( aData ) ; const aRelu = a . relu ( ) ; expectArraysClose ( await a . data ( ) , aData ) ; expectArraysClose ( await aRelu . data ( ) , aData ) ; } ) ; } ) ; describeWithFlags ( 'debug on webgl' , WEBGL_ENVS , ( ) => { beforeAll ( ( ) => { spyOn ( console , 'warn' ) ; tf . enableDebugMode ( ) ; } ) ; it ( 'debug mode errors when overflow in tensor construction' , ( ) => { const savedRenderFloat32Flag = tf . env ( ) . getBool ( 'WEBGL_RENDER_FLOAT32_ENABLED' ) ; tf . env ( ) . set ( 'WEBGL_RENDER_FLOAT32_ENABLED' , false ) ; const a = ( ) => tf . tensor1d ( [ 2 , Math . pow ( 2 , 17 ) ] , 'float32' ) ; expect ( a ) . toThrowError ( ) ; tf . env ( ) . set ( 'WEBGL_RENDER_FLOAT32_ENABLED' , savedRenderFloat32Flag ) ; } ) ; it ( 'debug mode errors when underflow in tensor construction' , ( ) => { const savedRenderFloat32Flag = tf . env ( ) . getBool ( 'WEBGL_RENDER_FLOAT32_ENABLED' ) ; tf . env ( ) . set ( 'WEBGL_RENDER_FLOAT32_ENABLED' , false ) ; const a = ( ) => tf . tensor1d ( [ 2 , 1e-8 ] , 'float32' ) ; expect ( a ) . toThrowError ( ) ; tf . env ( ) . set ( 'WEBGL_RENDER_FLOAT32_ENABLED' , savedRenderFloat32Flag ) ; } ) ; } ) ; const WEBGL1_ENVS = { flags : { 'WEBGL_VERSION' : 1 } , predicate : WEBGL_ENVS . predicate } ; const WEBGL2_ENVS = { flags : { 'WEBGL_VERSION' : 2 } , predicate : WEBGL_ENVS . predicate } ; describeWithFlags ( 'computeBytes counts bytes correctly' , WEBGL1_ENVS , ( ) => { it ( 'for all physical texture types' , ( ) => { const gpgpu = new GPGPUContext ( ) ; const shapeRC = [ 2 , 3 ] ; let bytesForTex = computeBytes ( shapeRC , PhysicalTextureType . UNPACKED_FLOAT16 , gpgpu . gl , gpgpu . textureConfig , false ) ; expect ( bytesForTex ) . toBe ( 96 ) ; bytesForTex = computeBytes ( shapeRC , PhysicalTextureType . UNPACKED_FLOAT32 , gpgpu . gl , gpgpu . textureConfig , false ) ; expect ( bytesForTex ) . toBe ( 96 ) ; bytesForTex = computeBytes ( shapeRC , PhysicalTextureType . PACKED_4X1_UNSIGNED_BYTE , gpgpu . gl , gpgpu . textureConfig , true ) ; expect ( bytesForTex ) . toBe ( 32 ) ; bytesForTex = computeBytes ( shapeRC , PhysicalTextureType . PACKED_2X2_FLOAT32 , gpgpu . gl , gpgpu . textureConfig , true ) ; expect ( bytesForTex ) . toBe ( 32 ) ; bytesForTex = computeBytes ( shapeRC , PhysicalTextureType . PACKED_2X2_FLOAT16 , gpgpu . gl , gpgpu . textureConfig , true ) ; expect ( bytesForTex ) . toBe ( 32 ) ; gpgpu . dispose ( ) ; } ) ; } ) ; describeWithFlags ( 'computeBytes counts bytes correctly' , WEBGL2_ENVS , ( ) => { it ( 'test every physical tex type input to computeBytes' , ( ) => { const gpgpu = new GPGPUContext ( ) ; const shapeRC = [ 2 , 3 ] ; let bytesForTex = computeBytes ( shapeRC , PhysicalTextureType . UNPACKED_FLOAT16 , gpgpu . gl , gpgpu . textureConfig , false ) ; expect ( bytesForTex ) . toBe ( 12 ) ; bytesForTex = computeBytes ( shapeRC , PhysicalTextureType . UNPACKED_FLOAT32 , gpgpu . gl , gpgpu . textureConfig , false ) ; expect ( bytesForTex ) . toBe ( 24 ) ; bytesForTex = computeBytes ( shapeRC , PhysicalTextureType . PACKED_4X1_UNSIGNED_BYTE , gpgpu . gl , gpgpu . textureConfig , true ) ; expect ( bytesForTex ) . toBe ( 8 ) ; bytesForTex = computeBytes ( shapeRC , PhysicalTextureType . PACKED_2X2_FLOAT32 , gpgpu . gl , gpgpu . textureConfig , true ) ; expect ( bytesForTex ) . toBe ( 32 ) ; bytesForTex = computeBytes ( shapeRC , PhysicalTextureType . PACKED_2X2_FLOAT16 , gpgpu . gl , gpgpu . textureConfig , true ) ; expect ( bytesForTex ) . toBe ( 16 ) ; gpgpu . dispose ( ) ; } ) ; } ) ; describeWithFlags ( 'aggressive texture deletion' , WEBGL_ENVS , ( ) => { it ( 'basic' , ( ) => { const savedDeleteThreshold = tf . env ( ) . get ( 'WEBGL_DELETE_TEXTURE_THRESHOLD' ) as number ; tf . env ( ) . set ( 'WEBGL_DELETE_TEXTURE_THRESHOLD' , 0 ) ; const a = tf . tensor2d ( [ 1 , 2 , 3 , 4 , 5 , 6 ] , [ 2 , 3 ] ) ; const b = tf . tensor2d ( [ 0 , 1 , - 3 , 2 , 2 , 1 ] , [ 3 , 2 ] ) ; tf . matMul ( a , b ) ; const startNumBytesAllocated = ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPUAllocated ; a . dispose ( ) ; b . dispose ( ) ; expect ( startNumBytesAllocated - ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPUAllocated ) . toBeGreaterThan ( 0 ) ; tf . env ( ) . set ( 'WEBGL_DELETE_TEXTURE_THRESHOLD' , savedDeleteThreshold ) ; } ) ; } ) ; describeWithFlags ( 'memory webgl' , WEBGL_ENVS , ( ) => { it ( 'unreliable is falsy/not present when all tensors are numeric' , ( ) => { tf . tensor ( 1 ) ; const mem = tf . memory ( ) ; expect ( mem . numTensors ) . toBe ( 1 ) ; expect ( mem . numDataBuffers ) . toBe ( 1 ) ; expect ( mem . numBytes ) . toBe ( 4 ) ; expect ( mem . unreliable ) . toBeFalsy ( ) ; } ) ; } ) ; describeWithFlags ( 'manual gl flush' , WEBGL_ENVS , ( ) => { it ( 'works when manual gl flush is enabled' , async ( ) => { const savedGlThreshold = tf . env ( ) . get ( 'WEBGL_FLUSH_THRESHOLD' ) as number ; tf . env ( ) . set ( 'WEBGL_FLUSH_THRESHOLD' , 1 ) ; const a = tf . tensor2d ( [ 1 , 2 , 3 , 4 , 5 , 6 ] , [ 2 , 3 ] ) ; const b = tf . tensor2d ( [ 1 , 1 , - 3 , 2 , 2 , 1 ] , [ 2 , 3 ] ) ; const result = tf . div ( tf . div ( tf . mul ( a , b ) , a ) , b ) ; expectArraysClose ( await result . data ( ) , [ 1 , 1 , 1 , 1 , 1 , 1 ] ) ; tf . env ( ) . set ( 'WEBGL_FLUSH_THRESHOLD' , savedGlThreshold ) ; } ) ; } ) ; describeWithFlags ( 'backend without render float32 support' , WEBGL_ENVS , ( ) => { const savedRenderFloat32Flag = tf . env ( ) . getBool ( 'WEBGL_RENDER_FLOAT32_ENABLED' ) ; const customWebGLBackendName = 'half-float-webgl' ; beforeAll ( ( ) => { tf . env ( ) . set ( 'WEBGL_RENDER_FLOAT32_ENABLED' , false ) ; } ) ; beforeEach ( ( ) => { tf . copyRegisteredKernels ( 'webgl' , customWebGLBackendName ) ; tf . registerBackend ( customWebGLBackendName , ( ) => new MathBackendWebGL ( null ) ) ; } ) ; afterEach ( ( ) => { tf . removeBackend ( customWebGLBackendName ) ; } ) ; afterAll ( ( ) => { tf . env ( ) . set ( 'WEBGL_RENDER_FLOAT32_ENABLED' , savedRenderFloat32Flag ) ; } ) ; it ( 'basic usage' , async ( ) => { tf . setBackend ( customWebGLBackendName ) ; const a = tf . tensor2d ( [ 1 , 2 ] , [ 1 , 2 ] ) ; const b = tf . tensor2d ( [ 1 , 2 ] , [ 1 , 2 ] ) ; const c = tf . add ( a , b ) ; expectArraysClose ( await c . data ( ) , [ 2 , 4 ] ) ; } ) ; it ( 'disposing tensors should not cause errors' , ( ) => { tf . setBackend ( customWebGLBackendName ) ; expect ( ( ) => tf . tidy ( ( ) => { const a = tf . tensor2d ( [ 1 , 2 ] , [ 1 , 2 ] ) ; const b = tf . tensor2d ( [ 1 , 2 ] , [ 1 , 2 ] ) ; const c = tf . add ( a , b ) ; c . dataSync ( ) ; return c . add ( tf . tensor2d ( [ 2 , 4 ] , [ 1 , 2 ] ) ) ; } ) ) . not . toThrowError ( ) ; } ) ; } ) ; describeWithFlags ( 'time webgl' , WEBGL_ENVS , ( ) => { it ( 'upload + compute' , async ( ) => { const a = tf . zeros ( [ 10 , 10 ] ) ; const time = await tf . time ( ( ) => a . square ( ) ) as WebGLTimingInfo ; expect ( time . uploadWaitMs > 0 ) ; expect ( time . downloadWaitMs === 0 ) ; expect ( time . kernelMs > 0 ) ; expect ( time . wallMs >= time . kernelMs ) ; } ) ; it ( 'upload + compute + dataSync' , async ( ) => { const a = tf . zeros ( [ 10 , 10 ] ) ; const time = await tf . time ( ( ) => a . square ( ) . dataSync ( ) ) as WebGLTimingInfo ; expect ( time . uploadWaitMs > 0 ) ; expect ( time . downloadWaitMs > 0 ) ; expect ( time . kernelMs > 0 ) ; expect ( time . wallMs >= time . kernelMs ) ; } ) ; it ( 'upload + compute + data' , async ( ) => { const a = tf . zeros ( [ 10 , 10 ] ) ; const time = await tf . time ( async ( ) => a . square ( ) . data ( ) ) as WebGLTimingInfo ; expect ( time . uploadWaitMs > 0 ) ; expect ( time . downloadWaitMs > 0 ) ; expect ( time . kernelMs > 0 ) ; expect ( time . wallMs >= time . kernelMs ) ; } ) ; it ( 'preupload (not included) + compute + data' , async ( ) => { const a = tf . zeros ( [ 10 , 10 ] ) ; a . square ( ) ; const time = await tf . time ( ( ) => a . sqrt ( ) ) as WebGLTimingInfo ; expect ( time . uploadWaitMs === 0 ) ; expect ( time . downloadWaitMs === 0 ) ; expect ( time . kernelMs > 0 ) ; expect ( time . wallMs >= time . kernelMs ) ; } ) ; it ( 'returns error for kernelMs if query timer extension is unavailable' , async ( ) => { const savedQueryReliableValue = tf . env ( ) . get ( 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE' ) ; tf . env ( ) . set ( 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE' , false ) ; const a = tf . zeros ( [ 10 , 10 ] ) ; const time = await tf . backend ( ) . time ( ( ) => a . sqrt ( ) ) as WebGLTimingInfo ; expect ( time . kernelMs ) . toEqual ( { error : 'WebGL query timers are not supported in this environment.' } ) ; tf . env ( ) . set ( 'WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE' , savedQueryReliableValue ) ; } ) ; } ) ; describeWithFlags ( 'keeping data on gpu ' , WEBGL2_ENVS , ( ) => { let flag ; beforeAll ( ( ) => { flag = tf . env ( ) . getBool ( 'WEBGL_CPU_FORWARD' ) ; tf . env ( ) . set ( 'WEBGL_CPU_FORWARD' , false ) ; } ) ; afterAll ( ( ) => { tf . env ( ) . set ( 'WEBGL_CPU_FORWARD' , flag ) ; } ) ; it ( 'has a valid texture for dtype=float32.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const texShape = [ 2 , 2 ] ; const a = tf . tensor ( data , [ 1 , 3 , 4 ] ) ; const b = tf . add ( a , 0 ) ; const res = b . dataToGPU ( ) ; expectArraysEqual ( res . texShape , texShape ) ; const webGLBackend = tf . backend ( ) as MathBackendWebGL ; const buffer = webGLBackend . gpgpu . createBufferFromTexture ( res . texture , res . texShape [ 0 ] , res . texShape [ 1 ] ) ; const vals = webGLBackend . gpgpu . downloadFloat32MatrixFromBuffer ( buffer , 12 ) ; expectArraysEqual ( vals , data ) ; } ) ; it ( 'uses user defined texShape.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const a = tf . tensor ( data , [ 1 , 3 , 4 ] ) ; const b = tf . add ( a , 0 ) ; const texShape = [ 1 , 3 ] as [ number , number ] ; const res = b . dataToGPU ( { customTexShape : texShape } ) ; expectArraysEqual ( res . texShape , texShape ) ; const webGLBackend = tf . backend ( ) as MathBackendWebGL ; const buffer = webGLBackend . gpgpu . createBufferFromTexture ( res . texture , res . texShape [ 0 ] , res . texShape [ 1 ] ) ; const vals = webGLBackend . gpgpu . downloadFloat32MatrixFromBuffer ( buffer , 12 ) ; expectArraysEqual ( vals , data ) ; } ) ; it ( 'has a valid texture for dtype=int32.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const texShape = [ 2 , 2 ] ; const a = tf . tensor ( data , [ 1 , 3 , 4 ] , 'int32' ) ; const b = tf . add ( a , 0 ) ; const res = b . dataToGPU ( ) ; expectArraysEqual ( res . texShape , texShape ) ; const webGLBackend = tf . backend ( ) as MathBackendWebGL ; const buffer = webGLBackend . gpgpu . createBufferFromTexture ( res . texture , res . texShape [ 0 ] , res . texShape [ 1 ] ) ; const vals = webGLBackend . gpgpu . downloadFloat32MatrixFromBuffer ( buffer , 12 ) ; expectArraysEqual ( vals , data ) ; } ) ; it ( 'has no memory leak.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const a = tf . tensor ( data , [ 1 , 3 , 4 ] ) ; const b = tf . add ( a , 0 ) ; const webGLBackend = tf . backend ( ) as MathBackendWebGL ; const startNumBytes = ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPU ; const startTensor = tf . memory ( ) . numTensors ; const startDataBuckets = webGLBackend . numDataIds ( ) ; const res = b . dataToGPU ( ) ; res . tensorRef . dispose ( ) ; const endNumBytes = ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPU ; const endTensor = tf . memory ( ) . numTensors ; const endDataBuckets = webGLBackend . numDataIds ( ) ; expect ( endNumBytes ) . toEqual ( startNumBytes ) ; expect ( endTensor ) . toEqual ( startTensor ) ; expect ( endDataBuckets ) . toEqual ( startDataBuckets ) ; } ) ; it ( 'can be used in tidy.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const webGLBackend = tf . backend ( ) as MathBackendWebGL ; const startTensor = tf . memory ( ) . numTensors ; const startDataBuckets = webGLBackend . numDataIds ( ) ; const result = tf . tidy ( ( ) => { const a = tf . tensor ( data , [ 1 , 3 , 4 ] ) ; const b = tf . add ( a , 0 ) ; return b . dataToGPU ( ) as { } as tf . Tensor ; } ) ; const endTensor = tf . memory ( ) . numTensors ; const endDataBuckets = webGLBackend . numDataIds ( ) ; expect ( endTensor ) . toEqual ( startTensor + 1 ) ; expect ( endDataBuckets ) . toEqual ( startDataBuckets + 1 ) ; const res = result as { } as GPUData ; const buffer = webGLBackend . gpgpu . createBufferFromTexture ( res . texture , res . texShape [ 0 ] , res . texShape [ 1 ] ) ; const vals = webGLBackend . gpgpu . downloadFloat32MatrixFromBuffer ( buffer , 12 ) ; expectArraysEqual ( vals , data ) ; } ) ; it ( 'tidy has no memory leak.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const webGLBackend = tf . backend ( ) ; const startTensor = tf . memory ( ) . numTensors ; const startDataBuckets = webGLBackend . numDataIds ( ) ; tf . tidy ( ( ) => { const a = tf . tensor ( data , [ 1 , 3 , 4 ] ) ; const b = tf . add ( a , 0 ) ; b . dataToGPU ( ) ; return b ; } ) ; const endTensor = tf . memory ( ) . numTensors ; const endDataBuckets = webGLBackend . numDataIds ( ) ; expect ( endTensor ) . toEqual ( startTensor + 1 ) ; expect ( endDataBuckets ) . toEqual ( startDataBuckets + 1 ) ; } ) ; it ( 'throws error when user defined texShape is too small.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const a = tf . tensor ( data , [ 1 , 3 , 4 ] ) ; const b = tf . add ( a , 0 ) ; expect ( ( ) => { b . dataToGPU ( { customTexShape : [ 1 , 1 ] } ) ; } ) . toThrowError ( ) ; } ) ; } ) ; describeWithFlags ( 'keeping data on gpu ' , WEBGL1_ENVS , ( ) => { let flag ; const webGLBackend = ( tf . backend ( ) as MathBackendWebGL ) ; beforeAll ( ( ) => { flag = tf . env ( ) . getBool ( 'WEBGL_CPU_FORWARD' ) ; tf . env ( ) . set ( 'WEBGL_CPU_FORWARD' , false ) ; } ) ; afterAll ( ( ) => { tf . env ( ) . set ( 'WEBGL_CPU_FORWARD' , flag ) ; } ) ; it ( 'has a valid texture for dtype=float32.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const texShape = [ 2 , 2 ] ; const size = 12 ; const a = tf . tensor ( data , [ 1 , 3 , 4 ] ) ; const b = tf . add ( a , 0 ) ; const res = b . dataToGPU ( ) ; expectArraysEqual ( res . texShape , texShape ) ; if ( tf . env ( ) . getBool ( 'WEBGL_DOWNLOAD_FLOAT_ENABLED' ) ) { const tmpData = webGLBackend . texData . get ( res . tensorRef . dataId ) ; const vals = webGLBackend . gpgpu . downloadMatrixFromPackedTexture ( tmpData . texture . texture , ... tmpData . texture . texShape ) . subarray ( 0 , size ) ; expectArraysEqual ( vals , data ) ; } } ) ; it ( 'uses user defined texShape.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const a = tf . tensor ( data , [ 1 , 3 , 4 ] ) ; const b = tf . add ( a , 0 ) ; const texShape = [ 1 , 3 ] as [ number , number ] ; const size = 12 ; const res = b . dataToGPU ( { customTexShape : texShape } ) ; expectArraysEqual ( res . texShape , texShape ) ; if ( tf . env ( ) . getBool ( 'WEBGL_DOWNLOAD_FLOAT_ENABLED' ) ) { const tmpData = webGLBackend . texData . get ( res . tensorRef . dataId ) ; const vals = webGLBackend . gpgpu . downloadMatrixFromPackedTexture ( tmpData . texture . texture , ... tmpData . texture . texShape ) . subarray ( 0 , size ) ; expectArraysEqual ( vals , data ) ; } } ) ; it ( 'has a valid texture for dtype=int32.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const texShape = [ 2 , 2 ] ; const size = 12 ; const a = tf . tensor ( data , [ 1 , 3 , 4 ] , 'int32' ) ; const b = tf . add ( a , 0 ) ; const res = b . dataToGPU ( ) ; expectArraysEqual ( res . texShape , texShape ) ; if ( tf . env ( ) . getBool ( 'WEBGL_DOWNLOAD_FLOAT_ENABLED' ) ) { const tmpData = webGLBackend . texData . get ( res . tensorRef . dataId ) ; const vals = webGLBackend . gpgpu . downloadMatrixFromPackedTexture ( tmpData . texture . texture , ... tmpData . texture . texShape ) . subarray ( 0 , size ) ; expectArraysEqual ( vals , data ) ; } } ) ; it ( 'has no memory leak.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const a = tf . tensor ( data , [ 1 , 3 , 4 ] ) ; const b = tf . add ( a , 0 ) ; const webGLBackend = tf . backend ( ) as MathBackendWebGL ; const startNumBytes = ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPU ; const startTensor = tf . memory ( ) . numTensors ; const startDataBuckets = webGLBackend . numDataIds ( ) ; const res = b . dataToGPU ( ) ; res . tensorRef . dispose ( ) ; const endNumBytes = ( tf . memory ( ) as WebGLMemoryInfo ) . numBytesInGPU ; const endTensor = tf . memory ( ) . numTensors ; const endDataBuckets = webGLBackend . numDataIds ( ) ; expect ( endNumBytes ) . toEqual ( startNumBytes ) ; expect ( endTensor ) . toEqual ( startTensor ) ; expect ( endDataBuckets ) . toEqual ( startDataBuckets ) ; } ) ; it ( 'throws error when user defined texShape is too small.' , ( ) => { const data = [ 1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 ] ; const a = tf . tensor ( data , [ 1 , 3 , 4 ] ) ; const b = tf . add ( a , 0 ) ; expect ( ( ) => { b . dataToGPU ( { customTexShape : [ 1 , 1 ] } ) ; } ) . toThrowError ( ) ; } ) ; } ) ; describeWithFlags ( 'caching on cpu' , WEBGL_ENVS , ( ) => { const customBackendName = 'cache-on-cpu' ; beforeAll ( ( ) => { tf . env ( ) . set ( 'WEBGL_CPU_FORWARD' , false ) ; const kernelFunc = tf . getKernel ( 'Square' , 'webgl' ) . kernelFunc ; tf . registerKernel ( { kernelName : 'Square' , backendName : customBackendName , kernelFunc } ) ; } ) ; afterAll ( ( ) => { tf . unregisterKernel ( 'Square' , customBackendName ) ; } ) ; it ( 'caches on cpu after async read' , async ( ) => { const backend = new MathBackendWebGL ( ) ; tf . registerBackend ( customBackendName , ( ) => backend ) ; tf . setBackend ( customBackendName ) ; const t = tf . square ( 2 ) ; const info = backend . getDataInfo ( t . dataId ) ; expect ( info . values == null ) . toBe ( true ) ; await t . data ( ) ; expect ( info . values ) . not . toBe ( null ) ; tf . removeBackend ( customBackendName ) ; } ) ; it ( 'caches on cpu after sync read' , ( ) => { const backend = new MathBackendWebGL ( ) ; tf . registerBackend ( customBackendName , ( ) => backend ) ; tf . setBackend ( customBackendName ) ; const t = tf . square ( 2 ) ; const info = backend . getDataInfo ( t . dataId ) ; expect ( info . values == null ) . toBe ( true ) ; t . dataSync ( ) ; expect ( info . values ) . not . toBe ( null ) ; tf . removeBackend ( customBackendName ) ; } ) ; } ) ; describeWithFlags ( 'WebGL backend has sync init' , WEBGL_ENVS , ( ) => { it ( 'can do matmul without waiting for ready' , async ( ) => { const customWebGLBackendName = 'my-webgl' ; tf . copyRegisteredKernels ( 'webgl' , customWebGLBackendName ) ; tf . registerBackend ( customWebGLBackendName , ( ) => { return new MathBackendWebGL ( ) ; } ) ; tf . setBackend ( customWebGLBackendName ) ; const a = tf . tensor1d ( [ 5 ] ) ; const b = tf . tensor1d ( [ 3 ] ) ; const res = tf . dot ( a , b ) ; expectArraysClose ( await res . data ( ) , 15 ) ; tf . dispose ( [ a , b , res ] ) ; tf . removeBackend ( customWebGLBackendName ) ; } ) ; } ) ; describeWithFlags ( 'custom canvas ' , WEBGL_ENVS , ( ) => { const customBackendName = 'custom-webgl' ; let flag ; beforeAll ( ( ) => { flag = tf . env ( ) . getBool ( 'WEBGL_CPU_FORWARD' ) ; tf . env ( ) . set ( 'WEBGL_CPU_FORWARD' , false ) ; const kernelFunc = tf . getKernel ( 'Square' , 'webgl' ) . kernelFunc ; tf . registerKernel ( { kernelName : 'Square' , backendName : customBackendName , kernelFunc } ) ; } ) ; afterAll ( ( ) => { tf . env ( ) . set ( 'WEBGL_CPU_FORWARD' , flag ) ; tf . unregisterKernel ( 'Square' , customBackendName ) ; } ) ; it ( 'works.' , ( ) => { const customCanvas = document . createElement ( 'canvas' ) ; customCanvas . width = 300 ; customCanvas . height = 200 ; const backend = new MathBackendWebGL ( customCanvas ) ; tf . registerBackend ( customBackendName , ( ) => backend ) ; tf . setBackend ( customBackendName ) ; const t = tf . square ( 2 ) ; const data = t . dataSync ( ) ; expectArraysEqual ( data , [ 4 ] ) ; tf . removeBackend ( customBackendName ) ; } ) ; } ) ;