license import { env , KernelConfig , KernelFunc , NumericDataType , TensorInfo , TopK , TopKAttrs , TopKInputs , TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { topKImplCPU } from '../kernel_utils/shared' ; import { MergeProgram , SwapProgram } from '../top_k_gpu' ; import { fill } from './Fill' ; import { gatherV2 } from './GatherV2' ; import { reshape } from './Reshape' ; import { slice } from './Slice' ; function disposeIntermediateTensorInfoOrNull ( backend , tensorInfo ) { if ( tensorInfo !== null ) { backend . disposeIntermediateTensorInfo ( tensorInfo ) ; } } function roundUpToPow2 ( num ) { let pow2 = 1 ; while ( pow2 < num ) { pow2 *= 2 ; } return pow2 ; } export function topK ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { k , sorted } = attrs ; const TOPK_LAST_DIM_CPU_HANDOFF_SIZE_THRESHOLD = env ( ) . getNumber ( 'TOPK_LAST_DIM_CPU_HANDOFF_SIZE_THRESHOLD' ) ; const TOPK_K_CPU_HANDOFF_THRESHOLD = env ( ) . getNumber ( 'TOPK_K_CPU_HANDOFF_THRESHOLD' ) ; const xShape = x . shape ; const lastDim = xShape [ xShape . length - 1 ] ; if ( backend . shouldExecuteOnCPU ( [ x ] ) || lastDim < TOPK_LAST_DIM_CPU_HANDOFF_SIZE_THRESHOLD || k > TOPK_K_CPU_HANDOFF_THRESHOLD ) { const xVals = backend . readSync ( x . dataId ) as TypedArray ; const [ allTopKVals , allTopKIndices ] = topKImplCPU ( xVals , xShape , x . dtype as NumericDataType , k , sorted ) ; return [ backend . makeTensorInfo ( allTopKVals . shape , allTopKVals . dtype , allTopKVals . values ) , backend . makeTensorInfo ( allTopKIndices . shape , allTopKIndices . dtype , allTopKIndices . values ) ] ; } if ( k === 0 ) { xShape [ xShape . length - 1 ] = 0 ; return [ backend . makeTensorInfo ( xShape , x . dtype , [ ] ) , backend . makeTensorInfo ( xShape , 'int32' , [ ] ) ] ; } if ( lastDim === 1 ) { return [ x , fill ( { attrs : { shape : xShape , dtype : 'int32' , value : 0 } , backend } ) ] ; } const xtexData = backend . texData . get ( x . dataId ) ; const xIsPacked = xtexData !== null && xtexData . isPacked ; const xUnPacked = xIsPacked ? backend . unpackTensor ( x ) : x ; const xSize = util . sizeFromShape ( xShape ) ; const batch = xSize / lastDim ; const x2D = reshape ( { inputs : { x : xUnPacked } , attrs : { shape : [ batch , lastDim ] } , backend } ) ; if ( xIsPacked ) { disposeIntermediateTensorInfoOrNull ( backend , xUnPacked ) ; } const kPow2 = roundUpToPow2 ( k ) ; const lastDimPow2 = roundUpToPow2 ( lastDim ) ; let indices = null ; const getInputs = ( ) => indices === null ? [ x2D , x2D ] : [ x2D , indices ] ; const runSwap = ( dir , inc , shape ) => { const inputs = getInputs ( ) ; const program = new SwapProgram ( shape ) ; const fistPass = indices === null ? 1 : 0 ; const customValues = [ [ lastDim ] , [ fistPass ] , [ Number . NEGATIVE_INFINITY ] , [ dir ] , [ inc ] ] ; const prevIndices = indices ; indices = backend . runWebGLProgram ( program , inputs , 'int32' , customValues ) ; disposeIntermediateTensorInfoOrNull ( backend , prevIndices ) ; } ; for ( let len = 1 ; len < kPow2 ; len *= 2 ) { const dir = len * 2 ; for ( let inc = len ; inc >= 1 ; inc /= 2 ) { runSwap ( dir , inc , [ batch , lastDimPow2 ] ) ; } } for ( let indicesSize = lastDimPow2 ; indicesSize > kPow2 ; indicesSize /= 2 ) { const inputs = getInputs ( ) ; const mergeProgram = new MergeProgram ( [ batch , indicesSize / 2 ] ) ; const firstPass = indices === null ? 1 : 0 ; const customValues = [ [ lastDim ] , [ firstPass ] , [ kPow2 ] ] ; const prevIndices = indices ; indices = backend . runWebGLProgram ( mergeProgram , inputs , 'int32' , customValues ) ; disposeIntermediateTensorInfoOrNull ( backend , prevIndices ) ; const len = kPow2 / 2 ; const dir = len * 2 ; for ( let inc = len ; inc >= 1 ; inc /= 2 ) { runSwap ( dir , inc , indices . shape ) ; } } let prevIndices = indices ; indices = slice ( { inputs : { x : indices } , backend , attrs : { begin : 0 , size : [ batch , k ] } } ) ; disposeIntermediateTensorInfoOrNull ( backend , prevIndices ) ; let values = gatherV2 ( { inputs : { x : x2D , indices } , backend , attrs : { axis : 1 , batchDims : 1 } } ) ; disposeIntermediateTensorInfoOrNull ( backend , x2D ) ; const newShape = xShape . slice ( 0 , - 1 ) ; newShape . push ( k ) ; prevIndices = indices ; indices = reshape ( { inputs : { x : indices } , attrs : { shape : newShape } , backend } ) ; disposeIntermediateTensorInfoOrNull ( backend , prevIndices ) ; const prevValues = values ; values = reshape ( { inputs : { x : values } , attrs : { shape : newShape } , backend } ) ; disposeIntermediateTensorInfoOrNull ( backend , prevValues ) ; return [ values , indices ] ; } export const topKConfig = { kernelName : TopK , backendName : 'webgl' , kernelFunc : topK as { } as KernelFunc } ;