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)
Thanks yunfei for implementing the
split
kernel. Now, onlytopK
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.
@carrycooldude Below links implement these kernels on webgl backend, you could refer them. https://github.com/tensorflow/tfjs/blob/master/tfjs-backend-webgl/src/kernels/TopK.ts https://github.com/tensorflow/tfjs/blob/master/tfjs-backend-webgl/src/kernels/SplitV.ts https://github.com/tensorflow/tfjs/blob/master/tfjs-backend-webgl/src/kernels/FlipLeftRight.ts https://github.com/tensorflow/tfjs/blob/master/tfjs-backend-webgl/src/kernels/RotateWithOffset.ts
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.