decode.glsl (5280B)
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 = 16) readonly 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 u32 rf_offset = TransmitCount * ChunkChannelCount * gl_WorkGroupID.z + TransmitCount * channel; 32 for (u32 index = gl_LocalInvocationID.x; index < TransmitCount; index += gl_WorkGroupSize.x) 33 rf[gl_LocalInvocationID.y][index] = RF(rf_buffer).x[rf_offset + index]; 34 35 barrier(); 36 37 OutputDataType result[ToProcess]; 38 if (time_sample < OutputTransmitStride) { 39 for (s32 i = 0; i < ToProcess; i++) 40 result[i] = OutputDataType(0); 41 42 F16 h = F16(HeapBase + Hadamard); 43 for (s32 j = 0; j < TransmitCount; j++) { 44 OutputDataType s = OutputDataType(rf[gl_LocalInvocationID.y][j]); 45 for (s32 i = 0; i < ToProcess; i++) 46 result[i] += s * h.x[TransmitCount * j + (i + transmit)]; 47 } 48 49 for (uint i = 0; i < ToProcess; i++) 50 result[i] /= float(TransmitCount); 51 } 52 53 /* NOTE(rnp): DO NOT combine with above; compiler shits the bed on TransmitCount == 80 54 * and it kills performance. reinvestigate when we further optimize */ 55 if (time_sample < OutputTransmitStride) { 56 uint out_off = OutputChannelStride * channel + 57 OutputTransmitStride * transmit + 58 OutputSampleStride * time_sample; 59 60 for (uint i = 0; i < ToProcess; i++, out_off += OutputTransmitStride) 61 if (TransmitCount % (gl_WorkGroupSize.x * ToProcess) == 0 || transmit + i < TransmitCount) 62 Output(output_buffer).x[out_off] = result[i]; 63 } 64 } 65 66 #if CooperativeMatrix 67 void run_decode_coop_shmem(void) 68 { 69 } 70 71 void run_decode_coop(void) 72 { 73 u32vec2 tile_index = gl_WorkGroupID.xy; 74 u32 time_sample = gl_WorkGroupID.z; 75 76 coopmat<f16, gl_ScopeSubgroup, CooperativeMatrixM, CooperativeMatrixK, gl_MatrixUseA> rf_matrix; 77 coopmat<f16, gl_ScopeSubgroup, CooperativeMatrixK, CooperativeMatrixN, gl_MatrixUseB> hadamard_matrix; 78 coopmat<f32, gl_ScopeSubgroup, CooperativeMatrixM, CooperativeMatrixN, gl_MatrixUseAccumulator> result; 79 result = coopmat<f32, gl_ScopeSubgroup, CooperativeMatrixM, CooperativeMatrixN, gl_MatrixUseAccumulator>(0.0f); 80 81 u32 result_row = CooperativeMatrixM * tile_index.y; 82 u32 result_col = CooperativeMatrixN * tile_index.x; 83 84 u32 offset = ChunkChannelCount * TransmitCount * time_sample; 85 86 F16 h = F16(HeapBase + Hadamard); 87 for (u32 k = 0; k < TransmitCount; k += CooperativeMatrixK) { 88 u32 rf_tile_row = CooperativeMatrixM * tile_index.y; 89 u32 rf_tile_col = k; 90 coopMatLoad(rf_matrix, RF(rf_buffer).x, offset + TransmitCount * rf_tile_row + rf_tile_col, 91 TransmitCount, gl_CooperativeMatrixLayoutRowMajor); 92 93 u32 hadamard_tile_row = k; 94 u32 hadamard_tile_col = CooperativeMatrixN * tile_index.x; 95 coopMatLoad(hadamard_matrix, h.x, TransmitCount * hadamard_tile_row + hadamard_tile_col, 96 TransmitCount, gl_CooperativeMatrixLayoutRowMajor); 97 98 result = coopMatMulAdd(rf_matrix, hadamard_matrix, result); 99 } 100 101 for (s32 i = 0; i < result.length(); i++) 102 result[i] = result[i] / f32(TransmitCount); 103 104 Output out_buffer = Output(output_buffer); 105 coopMatStore(result, out_buffer.x, offset + TransmitCount * result_row + result_col, 106 TransmitCount, gl_CooperativeMatrixLayoutRowMajor); 107 } 108 #endif 109 110 void run_decode_small(void) 111 { 112 u32 time_sample = gl_GlobalInvocationID.x; 113 u32 channel = gl_GlobalInvocationID.y; 114 u32 rf_offset = TransmitCount * ChunkChannelCount * time_sample + TransmitCount * channel; 115 116 if (time_sample < OutputTransmitStride) { 117 InputDataType rf[TransmitCount]; 118 for (s32 j = 0; j < TransmitCount; j++) 119 rf[j] = RF(rf_buffer).x[rf_offset + j]; 120 121 OutputDataType result[TransmitCount]; 122 for (s32 j = 0; j < TransmitCount; j++) 123 result[j] = OutputDataType(0); 124 125 F16 h = F16(HeapBase + Hadamard); 126 for (s32 i = 0; i < TransmitCount; i++) { 127 OutputDataType s = OutputDataType(rf[i]); 128 for (s32 j = 0; j < TransmitCount; j++) { 129 result[j] += s * h.x[TransmitCount * i + j]; 130 } 131 } 132 133 for (int i = 0; i < TransmitCount; i++) 134 result[i] /= float(TransmitCount); 135 136 uint out_off = OutputChannelStride * channel + 137 OutputSampleStride * time_sample; 138 for (int i = 0; i < TransmitCount; i++, out_off += OutputTransmitStride) 139 Output(output_buffer).x[out_off] = result[i]; 140 } 141 } 142 143 void main() 144 { 145 switch (DecodeMode) { 146 case DecodeMode_Hadamard:{ 147 #if CooperativeMatrix 148 if (UseSharedMemory) run_decode_coop_shmem(); 149 else run_decode_coop(); 150 #else 151 if (UseSharedMemory) run_decode_large(); 152 else run_decode_small(); 153 #endif 154 }break; 155 } 156 }