webgpu: Add shaderKey as the program member #2451
Conversation
This change is to reduce the overhead of shaderKey comparison. Previously, we used a very long string 'userCode' as part of the shader key which is time consuming when search the shader key.
qjia7
changed the title
|
@annxingyuan Please have a preview if it's in the right direction. It seems that latest webgpu backend has some regressions. Lots of tests will fail because of SPIRV validation error. I will look at that first. |
annxingyuan
left a comment
annxingyuan
left a comment
There was a problem hiding this comment.
@qjia7 This is great! However I would like this to be opt-in for each webgpu program. Can you make sure shaderKey is an optional property and that we default to the user code when shaderKey is undefined?
Also please keep me posted re the SPIRV validation errors - thanks so much for looking into those!
Please see if the latest change is what you want. I am wondering maybe in future we can change userCode to a function getUserCode(). In this case, we can move a lot logic to getUserCode and don't need to do them every time when we call the program constructor.
Filed a bug to SPIRV-Tools. See KhronosGroup/SPIRV-Tools#3085 |
annxingyuan
left a comment
annxingyuan
left a comment
There was a problem hiding this comment.
@qjia7 Thank you this is great!
FIX hat happens: Run case one by one (yarn test --grep=), pass. But when run several together. Some case may fail. Example case: Case may be failed due to cache reasons(div/Arithmetic_test.ts): 'gradient: Tensor1D with int32', // Actual != expected. 'gradient: Tensor2D', // Actual != expected. 'gradient: Tensor2D / Tensor2D w/ broadcast', // Actual != expected. Root cause: This problem is introduced by two reasons: 1), makeShaderKey called before the real shader generated(compileProgram and makeShader). Which means, two different shaders may possibly share the same shader key. 2), The shader key optimization(tensorflow#2451), which use several key properties of the shader instead of the full shader source as shader key. For most case, especially all properties are self-contained in the shader`s WebGPUProgram file, this works. But when some properties are originated from the shader_preprocessor.ts, two different shader may be collided with the same shader key. Currently the fix is to disable the shader key and shader cache first. The follow up todo is: tensorflow#2669. Also add square and neg operators.
3 participants
This change is to reduce the overhead of shaderKey comparison.
Previously, we used a very long string 'userCode' as part
of the shader key which is time consuming when search the shader key.
To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.
This change is