license import { BackendTimer , BackendTimingInfo } from './backends/backend' ; import { env } from './environment' ; import { Tensor } from './tensor' ; import { NamedTensorMap } from './tensor_types' ; import { DataType , DataTypeMap , TypedArray } from './types' ; import * as util from './util' ; export type KernelProfile = { kernelName : string ; outputs : Tensor [ ] ; inputs : NamedTensorMap ; timeMs : Promise < number | { error : string ; } > ; extraInfo : Promise < string > ; } ; export class Profiler { constructor ( private backendTimer , private logger ? ) { if ( logger == null ) { this . logger = new Logger ( ) ; } } profileKernel ( kernelName , inputs , f ) { let outputs ; const holdResultWrapperFn = ( ) => { outputs = f ( ) ; } ; let timer ; const start = util . now ( ) ; if ( this . backendTimer . timerAvailable ( ) ) { timer = this . backendTimer . time ( holdResultWrapperFn ) ; } else { holdResultWrapperFn ( ) ; for ( const output of outputs ) { output . dataSync ( ) ; } timer = Promise . resolve ( { kernelMs : util . now ( ) - start } ) ; } if ( env ( ) . getBool ( 'CHECK_COMPUTATION_FOR_ERRORS' ) ) { for ( let i = 0 ; i < outputs . length ; i ++ ) { const output = outputs [ i ] ; output . data ( ) . then ( tensorVals => { checkComputationForErrors ( tensorVals , output . dtype , kernelName ) ; } ) ; } } const kernelProfile = { kernelName , outputs , inputs , timeMs : timer . then ( timing => timing . kernelMs ) , extraInfo : timer . then ( timing => timing . getExtraProfileInfo != null ? timing . getExtraProfileInfo ( ) : '' ) } ; return kernelProfile ; } logKernelProfile ( kernelProfile ) { const { kernelName , outputs , timeMs , inputs , extraInfo } = kernelProfile ; outputs . forEach ( result => { Promise . all ( [ result . data ( ) , timeMs , extraInfo ] ) . then ( valueContainer => { this . logger . logKernelProfile ( kernelName , result , valueContainer [ 0 ] , valueContainer [ 1 ] , inputs , valueContainer [ 2 ] ) ; } ) ; } ) ; } } export function checkComputationForErrors < D extends DataType > ( vals , dtype , kernelName ) { if ( dtype !== 'float32' ) { return false ; } for ( let i = 0 ; i < vals . length ; i ++ ) { const num = vals [ i ] as number ; if ( isNaN ( num ) || ! isFinite ( num ) ) { console . warn ( `Found ${ num } in the result of '${ kernelName }'` ) ; return true ; } } return false ; } export class Logger { logKernelProfile ( name , result , vals , timeMs , inputs , extraInfo ? ) { const time = typeof timeMs === 'number' ? util . rightPad ( `${ timeMs }ms` , 9 ) : timeMs [ 'error' ] ; const paddedName = util . rightPad ( name , 25 ) ; const rank = result . rank ; const size = result . size ; const shape = util . rightPad ( result . shape . toString ( ) , 14 ) ; let inputShapesDescription = '' ; for ( const name in inputs ) { const input = inputs [ name ] ; if ( input != null ) { const inputShape = input . shape || result . shape ; const inputRank = inputShape . length ; inputShapesDescription += `${ name }: ${ inputRank }D ${ inputRank > 0 ? inputShape : '' } ` ; } } console . log ( `%c${ paddedName }\t%c${ time }\t%c${ rank }D ${ shape }\t%c${ size }\t%c${ inputShapesDescription }\t%c${ extraInfo }` , 'font-weight:bold' , 'color:red' , 'color:blue' , 'color: orange' , 'color: green' , 'color: steelblue' ) ; } }