license import { backend_util , BinaryInputs , DataType , env , KernelFunc , TypedArray , UnaryInputs , upcastType } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { BinaryOpProgram } from '../binaryop_gpu' ; import { BinaryOpPackedProgram } from '../binaryop_packed_gpu' ; import { complex } from '../kernels/Complex' ; import { LEAKYRELU , LEAKYRELU_PACKED } from '../kernels/LeakyRelu' ; import { PRELU , PRELU_PACKED } from '../kernels/Prelu' ; import * as unary_op from '../unaryop_gpu' ; import { UnaryOpProgram } from '../unaryop_gpu' ; import * as unary_packed_op from '../unaryop_packed_gpu' ; import { UnaryOpPackedProgram } from '../unaryop_packed_gpu' ; import { SimpleBinaryKernelImplCPU , SimpleUnaryKernelImplCPU } from './shared' ; export const CHECK_NAN_SNIPPET_UNARY = `if (isnan(x)) return x;` ; export const CHECK_NAN_SNIPPET_BINARY = `
  if (isnan(a)) return a;
  if (isnan(b)) return b;
` ; export const CHECK_NAN_SNIPPET_BINARY_PACKED = `
  result.r = isNaN.r > 0. ? NAN : result.r;
  result.g = isNaN.g > 0. ? NAN : result.g;
  result.b = isNaN.b > 0. ? NAN : result.b;
  result.a = isNaN.a > 0. ? NAN : result.a;
` ; type UnaryKernelFuncConfig = { opSnippet : string ; packedOpSnippet ? : string ; cpuKernelImpl ? : SimpleUnaryKernelImplCPU ; dtype ? : DataType ; } ; param opSnippet param packedOpSnippet param dtype export function unaryKernelFunc ( { opSnippet , packedOpSnippet , cpuKernelImpl , dtype } ) { return ( { inputs , backend } ) => { const { x } = inputs as UnaryInputs ; const webglBackend = backend as MathBackendWebGL ; const $dtype = dtype || x . dtype ; if ( webglBackend . shouldExecuteOnCPU ( [ x ] ) && cpuKernelImpl != null ) { const xData = webglBackend . texData . get ( x . dataId ) ; const outValues = cpuKernelImpl ( xData . values as TypedArray , $dtype ) ; return webglBackend . makeTensorInfo ( x . shape , $dtype , outValues ) ; } const shouldUsePackedProgram = env ( ) . getBool ( 'WEBGL_PACK_UNARY_OPERATIONS' ) && packedOpSnippet != null ; let program ; if ( shouldUsePackedProgram ) { program = new UnaryOpPackedProgram ( x . shape , packedOpSnippet ) ; } else { program = new UnaryOpProgram ( x . shape , opSnippet ) ; } return webglBackend . runWebGLProgram ( program , [ x ] , $dtype ) ; } ; } type BinaryKernelFuncConfig = { opSnippet : string ; packedOpSnippet ? : string ; checkOutOfBounds ? : boolean ; supportsComplex ? : boolean ; cpuKernelImpl ? : SimpleBinaryKernelImplCPU ; dtype ? : DataType ; } ; param opSnippet param packedOpSnippet param checkOutOfBoundsForPackedProgram param dtype export function binaryKernelFunc ( { opSnippet , packedOpSnippet , checkOutOfBounds = false , supportsComplex = false , cpuKernelImpl , dtype } ) { return ( { inputs , backend } ) => { const { a , b } = inputs as BinaryInputs ; const webglBackend = backend as MathBackendWebGL ; if ( supportsComplex && a . dtype === 'complex64' ) { const aData = webglBackend . texData . get ( a . dataId ) ; const bData = webglBackend . texData . get ( b . dataId ) ; const [ real , imag ] = [ [ aData . complexTensorInfos . real , bData . complexTensorInfos . real ] , [ aData . complexTensorInfos . imag , bData . complexTensorInfos . imag ] ] . map ( complexParts => { const [ aPart , bPart ] = complexParts ; const aHandle = { dataId : aPart . dataId , dtype : aPart . dtype , shape : a . shape } ; const bHandle = { dataId : bPart . dataId , dtype : bPart . dtype , shape : b . shape } ; const program = new BinaryOpProgram ( opSnippet , a . shape , b . shape ) ; return webglBackend . runWebGLProgram ( program , [ aHandle , bHandle ] , upcastType ( aPart . dtype , bPart . dtype ) ) ; } ) ; const complexOutput = complex ( { inputs : { real , imag } , backend : webglBackend } ) ; webglBackend . disposeIntermediateTensorInfo ( real ) ; webglBackend . disposeIntermediateTensorInfo ( imag ) ; return complexOutput ; } const $dtype = dtype || upcastType ( a . dtype , b . dtype ) ; if ( ( a . dtype === 'string' || b . dtype === 'string' || webglBackend . shouldExecuteOnCPU ( [ a , b ] ) ) && cpuKernelImpl != null ) { const aVals = webglBackend . texData . get ( a . dataId ) . values as TypedArray ; const bVals = webglBackend . texData . get ( b . dataId ) . values as TypedArray ; const decodedAVals = a . dtype === 'string' ? backend_util . fromUint8ToStringArray ( aVals as any as Uint8Array [ ] ) : aVals ; const decodedBVals = a . dtype === 'string' ? backend_util . fromUint8ToStringArray ( bVals as any as Uint8Array [ ] ) : bVals ; const [ outValues , outShape ] = cpuKernelImpl ( a . shape , b . shape , decodedAVals , decodedBVals , $dtype ) ; const out = webglBackend . makeTensorInfo ( outShape , $dtype ) ; const outData = webglBackend . texData . get ( out . dataId ) ; outData . values = outValues ; return out ; } const shouldUsePackedProgram = env ( ) . getBool ( 'WEBGL_PACK_BINARY_OPERATIONS' ) && packedOpSnippet != null ; let program ; if ( shouldUsePackedProgram ) { program = new BinaryOpPackedProgram ( packedOpSnippet , a . shape , b . shape , checkOutOfBounds ) ; } else { program = new BinaryOpProgram ( opSnippet , a . shape , b . shape ) ; } return webglBackend . runWebGLProgram ( program , [ a , b ] , $dtype ) ; } ; } export function mapActivationToShaderProgram ( activation , packed = false ) { if ( activation === 'linear' ) { if ( packed ) { return unary_packed_op . LINEAR ; } return unary_op . LINEAR ; } else if ( activation === 'relu' ) { if ( packed ) { return unary_packed_op . RELU ; } return unary_op . RELU ; } else if ( activation === 'elu' ) { if ( packed ) { return unary_packed_op . ELU ; } return unary_op . ELU ; } else if ( activation === 'relu6' ) { if ( packed ) { return unary_packed_op . RELU6 ; } return unary_op . RELU6 ; } else if ( activation === 'prelu' ) { if ( packed ) { return PRELU_PACKED ; } return PRELU ; } else if ( activation === 'leakyrelu' ) { if ( packed ) { return LEAKYRELU_PACKED ; } return LEAKYRELU ; } else if ( activation === 'sigmoid' ) { if ( packed ) { return unary_packed_op . SIGMOID ; } return unary_op . SIGMOID ; } throw new Error ( `Activation ${ activation } has not been implemented for the WebGL backend.` ) ; }