webgpu: add a naive implementation of conv3d #7016
Conversation
xhcao
force-pushed
the
conv3d
branch
2 times, most recently
from
3f7772d to
d69431f
Compare
gyagp
left a comment
gyagp
left a comment
There was a problem hiding this comment.
LGTM with some nits. Overall, we should use more WGSL syntax sugar, like ++ and +=.
| let inputDepthVec4Remainder = uniforms.xShape.u % 4; | ||
| var dotProd = 0.0; | ||
| for (var wF = 0; wF < uniforms.filterDims[0]; wF = wF + 1) { |
There was a problem hiding this comment.
| for (var wF = 0; wF < uniforms.filterDims[0]; wF = wF + 1) { | |
| for (var wF = 0; wF < uniforms.filterDims[0]; wF++) { |
| continue; | ||
| } | ||
| for (var wR = 0; wR < uniforms.filterDims[1]; wR = wR + 1) { |
| continue; | ||
| } | ||
| for (var wC = 0; wC < uniforms.filterDims[2]; wC = wC + 1) { |
| continue; | ||
| } | ||
| for (var d1 = 0; d1 < inputDepthNearestVec4; d1 = d1 + 4) { |
| getW(wF, wR, wC, d1 + 3, d2) | ||
| ); | ||
| dotProd = dotProd + dot(xValues, wValues); |
| getW(wF, wR, wC, inputDepthNearestVec4, d2), | ||
| getW(wF, wR, wC, inputDepthNearestVec4 + 1, d2) | ||
| ); | ||
| dotProd = dotProd + dot(xValues, wValues); |
| getW(wF, wR, wC, inputDepthNearestVec4 + 1, d2), | ||
| getW(wF, wR, wC, inputDepthNearestVec4 + 2, d2) | ||
| ); | ||
| dotProd = dotProd + dot(xValues, wValues); |
|
|
||
| const numCoords = indicesArr.length; | ||
| const shape = indicesArr.map(d => `${variableName}[${d}]`); | ||
| const indicesStr = ['.x', '.y', '.z', '.w', '.u', '.v']; |
There was a problem hiding this comment.
Why is the sequence wuv instead of uvw?
You may just define indicesStr as 'xyzuvw', then use it as ${variableName}.${indicesStr[d]}.
There was a problem hiding this comment.
Hi, @gyagp, xyzw is used for single component selection by WGSL https://gpuweb.github.io/gpuweb/wgsl/#:~:text=7.7.1.1.%20Vector%20Single%20Component%20Selection
TFJS webgpu backend extends the vec4 to vec5 and vec6
uv to select the fifth and sixth components.
All other comments are addressed, thank you.
| @@ -0,0 +1,129 @@ | |||
| /** | |||
| * @license | |||
| * Copyright 2022 Google LLC. | |||
There was a problem hiding this comment.
nit: 2022 -> 2023 and similar for other places
| variableNames = ['x', 'W']; | ||
| uniforms = | ||
| 'filterDims: vec3<i32>, pad: vec3<i32>, strides: vec3<i32>, dilations: vec3<i32>,'; | ||
| workgroupSize: [number, number, number] = [4, 4, 8]; |
There was a problem hiding this comment.
Can we use the flat dispatch layout here? Previous experience shows that flat dispatch layout always has good performance if the algorithm is irrelevant with dispatch layout. You can have a micro-benchmark to verify it.
3 participants
To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.
This change is