feat: cast-forward-f32-f32.

Author: b4rtaz

src/nn/nn-vulkan.cpp

@@ -531,7 +531,8 @@ static void resolveShaderGroups(const NnOpConfig *opConfig, const NnUint batchSi
         if (outputSize.floatType == F_Q80) {
             groupCount[2] = outputSize.x / Q80_BLOCK_SIZE;
         } else {
-            groupCount[2] = 32;
+            constexpr NnUint chunkSize = 4;
+            groupCount[2] = outputSize.x / chunkSize;
         }
     } else if (opConfig->code == OP_MERGE_ADD) {
         if (inputSize.floatType == F_Q80) {

src/nn/vulkan/cast-forward-f32-f32.comp

@@ -1,9 +1,11 @@
 #version 450
 
-#define N_THREADS 256
+#extension GL_EXT_control_flow_attributes : enable
+
+#define CHUNK_SIZE 4
 #define N_BATCHES 32
 
-layout(local_size_x = N_THREADS, local_size_y = 1, local_size_z = 1) in;
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
 
 struct BatchInfo {
     uint inputOffset;
@@ -17,20 +19,14 @@ layout(binding = 1) writeonly buffer outputBuffer { float y[]; };
 layout(binding = 2) readonly uniform batchInfosBuffer { BatchInfo infos[N_BATCHES]; };
 
 void main() {
-    const uint threadIndex = gl_LocalInvocationID.x;
-    const uint nWorkGroups = gl_NumWorkGroups.z;
     const uint batchIndex = gl_WorkGroupID.y;
-    const uint workGroupIndex = gl_WorkGroupID.z;
-
+    const uint chunkIndex = gl_WorkGroupID.z;
     const BatchInfo info = infos[batchIndex];
-    const uint slice = info.inputSizeX / nWorkGroups;
-    const uint rest = info.inputSizeX % nWorkGroups;
-    const uint offset = workGroupIndex * slice + min(rest, workGroupIndex);
-    const uint dim = slice + (workGroupIndex < rest ? 1 : 0);
+    const uint offset = chunkIndex * CHUNK_SIZE;
     const uint xOffset = info.inputOffset + offset;
     const uint yOffset = info.outputOffset + offset;
 
-    for (uint i = threadIndex; i < dim; i += N_THREADS) {
+    [[unroll]] for (uint i = 0; i < CHUNK_SIZE; i++) {
         y[yOffset + i] = x[xOffset + i];
     }
 }