|
#version 450 |
|
|
|
#extension GL_EXT_shader_16bit_storage : require |
|
|
|
layout (push_constant) uniform parameter |
|
{ |
|
uint batch_offset; uint offset_delta; |
|
uint IC; |
|
uint IW; uint IH; |
|
uint OW; uint OH; |
|
uint KW; uint KH; |
|
uint pelements; |
|
uint CHW; |
|
int s0; int s1; |
|
int p0; int p1; |
|
int d0; int d1; |
|
} p; |
|
|
|
#include "types.comp" |
|
|
|
#define BLOCK_SIZE 256 |
|
|
|
layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; |
|
|
|
layout (binding = 0) readonly buffer X {A_TYPE data_a[];}; |
|
layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; |
|
|
|
void main() { |
|
const uint i = gl_GlobalInvocationID.x; |
|
if (i >= p.pelements) { |
|
return; |
|
} |
|
|
|
const uint ksize = p.OW * (p.KH > 1 ? p.KW : 1); |
|
const uint kx = i / ksize; |
|
const uint kd = kx * ksize; |
|
const uint ky = (i - kd) / p.OW; |
|
const uint ix = i % p.OW; |
|
|
|
const uint oh = gl_GlobalInvocationID.y; |
|
const uint batch = gl_GlobalInvocationID.z / p.IC; |
|
const uint ic = gl_GlobalInvocationID.z % p.IC; |
|
|
|
const uint iiw = ix * p.s0 + kx * p.d0 - p.p0; |
|
const uint iih = oh * p.s1 + ky * p.d1 - p.p1; |
|
|
|
const uint offset_dst = |
|
((batch * p.OH + oh) * p.OW + ix) * p.CHW + |
|
(ic * (p.KW * p.KH) + ky * p.KW + kx); |
|
|
|
if (iih < 0 || iih >= p.IH || iiw < 0 || iiw >= p.IW) { |
|
data_d[offset_dst] = D_TYPE(0.0f); |
|
} else { |
|
const uint offset_src = ic * p.offset_delta + batch * p.batch_offset; |
|
data_d[offset_dst] = D_TYPE(data_a[offset_src + iih * p.IW + iiw]); |
|
} |
|
} |
|
|
