decode.glsl (5500B)
1 /* See LICENSE for license details. */ 2 3 #if CooperativeMatrix 4 #extension GL_KHR_cooperative_matrix : require 5 #extension GL_KHR_memory_scope_semantics : require 6 #endif 7 8 layout(std430, buffer_reference, buffer_reference_align = 64) restrict readonly buffer RF { 9 InputDataType x[]; 10 }; 11 12 layout(std430, buffer_reference, buffer_reference_align = 64) restrict writeonly buffer Output { 13 OutputDataType x[]; 14 }; 15 16 layout(std430, buffer_reference) buffer F16 { f16 x[]; }; 17 18 OutputDataType sample_rf_data(u32 index) 19 { 20 OutputDataType result = OutputDataType(RF(rf_buffer).x[index]); 21 return result; 22 } 23 24 shared InputDataType rf[gl_WorkGroupSize.y][TransmitCount]; 25 void run_decode_large(void) 26 { 27 u32 transmit = gl_GlobalInvocationID.x * ToProcess; 28 u32 channel = gl_GlobalInvocationID.y; 29 u32 time_sample = gl_GlobalInvocationID.z; 30 31 const u32 samples_per_thread = TransmitCount / gl_WorkGroupSize.x; 32 const u32 leftover_samples = TransmitCount % gl_WorkGroupSize.x; 33 34 u32 thread_index_x = gl_LocalInvocationID.x; 35 u32 samples_this_thread = samples_per_thread + u32(thread_index_x < leftover_samples); 36 37 u32 rf_offset = TransmitCount * ChunkChannelCount * gl_WorkGroupID.z + TransmitCount * channel; 38 39 for (u32 i = 0; i < samples_this_thread; i++) { 40 u32 index = i * gl_WorkGroupSize.x + thread_index_x; 41 rf[gl_LocalInvocationID.y][index] = RF(rf_buffer).x[rf_offset + index]; 42 } 43 44 barrier(); 45 46 OutputDataType result[ToProcess]; 47 if (time_sample < OutputTransmitStride) { 48 for (s32 i = 0; i < ToProcess; i++) 49 result[i] = OutputDataType(0); 50 51 F16 h = F16(Hadamard); 52 for (s32 j = 0; j < TransmitCount; j++) { 53 OutputDataType s = OutputDataType(rf[gl_LocalInvocationID.y][j]); 54 for (s32 i = 0; i < ToProcess; i++) 55 result[i] += s * h.x[TransmitCount * j + (i + transmit)]; 56 } 57 58 for (uint i = 0; i < ToProcess; i++) 59 result[i] /= float(TransmitCount); 60 } 61 62 /* NOTE(rnp): DO NOT combine with above; compiler shits the bed on TransmitCount == 80 63 * and it kills performance. reinvestigate when we further optimize */ 64 if (time_sample < OutputTransmitStride) { 65 uint out_off = OutputChannelStride * channel + 66 OutputTransmitStride * transmit + 67 OutputSampleStride * time_sample; 68 69 for (uint i = 0; i < ToProcess; i++, out_off += OutputTransmitStride) 70 if (TransmitCount % (gl_WorkGroupSize.x * ToProcess) == 0 || transmit + i < TransmitCount) 71 Output(output_buffer).x[out_off] = result[i]; 72 } 73 } 74 75 #if CooperativeMatrix 76 void run_decode_coop_shmem(void) 77 { 78 } 79 80 void run_decode_coop(void) 81 { 82 u32vec2 tile_index = gl_WorkGroupID.xy; 83 u32 time_sample = gl_WorkGroupID.z; 84 85 coopmat<f16, gl_ScopeSubgroup, CooperativeMatrixM, CooperativeMatrixK, gl_MatrixUseA> rf_matrix; 86 coopmat<f16, gl_ScopeSubgroup, CooperativeMatrixK, CooperativeMatrixN, gl_MatrixUseB> hadamard_matrix; 87 coopmat<f32, gl_ScopeSubgroup, CooperativeMatrixM, CooperativeMatrixN, gl_MatrixUseAccumulator> result; 88 result = coopmat<f32, gl_ScopeSubgroup, CooperativeMatrixM, CooperativeMatrixN, gl_MatrixUseAccumulator>(0.0f); 89 90 u32 result_row = CooperativeMatrixM * tile_index.y; 91 u32 result_col = CooperativeMatrixN * tile_index.x; 92 93 u32 offset = ChunkChannelCount * TransmitCount * time_sample; 94 95 F16 h = F16(Hadamard); 96 for (u32 k = 0; k < TransmitCount; k += CooperativeMatrixK) { 97 u32 rf_tile_row = CooperativeMatrixM * tile_index.y; 98 u32 rf_tile_col = k; 99 coopMatLoad(rf_matrix, RF(rf_buffer).x, offset + TransmitCount * rf_tile_row + rf_tile_col, 100 TransmitCount, gl_CooperativeMatrixLayoutRowMajor); 101 102 u32 hadamard_tile_row = k; 103 u32 hadamard_tile_col = CooperativeMatrixN * tile_index.x; 104 coopMatLoad(hadamard_matrix, h.x, TransmitCount * hadamard_tile_row + hadamard_tile_col, 105 TransmitCount, gl_CooperativeMatrixLayoutRowMajor); 106 107 result = coopMatMulAdd(rf_matrix, hadamard_matrix, result); 108 } 109 110 for (s32 i = 0; i < result.length(); i++) 111 result[i] = result[i] / f32(TransmitCount); 112 113 Output out_buffer = Output(output_buffer); 114 coopMatStore(result, out_buffer.x, offset + TransmitCount * result_row + result_col, 115 TransmitCount, gl_CooperativeMatrixLayoutRowMajor); 116 } 117 #endif 118 119 void run_decode_small(void) 120 { 121 u32 time_sample = gl_GlobalInvocationID.x; 122 u32 channel = gl_GlobalInvocationID.y; 123 u32 rf_offset = TransmitCount * ChunkChannelCount * time_sample + TransmitCount * channel; 124 125 if (time_sample < OutputTransmitStride) { 126 InputDataType rf[TransmitCount]; 127 for (s32 j = 0; j < TransmitCount; j++) 128 rf[j] = RF(rf_buffer).x[rf_offset + j]; 129 130 OutputDataType result[TransmitCount]; 131 for (s32 j = 0; j < TransmitCount; j++) 132 result[j] = OutputDataType(0); 133 134 F16 h = F16(Hadamard); 135 for (s32 i = 0; i < TransmitCount; i++) { 136 OutputDataType s = OutputDataType(rf[i]); 137 for (s32 j = 0; j < TransmitCount; j++) { 138 result[j] += s * h.x[TransmitCount * i + j]; 139 } 140 } 141 142 for (int i = 0; i < TransmitCount; i++) 143 result[i] /= float(TransmitCount); 144 145 uint out_off = OutputChannelStride * channel + 146 OutputSampleStride * time_sample; 147 for (int i = 0; i < TransmitCount; i++, out_off += OutputTransmitStride) 148 Output(output_buffer).x[out_off] = result[i]; 149 } 150 } 151 152 void main() 153 { 154 switch (DecodeMode) { 155 case DecodeMode_Hadamard:{ 156 #if CooperativeMatrix 157 if (UseSharedMemory) run_decode_coop_shmem(); 158 else run_decode_coop(); 159 #else 160 if (UseSharedMemory) run_decode_large(); 161 else run_decode_small(); 162 #endif 163 }break; 164 } 165 }