vulkan : kernels for depthwise 2D convolution (CONV_2D_DW) #1204
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,104 @@ | ||
| #version 450 | ||
|
|
||
| #include "types.comp" | ||
|
|
||
| layout (push_constant) uniform parameter | ||
| { | ||
| uint ne; | ||
| uint batches; | ||
| uint channels; | ||
| uint dst_w; | ||
| uint dst_h; | ||
| uint src_w; | ||
| uint src_h; | ||
| uint knl_w; | ||
| uint knl_h; | ||
| int stride_x; | ||
| int stride_y; | ||
| int pad_x; | ||
| int pad_y; | ||
| int dilation_x; | ||
| int dilation_y; | ||
| } p; | ||
|
|
||
| layout (binding = 0) readonly buffer A {A_TYPE knl_data[];}; | ||
| layout (binding = 1) readonly buffer B {B_TYPE src_data[];}; | ||
| layout (binding = 2) writeonly buffer D {D_TYPE dst_data[];}; | ||
|
|
||
| layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; | ||
|
|
||
| FLOAT_TYPE conv_2d_dw_whcn(uint idx) { | ||
| uint i0 = idx / p.dst_w; | ||
| uint dst_x = idx - i0 * p.dst_w; | ||
| uint i1 = i0 / p.dst_h; | ||
| uint dst_y = i0 - i1 * p.dst_h; | ||
| uint n = i1 / p.channels; | ||
| uint c = i1 - n * p.channels; | ||
|
|
||
| uint src_i = n * p.channels * p.src_h * p.src_w + c * p.src_h * p.src_w; | ||
| uint knl_i = c * p.knl_h * p.knl_w; | ||
|
|
||
| FLOAT_TYPE sum = 0.0; | ||
| for (uint knl_y = 0; knl_y < p.knl_h; ++knl_y) { | ||
| uint src_y = dst_y * p.stride_y - p.pad_y + knl_y * p.dilation_y; | ||
| if (src_y < 0 || src_y >= p.src_h) { | ||
| continue; | ||
| } | ||
| for (uint knl_x = 0; knl_x < p.knl_w; ++knl_x) { | ||
| uint src_x = dst_x * p.stride_x - p.pad_x + knl_x * p.dilation_x; | ||
| if (src_x < 0 || src_x >= p.src_w) { | ||
| continue; | ||
| } | ||
| FLOAT_TYPE v = FLOAT_TYPE(src_data[src_i + src_y * p.src_w + src_x]); | ||
| FLOAT_TYPE k = FLOAT_TYPE(knl_data[knl_i + knl_y * p.knl_w + knl_x]); | ||
| sum = fma(v, k, sum); | ||
| } | ||
| } | ||
| return sum; | ||
| } | ||
|
|
||
| FLOAT_TYPE conv_2d_dw_cwhn(uint idx) { | ||
| uint i0 = idx / p.channels; | ||
| uint c = idx - i0 * p.channels; | ||
| uint i1 = i0 / p.dst_w; | ||
| uint dst_x = i0 - i1 * p.dst_w; | ||
| uint n = i1 / p.dst_h; | ||
| uint dst_y = i1 - n * p.dst_h; | ||
|
|
||
| uint src_i = n * p.channels * p.src_h * p.src_w; | ||
| uint src_row = p.src_w * p.channels; | ||
| uint knl_row = p.knl_w * p.channels; | ||
|
|
||
| FLOAT_TYPE sum = 0.0; | ||
| for (uint knl_y = 0; knl_y < p.knl_h; ++knl_y) { | ||
| uint src_y = dst_y * p.stride_y + knl_y * p.dilation_y - p.pad_y; | ||
| if (src_y < 0 || src_y >= p.src_h) { | ||
| continue; | ||
| } | ||
| for (uint knl_x = 0; knl_x < p.knl_w; ++knl_x) { | ||
| uint src_x = dst_x * p.stride_x + knl_x * p.dilation_x - p.pad_x; | ||
| if (src_x < 0 || src_x >= p.src_w) { | ||
| continue; | ||
|
There was a problem hiding this comment. I guess padding is always considered to be with a value of zero, never replicating the border? There was a problem hiding this comment. Yes, there is no way to specify padding modes other than zero for convolutions so far. |
||
| } | ||
| FLOAT_TYPE v = FLOAT_TYPE(src_data[src_i + src_y * src_row + src_x * p.channels + c]); | ||
| FLOAT_TYPE k = FLOAT_TYPE(knl_data[ knl_y * knl_row + knl_x * p.channels + c]); | ||
| sum = fma(v, k, sum); | ||
| } | ||
| } | ||
| return sum; | ||
| } | ||
|
|
||
| void main() { | ||
| uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; | ||
| if (idx >= p.ne) { | ||
| return; | ||
| } | ||
|
|
||
| FLOAT_TYPE result = | ||
| #ifdef WHCN | ||
| conv_2d_dw_whcn(idx); | ||
| #else | ||
| conv_2d_dw_cwhn(idx); | ||
| #endif | ||
| dst_data[idx] = D_TYPE(result); | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Since src_x and src_y are unsigned, the
< 0conditions can be removed.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is it possible
src_yunderflows if pad_y is large enough?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, if there is padding the expression can become negative and wraps to a very large unsigned int, which will then be caught by the
>=check (for typical values). So in the end it does what's intended, and thesrc_y < 0check can be omitted.The alternative is to use signed int and keep the check (bit cleaner, more instructions). Using unsigned and the check as it is makes no sense, I'll fix that.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Removed the check and added a comment to indicate wrapping is intentional.