ogl_beamforming

Ultrasound Beamforming Implemented with OpenGL
git clone anongit@rnpnr.xyz:ogl_beamforming.git
Log | Files | Refs | Feed | Submodules | README | LICENSE

filter.glsl (4463B)


      1 /* See LICENSE for license details. */
      2 #if  (InputDataKind == DataKind_Int16Complex         || \
      3      (InputDataKind == DataKind_Int16 && Demodulate) || \
      4      (InputDataKind == DataKind_Float16 && Demodulate))
      5   #define SAMPLE_TYPE f16vec2
      6 #elif InputDataKind == DataKind_Int16
      7   #define SAMPLE_TYPE f16
      8 #elif InputDataKind == DataKind_Float32 && Demodulate
      9   #define SAMPLE_TYPE f32vec2
     10 #endif
     11 
     12 #ifndef SAMPLE_TYPE
     13   #define SAMPLE_TYPE InputDataType
     14 #endif
     15 
     16 #define ComplexSampleType (InputDataKind == DataKind_Float32Complex || \
     17                            InputDataKind == DataKind_Float16Complex || \
     18                            InputDataKind == DataKind_Int16Complex   || \
     19                            Demodulate)
     20 #if ComplexSampleType
     21   #define RESULT_TYPE f32vec2
     22 #else
     23   #define RESULT_TYPE f32
     24 #endif
     25 
     26 #if ComplexFilter
     27   #define FILTER_TYPE f32vec2
     28 #else
     29   #define FILTER_TYPE f32
     30 #endif
     31 
     32 #if ComplexFilter && ComplexSampleType
     33   #define apply_filter(iq, h) complex_mul(f32vec2(iq), f32vec2(h))
     34 #else
     35   #define apply_filter(iq, h) ((iq) * (h))
     36 #endif
     37 
     38 layout(std430, buffer_reference) restrict readonly  buffer Input  { InputDataType x[];  };
     39 layout(std430, buffer_reference) restrict writeonly buffer Output { OutputDataType x[]; };
     40 
     41 layout(std430, buffer_reference, buffer_reference_align = 64) restrict readonly buffer Filter {
     42 	FILTER_TYPE values[FilterLength];
     43 };
     44 
     45 f32vec2 complex_mul(f32vec2 a, f32vec2 b)
     46 {
     47 	mat2 m = mat2(b.x, b.y, -b.y, b.x);
     48 	f32vec2 result = m * a;
     49 	return result;
     50 }
     51 
     52 #if Demodulate
     53 SAMPLE_TYPE rotate_iq(SAMPLE_TYPE iq, uint index)
     54 {
     55 	float arg          = radians(360) * DemodulationFrequency * index / SamplingFrequency;
     56 	SAMPLE_TYPE result = SAMPLE_TYPE(complex_mul(iq, f32vec2(cos(arg), -sin(arg))));
     57 	return result;
     58 }
     59 #endif
     60 
     61 shared SAMPLE_TYPE rf[DecimationRate * gl_WorkGroupSize.x + FilterLength - 1];
     62 
     63 void main()
     64 {
     65 	uint out_sample = gl_GlobalInvocationID.x;
     66 	uint channel    = gl_GlobalInvocationID.y;
     67 	uint transmit   = gl_GlobalInvocationID.z;
     68 
     69 	uint thread_index = gl_LocalInvocationIndex;
     70 	uint thread_count = gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroupSize.z;
     71 	/////////////////////////
     72 	// NOTE: sample caching
     73 	{
     74 		bool offset_wraps = (DecimationRate * gl_WorkGroupID.x * gl_WorkGroupSize.x) < (FilterLength - 1);
     75 
     76 		u32 in_offset = InputDataKindByteSize * (InputChannelStride * channel + InputTransmitStride * transmit);
     77 		// NOTE(rnp): when demodulating we want to load 2 elements at a time but the
     78 		// input strides were specified in terms of a single element. therefore we
     79 		// must divide this by two. by doing this here we can gracefully handle
     80 		// the case where there are an odd number of samples (this drops the last one).
     81 		if (Demodulate)
     82 			in_offset /= 2;
     83 
     84 		// NOTE(rnp): broken out to avoid overflow from the subtraction
     85 		u64 input_address = input_buffer + in_offset;
     86 		input_address += InputDataKindByteSize * (DecimationRate * gl_WorkGroupID.x * gl_WorkGroupSize.x);
     87 		input_address -= InputDataKindByteSize * (FilterLength - 1);
     88 
     89 		uint total_samples       = rf.length();
     90 		uint samples_per_thread  = total_samples / thread_count;
     91 		uint leftover_count      = total_samples % thread_count;
     92 		uint samples_this_thread = samples_per_thread + uint(thread_index < leftover_count);
     93 
     94 		const SAMPLE_TYPE scale = SAMPLE_TYPE(bool(ComplexFilter) ? 1 : sqrt(2.0f));
     95 		for (uint i = 0; i < samples_this_thread; i++) {
     96 			uint index = thread_count * i + thread_index;
     97 			SAMPLE_TYPE s = SAMPLE_TYPE(0);
     98 			if (!offset_wraps || index >= FilterLength - 1) {
     99 				s = SAMPLE_TYPE(Input(input_address).x[index]);
    100 				#if Demodulate
    101 				s = scale * rotate_iq(s * SAMPLE_TYPE(1, -1), index);
    102 				#endif
    103 			}
    104 			rf[index] = s;
    105 		}
    106 	}
    107 	barrier();
    108 
    109 	Filter f = Filter(FilterCoefficients);
    110 	if (out_sample < SampleCount / DecimationRate) {
    111 		RESULT_TYPE result = RESULT_TYPE(0);
    112 		u32 offset = DecimationRate * thread_index;
    113 		for (u32 j = 0; j < FilterLength; j++)
    114 			result += apply_filter(rf[offset + j], f.values[j]);
    115 
    116 		u32 out_offset = OutputChannelStride  * channel +
    117 		                 OutputTransmitStride * transmit +
    118 		                 OutputSampleStride   * out_sample;
    119 
    120 		if (BatchSampleCount != 0) {
    121 			// NOTE(rnp): deinterleave
    122 			Output(output_buffer).x[out_offset] = OutputDataType(result.x);
    123 			out_offset += BatchSampleCount;
    124 			Output(output_buffer).x[out_offset] = OutputDataType(result.y);
    125 		} else {
    126 			Output(output_buffer).x[out_offset] = OutputDataType(result);
    127 		}
    128 	}
    129 }