backend tfjs-backend-webgpu is missing several common kernel ops

after converting my existing app to run with tfjs-backend-webgpu built from main branch
(latest released version tfjs-backend-webgpu 0.0.1-alpha.7 is just too old),
i’ve run into issues running some of the models (all of them work perfectly using webgl backend):

Kernel ‘TopK’ not registered for backend ‘webgpu’ Kernel ‘SplitV’ not registered for backend ‘webgpu’ Kernel ‘FlipLeftRight’ not registered for backend ‘webgpu’ Kernel ‘RotateWithOffset’ not registered for backend ‘webgpu’

note: with remaining running models, with webgpu inference is 300-400ms while with webgl it’s 400-550ms
which makes webgpu ~30% faster than webgl!

btw, i’ll try using webgpu in web workers next as single thread does not come even close of saturating gpu
(average utilization is ~25%)

environment: chrome/94 canary with tfjs 3.8.0 and tjfs-backend-webgpu built from main

[xhcao]

Thanks yunfei for implementing the split kernel. Now, only topK kernel is unimplemented, but it is the most difficult one, the algorithm of webgl backend is based on the paper https://anilshanbhag.in/static/papers/gputopk_sigmod18.pdf, I will investigate and implement this operator with @carrycooldude, and make some optimizations in future according to webgpu features.

[qjia7]

@carrycooldude You can refer to this one #5585, which implements a new kernel DepthToSpace for webgpu.