decode.glsl (5565B)
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_reference_align = 64) restrict readonly buffer Hadamard { 17 f16 x[]; 18 }; 19 20 OutputDataType sample_rf_data(u32 index) 21 { 22 OutputDataType result = OutputDataType(RF(rf_buffer).x[index]); 23 return result; 24 } 25 26 shared InputDataType rf[gl_WorkGroupSize.y][TransmitCount]; 27 void run_decode_large(void) 28 { 29 u32 transmit = gl_GlobalInvocationID.x * ToProcess; 30 u32 channel = gl_GlobalInvocationID.y; 31 u32 time_sample = gl_GlobalInvocationID.z; 32 33 const u32 samples_per_thread = TransmitCount / gl_WorkGroupSize.x; 34 const u32 leftover_samples = TransmitCount % gl_WorkGroupSize.x; 35 36 u32 thread_index_x = gl_LocalInvocationID.x; 37 u32 samples_this_thread = samples_per_thread + u32(thread_index_x < leftover_samples); 38 39 u32 rf_offset = TransmitCount * ChunkChannelCount * gl_WorkGroupID.z + TransmitCount * channel; 40 41 for (u32 i = 0; i < samples_this_thread; i++) { 42 u32 index = i * gl_WorkGroupSize.x + thread_index_x; 43 rf[gl_LocalInvocationID.y][index] = RF(rf_buffer).x[rf_offset + index]; 44 } 45 46 barrier(); 47 48 OutputDataType result[ToProcess]; 49 if (time_sample < OutputTransmitStride) { 50 for (s32 i = 0; i < ToProcess; i++) 51 result[i] = OutputDataType(0); 52 53 for (s32 j = 0; j < TransmitCount; j++) { 54 OutputDataType s = OutputDataType(rf[gl_LocalInvocationID.y][j]); 55 for (s32 i = 0; i < ToProcess; i++) 56 result[i] += s * Hadamard(hadamard_buffer).x[TransmitCount * j + (i + transmit)]; 57 } 58 59 for (uint i = 0; i < ToProcess; i++) 60 result[i] /= float(TransmitCount); 61 } 62 63 /* NOTE(rnp): DO NOT combine with above; compiler shits the bed on TransmitCount == 80 64 * and it kills performance. reinvestigate when we further optimize */ 65 if (time_sample < OutputTransmitStride) { 66 uint out_off = OutputChannelStride * channel + 67 OutputTransmitStride * transmit + 68 OutputSampleStride * time_sample; 69 70 for (uint i = 0; i < ToProcess; i++, out_off += OutputTransmitStride) 71 if (TransmitCount % (gl_WorkGroupSize.x * ToProcess) == 0 || transmit + i < TransmitCount) 72 Output(output_buffer).x[out_off] = result[i]; 73 } 74 } 75 76 #if CooperativeMatrix 77 void run_decode_coop_shmem(void) 78 { 79 } 80 81 void run_decode_coop(void) 82 { 83 u32vec2 tile_index = gl_WorkGroupID.xy; 84 u32 time_sample = gl_WorkGroupID.z; 85 86 coopmat<f16, gl_ScopeSubgroup, CooperativeMatrixM, CooperativeMatrixK, gl_MatrixUseA> rf_matrix; 87 coopmat<f16, gl_ScopeSubgroup, CooperativeMatrixK, CooperativeMatrixN, gl_MatrixUseB> hadamard_matrix; 88 coopmat<f32, gl_ScopeSubgroup, CooperativeMatrixM, CooperativeMatrixN, gl_MatrixUseAccumulator> result; 89 result = coopmat<f32, gl_ScopeSubgroup, CooperativeMatrixM, CooperativeMatrixN, gl_MatrixUseAccumulator>(0.0f); 90 91 u32 result_row = CooperativeMatrixM * tile_index.y; 92 u32 result_col = CooperativeMatrixN * tile_index.x; 93 94 u32 offset = ChunkChannelCount * TransmitCount * time_sample; 95 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, Hadamard(hadamard_buffer).x, 105 TransmitCount * hadamard_tile_row + hadamard_tile_col, TransmitCount, 106 gl_CooperativeMatrixLayoutRowMajor); 107 108 result = coopMatMulAdd(rf_matrix, hadamard_matrix, result); 109 } 110 111 for (s32 i = 0; i < result.length(); i++) 112 result[i] = result[i] / f32(TransmitCount); 113 114 Output out_buffer = Output(output_buffer); 115 coopMatStore(result, out_buffer.x, offset + TransmitCount * result_row + result_col, 116 TransmitCount, gl_CooperativeMatrixLayoutRowMajor); 117 } 118 #endif 119 120 void run_decode_small(void) 121 { 122 u32 time_sample = gl_GlobalInvocationID.x; 123 u32 channel = gl_GlobalInvocationID.y; 124 u32 rf_offset = TransmitCount * ChunkChannelCount * time_sample + TransmitCount * channel; 125 126 if (time_sample < OutputTransmitStride) { 127 InputDataType rf[TransmitCount]; 128 for (s32 j = 0; j < TransmitCount; j++) 129 rf[j] = RF(rf_buffer).x[rf_offset + j]; 130 131 OutputDataType result[TransmitCount]; 132 for (s32 j = 0; j < TransmitCount; j++) 133 result[j] = OutputDataType(0); 134 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 * Hadamard(hadamard_buffer).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 }