1  #extension GL_EXT_shader_16bit_storage : require
 2  
 3  layout (push_constant) uniform parameter
 4  {
 5      uint ne;
 6      uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03;
 7      uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13;
 8      uint d_offset;
 9      float param1; float param2;
10  } p;
11  
12  layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
13  
14  layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
15  layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
16  
17  uint src0_idx(uint idx) {
18      const uint i03 = idx / (p.ne02*p.ne01*p.ne00);
19      const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00;
20      const uint i02 = (idx - i03_offset) / (p.ne01*p.ne00);
21      const uint i02_offset = i02*p.ne01*p.ne00;
22      const uint i01 = (idx - i03_offset - i02_offset) / p.ne00;
23      const uint i00 = idx - i03_offset - i02_offset - i01*p.ne00;
24      return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00;
25  }
26  
27  uint dst_idx(uint idx) {
28      const uint i13 = idx / (p.ne12*p.ne11*p.ne10);
29      const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10;
30      const uint i12 = (idx - i13_offset) / (p.ne11*p.ne10);
31      const uint i12_offset = i12*p.ne11*p.ne10;
32      const uint i11 = (idx - i13_offset - i12_offset) / p.ne10;
33      const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10;
34      return i13*p.nb13 + i12*p.nb12 + i11*p.nb11 + i10*p.nb10;
35  }