tfjs: 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

About this issue

  • Original URL
  • State: closed
  • Created 3 years ago
  • Comments: 21 (9 by maintainers)

Most upvoted comments

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.

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

Could I take over these two kernels (‘FlipLeftRight’ and ‘RotateWithOffset’)? We could implement them synchronously.

Hi, @carrycooldude, how about the process of these kernels? Currently, I have spare time, and could help you to implement the kernels, which have been not started to implement.

Hi, @carrycooldude Are you implementing these kernel ops? You could submit several patches, and each patch only implements one kernel op. I know It is difficult to implement these ops, if you have some problems, we could discuss the solutions.