Unpack input in TopK WebGL kernel #5286
Conversation
There was a problem hiding this comment.
I think the current mechanism can handle this case:
The GPGPUProgram defaults packedInputs and packedOutput to false: https://github.com/tensorflow/tfjs/blob/master/tfjs-backend-webgl/src/gpgpu_math.ts#L31-L33
When an input is packed, and the program expects an unpacked input, it will do the unpacking here: https://github.com/tensorflow/tfjs/blob/master/tfjs-backend-webgl/src/backend_webgl.ts#L823
Maybe I missed anything?
Reviewable status: 0 of 1 approvals obtained (waiting on @pyu10055)
|
The issue isn't that its unpacking, it's that it unpacks for every shader so it's repeating computation |
There was a problem hiding this comment.
Do you mean, because TopK call the shaders repeatedly in a loop, and so it is doing unpacking and packing in each iteration?
Reviewable status: 0 of 1 approvals obtained (waiting on @pyu10055)
|
Yes, i think in other kernels shaders aren't usually repeated and if they are it will be on new input not the original input, so they don't have that issue |
There was a problem hiding this comment.
Make sense. In that case, should we do the reshape after the unpack, because reshape of unpacked tensor is not expensive. Otherwise, LGTM. Also curious how much perf improvement with this change?
Reviewable status:
complete! 1 of 1 approvals obtained (waiting on @pyu10055)
There was a problem hiding this comment.
Reviewed 3 of 3 files at r1.
Reviewable status:
Great catch, will swap the reshape. The change was about a ~20% improvement on packed inputs (0 for unpacked since there's no issues with packing there) |
This prevents each shader from having to unpack the input again due to lazy unpacking
To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.
This change is