ogl_beamforming

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

ogl_beamformer_lib.c (27184B)


      1 /* See LICENSE for license details. */
      2 #define BEAMFORMER_IMPORT static
      3 
      4 #ifndef BASE_PLATFORM_H
      5 #define BASE_PLATFORM_NO_MAIN 1
      6 #define BASE_EXPORT static
      7 #endif
      8 
      9 #include "../util.h"
     10 
     11 #include "../generated/beamformer.c"
     12 #include "ogl_beamformer_lib_base.h"
     13 
     14 #if OS_LINUX
     15 #include "../base_linux.c"
     16 #elif OS_WINDOWS
     17 #include "../base_win32.c"
     18 
     19 W32(iptr) OpenFileMappingA(u32, b32, c8 *);
     20 
     21 #else
     22 #error Unsupported Platform
     23 #endif
     24 
     25 #include "../util_os.c"
     26 #include "../beamformer_compute_stats.c"
     27 #include "../beamformer_shared_memory.c"
     28 
     29 global struct {
     30 	BeamformerSharedMemory *bp;
     31 	i32                     timeout_ms;
     32 	BeamformerLibErrorKind  last_error;
     33 	i64                     shared_memory_size;
     34 } g_beamformer_library_context;
     35 
     36 #if OS_LINUX
     37 
     38 function str8
     39 os_open_shared_memory_area(char *name)
     40 {
     41 	str8 result = {0};
     42 	i32 fd = shm_open(name, O_RDWR, S_IRUSR|S_IWUSR);
     43 	if (fd > 0) {
     44 		struct stat sb;
     45 		if (fstat(fd, &sb) != -1) {
     46 			void *new = mmap(0, sb.st_size, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0);
     47 			if (new != MAP_FAILED) {
     48 				result.data   = new;
     49 				result.length = sb.st_size;
     50 			}
     51 		}
     52 		close(fd);
     53 	}
     54 	return result;
     55 }
     56 
     57 function void
     58 os_close_shared_memory_area(void *memory, i64 size)
     59 {
     60 	munmap(memory, size);
     61 }
     62 
     63 #elif OS_WINDOWS
     64 
     65 W32(u64) VirtualQuery(void *base_address, void *memory_basic_info, u64 memory_basic_info_size);
     66 W32(b32) UnmapViewOfFile(void *);
     67 
     68 function b32
     69 os_reserve_region_locks(void)
     70 {
     71 	u8 buffer[1024];
     72 	Stream sb = {.data = buffer, .cap = countof(buffer)};
     73 	stream_append_str8(&sb, str8(OS_SHARED_MEMORY_NAME "_lock_"));
     74 
     75 	i32 start_index    = sb.widx;
     76 	u32 reserved_count = 0;
     77 	for EachElement(os_w32_shared_memory_semaphores, it) {
     78 		stream_reset(&sb, start_index);
     79 		stream_append_u64(&sb, it);
     80 		stream_append_byte(&sb, 0);
     81 		os_w32_shared_memory_semaphores[it] = os_w32_create_semaphore((c8 *)sb.data, 1, 1);
     82 		if InvalidHandle(os_w32_shared_memory_semaphores[it])
     83 			break;
     84 		reserved_count++;
     85 	}
     86 
     87 	b32 result = reserved_count == countof(os_w32_shared_memory_semaphores);
     88 	if (!result) {
     89 		for (u32 i = 0; i < reserved_count; i++)
     90 			CloseHandle(os_w32_shared_memory_semaphores[i].value[0]);
     91 	}
     92 
     93 	return result;
     94 }
     95 
     96 function str8
     97 os_open_shared_memory_area(char *name)
     98 {
     99 	struct alignas(16) {
    100 		void *BaseAddress;
    101 		void *AllocationBase;
    102 		u32   AllocationProtect;
    103 		u32   __alignment1;
    104 		u64   RegionSize;
    105 		u32   State;
    106 		u32   Protect;
    107 		u32   Type;
    108 		u32   __alignment2;
    109 	} memory_basic_info;
    110 
    111 	str8 result = {0};
    112 	iptr h = OpenFileMappingA(FILE_MAP_ALL_ACCESS, 0, name);
    113 	if (h != INVALID_FILE) {
    114 		// NOTE(rnp): a size of 0 maps the whole region, we can determine its size after
    115 		void *new = MapViewOfFile(h, FILE_MAP_ALL_ACCESS, 0, 0, 0);
    116 		if (new &&
    117 		    VirtualQuery(new, &memory_basic_info, sizeof(memory_basic_info)) == sizeof(memory_basic_info) &&
    118 		    os_reserve_region_locks())
    119 		{
    120 			result.data   = new;
    121 			result.length = (i64)memory_basic_info.RegionSize;
    122 		}
    123 
    124 		if (new && !result.data)
    125 			UnmapViewOfFile(new);
    126 
    127 		CloseHandle(h);
    128 	}
    129 	return result;
    130 }
    131 
    132 function void
    133 os_close_shared_memory_area(void *memory, i64 size)
    134 {
    135 	UnmapViewOfFile(memory);
    136 }
    137 
    138 #endif
    139 
    140 #define lib_error_check(c, e) lib_error_check_(c, BeamformerLibErrorKind_##e)
    141 function b32
    142 lib_error_check_(b32 condition, BeamformerLibErrorKind error_kind)
    143 {
    144 	b32 result = condition;
    145 	if (!result) g_beamformer_library_context.last_error = error_kind;
    146 	assert(result);
    147 	return result;
    148 }
    149 
    150 function b32
    151 check_shared_memory(void)
    152 {
    153 	b32 result = g_beamformer_library_context.bp != 0;
    154 	if unlikely(!g_beamformer_library_context.bp) {
    155 		str8 shared_memory = os_open_shared_memory_area(OS_SHARED_MEMORY_NAME);
    156 		if (lib_error_check(shared_memory.data != 0, SharedMemory)) {
    157 			BeamformerSharedMemory *bp = (BeamformerSharedMemory *)shared_memory.data;
    158 			result = lib_error_check(bp->version == BEAMFORMER_SHARED_MEMORY_VERSION, VersionMismatch);
    159 			if (result) {
    160 				g_beamformer_library_context.bp                 = bp;
    161 				g_beamformer_library_context.shared_memory_size = shared_memory.length;
    162 			} else {
    163 				os_close_shared_memory_area(shared_memory.data, shared_memory.length);
    164 			}
    165 		}
    166 	}
    167 
    168 	if likely(g_beamformer_library_context.bp)
    169 		result = lib_error_check(likely(!g_beamformer_library_context.bp->invalid), InvalidAccess);
    170 	return result;
    171 }
    172 
    173 function b32
    174 valid_parameter_block(u32 block)
    175 {
    176 	b32 result = check_shared_memory();
    177 	if (result) {
    178 		result = lib_error_check(block < g_beamformer_library_context.bp->reserved_parameter_blocks,
    179 		                         ParameterBlockUnallocated);
    180 	}
    181 	return result;
    182 }
    183 
    184 function BeamformWork *
    185 try_push_work_queue(void)
    186 {
    187 	BeamformWork *result = beamform_work_queue_push(&g_beamformer_library_context.bp->external_work_queue);
    188 	lib_error_check(result != 0, WorkQueueFull);
    189 	return result;
    190 }
    191 
    192 function b32
    193 lib_try_lock(i32 lock, i32 timeout_ms)
    194 {
    195 	b32 result = beamformer_shared_memory_take_lock(g_beamformer_library_context.bp, lock, (u32)timeout_ms);
    196 	lib_error_check(result, SyncVariable);
    197 	return result;
    198 }
    199 
    200 function void
    201 lib_release_lock(i32 lock)
    202 {
    203 	beamformer_shared_memory_release_lock(g_beamformer_library_context.bp, lock);
    204 }
    205 
    206 u32
    207 beamformer_get_api_version(void)
    208 {
    209 	return BEAMFORMER_SHARED_MEMORY_VERSION;
    210 }
    211 
    212 const char *
    213 beamformer_error_string(BeamformerLibErrorKind kind)
    214 {
    215 	#define X(type, num, string) string,
    216 	local_persist const char *error_string_table[] = {BEAMFORMER_LIB_ERRORS "invalid error kind"};
    217 	#undef X
    218 	return error_string_table[Min(kind, countof(error_string_table) - 1)];
    219 }
    220 
    221 BeamformerLibErrorKind
    222 beamformer_get_last_error(void)
    223 {
    224 	return g_beamformer_library_context.last_error;
    225 }
    226 
    227 const char *
    228 beamformer_get_last_error_string(void)
    229 {
    230 	return beamformer_error_string(beamformer_get_last_error());
    231 }
    232 
    233 void
    234 beamformer_set_global_timeout(u32 timeout_ms)
    235 {
    236 	g_beamformer_library_context.timeout_ms = timeout_ms;
    237 }
    238 
    239 b32
    240 beamformer_reserve_parameter_blocks(uint32_t count)
    241 {
    242 	b32 result = 0;
    243 	if (check_shared_memory() &&
    244 	    lib_error_check(count <= BeamformerMaxParameterBlocks, ParameterBlockOverflow))
    245 	{
    246 		g_beamformer_library_context.bp->reserved_parameter_blocks = count;
    247 		result = 1;
    248 	}
    249 	return result;
    250 }
    251 
    252 function b32
    253 validate_parameters(BeamformerParameters *bp)
    254 {
    255 	if (!lib_error_check(Between(bp->contrast_mode, 0, BeamformerContrastMode_Count - 1), InvalidContrastMode))
    256 		return 0;
    257 
    258 	u32 contrast_raw_sample_count = bp->acquisition_count * bp->sample_count * beamformer_contrast_mode_samples[bp->contrast_mode];
    259 	if (!lib_error_check(contrast_raw_sample_count <= bp->raw_data_dimensions.x, DataSizeMismatch))
    260 		return 0;
    261 
    262 	// NOTE(rnp): frame size checks
    263 	{
    264 		// TODO(rnp): this check is overly conservative, what if we are exporting something smaller than Float32Complex
    265 		u64 buffer_size     = g_beamformer_library_context.bp->beamformed_frame_buffer_size;
    266 		u64 frame_size      = Max(1, bp->output_points.x) * Max(1, bp->output_points.y) * Max(1, bp->output_points.z)
    267 		                      * beamformer_data_kind_byte_size[BeamformerDataKind_Float32Complex];
    268 		u64 incoherent_size = frame_size / 2;
    269 		if (bp->coherency_weighting)
    270 			buffer_size -= incoherent_size;
    271 
    272 		if (!lib_error_check(frame_size <= buffer_size, FrameSizeOverflow))
    273 			return 0;
    274 	}
    275 
    276 	return 1;
    277 }
    278 
    279 function b32
    280 validate_pipeline(i32 *shaders, u32 shader_count, BeamformerDataKind data_kind)
    281 {
    282 	b32 data_kind_test = Between(data_kind, 0, BeamformerDataKind_Count - 1);
    283 	if (!lib_error_check(data_kind_test, InvalidDataKind))
    284 		return 0;
    285 
    286 	if (!lib_error_check(shader_count <= BeamformerMaxComputeShaderStages, ComputeStageOverflow))
    287 		return 0;
    288 
    289 	for (u32 i = 0; i < shader_count; i++) {
    290 		b32 stage_test = Between(shaders[i], BeamformerShaderKind_ComputeFirst, BeamformerShaderKind_ComputeLast);
    291 		if (!lib_error_check(stage_test, InvalidComputeStage))
    292 			return 0;
    293 
    294 		if (shaders[i] == BeamformerShaderKind_Hilbert &&
    295 		    !lib_error_check(g_beamformer_library_context.bp->capabilities.hilbert != 0, InvalidComputeStage))
    296 			return 0;
    297 
    298 		if (shaders[i] == BeamformerShaderKind_Demodulate &&
    299 		    !lib_error_check(!beamformer_data_kind_complex[data_kind], InvalidDemodulationDataKind))
    300 		{
    301 			return 0;
    302 		}
    303 	}
    304 
    305 	b32 start_stage_test = shaders[0] == BeamformerShaderKind_Demodulate ||
    306 	                       shaders[0] == BeamformerShaderKind_Decode;
    307 	if (!lib_error_check(start_stage_test, InvalidStartShader))
    308 		return 0;
    309 
    310 	return 1;
    311 }
    312 
    313 u64
    314 beamformer_maximum_rf_data_size(void)
    315 {
    316 	u64 result = U64_MAX;
    317 	if (check_shared_memory()) {
    318 		Arena *sm = beamformer_shared_memory_scratch_arena(g_beamformer_library_context.bp,
    319 		                                                   g_beamformer_library_context.shared_memory_size);
    320 		result = Min(sm->reserved - sm->position, g_beamformer_library_context.bp->capabilities.max_rf_data_size);
    321 	}
    322 	return result;
    323 }
    324 
    325 u64
    326 beamformer_maximum_frames_for_parameters(BeamformerParameters *bp)
    327 {
    328 	u64 result = U64_MAX;
    329 	if (check_shared_memory() && validate_parameters(bp)) {
    330 		// TODO(rnp): overly conservative frame size check
    331 		u64 buffer_size     = g_beamformer_library_context.bp->beamformed_frame_buffer_size;
    332 		u64 frame_size      = Max(1, bp->output_points.x) * Max(1, bp->output_points.y) * Max(1, bp->output_points.z)
    333 		                      * beamformer_data_kind_byte_size[BeamformerDataKind_Float32Complex];
    334 		u64 incoherent_size = frame_size / 2;
    335 		if (bp->coherency_weighting)
    336 			buffer_size -= incoherent_size;
    337 		result = buffer_size / frame_size;
    338 	}
    339 	return result;
    340 }
    341 
    342 u64
    343 beamformer_maximum_frames_for_simple_parameters(BeamformerSimpleParameters *bp)
    344 {
    345 	u64 result = beamformer_maximum_frames_for_parameters((BeamformerParameters *)bp);
    346 	return result;
    347 }
    348 
    349 function b32
    350 parameter_block_region_upload(void *data, u32 size, u32 block, BeamformerParameterBlockRegions region_id,
    351                               u32 block_offset, i32 timeout_ms)
    352 {
    353 	i32 lock   = BeamformerSharedMemoryLockKind_Count + (i32)block;
    354 	b32 result = valid_parameter_block(block) && lib_try_lock(lock, timeout_ms);
    355 	if (result) {
    356 		memory_copy((u8 *)beamformer_parameter_block(g_beamformer_library_context.bp, block) + block_offset,
    357 		            data, size);
    358 		mark_parameter_block_region_dirty(g_beamformer_library_context.bp, block, region_id);
    359 		lib_release_lock(lock);
    360 	}
    361 	return result;
    362 }
    363 
    364 b32
    365 beamformer_set_pipeline_stage_parameters_at(u32 stage_index, i32 parameter, u32 block)
    366 {
    367 	u32 offset  = BeamformerParameterBlockRegionOffsets[BeamformerParameterBlockRegion_ComputePipeline];
    368 	offset     += offsetof(BeamformerComputePipeline, parameters);
    369 	offset     += (stage_index % BeamformerMaxComputeShaderStages) * sizeof(BeamformerShaderParameters);
    370 	b32 result  = parameter_block_region_upload(&parameter, sizeof(BeamformerShaderParameters), block,
    371 	                                            BeamformerParameterBlockRegion_ComputePipeline, offset,
    372 	                                            g_beamformer_library_context.timeout_ms);
    373 	return result;
    374 }
    375 
    376 b32
    377 beamformer_set_pipeline_stage_parameters(u32 stage_index, i32 parameter)
    378 {
    379 	b32 result = beamformer_set_pipeline_stage_parameters_at(stage_index, parameter, 0);
    380 	return result;
    381 }
    382 
    383 b32
    384 beamformer_push_pipeline_at(i32 *shaders, u32 shader_count, BeamformerDataKind data_kind, u32 block)
    385 {
    386 	b32 result = 0;
    387 	if (check_shared_memory() && validate_pipeline(shaders, shader_count, data_kind)) {
    388 		i32 lock = BeamformerSharedMemoryLockKind_Count + (i32)block;
    389 		if (valid_parameter_block(block) && lib_try_lock(lock, g_beamformer_library_context.timeout_ms)) {
    390 			BeamformerParameterBlock *b = beamformer_parameter_block(g_beamformer_library_context.bp, block);
    391 			memory_copy(&b->pipeline.shaders, shaders, shader_count * sizeof(*shaders));
    392 			mark_parameter_block_region_dirty(g_beamformer_library_context.bp, block,
    393 			                                  BeamformerParameterBlockRegion_ComputePipeline);
    394 			b->pipeline.shader_count = shader_count;
    395 			b->pipeline.data_kind    = data_kind;
    396 			lib_release_lock(lock);
    397 			result = 1;
    398 		}
    399 	}
    400 	return result;
    401 }
    402 
    403 b32
    404 beamformer_push_pipeline(i32 *shaders, u32 shader_count, BeamformerDataKind data_kind)
    405 {
    406 	b32 result = beamformer_push_pipeline_at(shaders, shader_count, data_kind, 0);
    407 	return result;
    408 }
    409 
    410 b32
    411 beamformer_create_filter(BeamformerFilterParameters *filter, u8 filter_slot, u8 parameter_block)
    412 {
    413 	b32 result = 0;
    414 	if (lib_error_check(filter->kind >= 0 && filter->kind < BeamformerFilterKind_Count, InvalidFilterKind)) {
    415 		if (check_shared_memory()) {
    416 			BeamformWork *work = try_push_work_queue();
    417 			if (work) {
    418 				BeamformerCreateFilterContext *ctx = &work->create_filter_context;
    419 				work->kind = BeamformerWorkKind_CreateFilter;
    420 				ctx->parameters      = *filter;
    421 				ctx->filter_slot     = filter_slot     % BeamformerFilterSlots;
    422 				ctx->parameter_block = parameter_block % BeamformerMaxParameterBlocks;
    423 				beamform_work_queue_push_commit(&g_beamformer_library_context.bp->external_work_queue);
    424 				result = 1;
    425 			}
    426 		}
    427 	}
    428 	return result;
    429 }
    430 
    431 function void
    432 beamformer_flush_commands(void)
    433 {
    434 	i32 lock = BeamformerSharedMemoryLockKind_DispatchCompute;
    435 	beamformer_shared_memory_take_lock(g_beamformer_library_context.bp, lock, 0);
    436 }
    437 
    438 #define BEAMFORMER_UPLOAD_FNS \
    439 	X(channel_mapping,               i16, 1, ChannelMapping) \
    440 	X(focal_vectors,                 f32, 2, FocalVectors)   \
    441 	X(sparse_elements,               i16, 1, SparseElements) \
    442 	X(transmit_receive_orientations, u8,  1, TransmitReceiveOrientations)
    443 
    444 #define X(name, dtype, elements, region_name) \
    445 b32 beamformer_push_##name ##_at(dtype *data, u32 count, u32 block) { \
    446 	b32 result = 0; \
    447 	if (lib_error_check(count <= countof(((BeamformerParameterBlock *)0)->name), BufferOverflow)) { \
    448 		result = parameter_block_region_upload(data, count * elements * sizeof(dtype), block, \
    449 		                                       BeamformerParameterBlockRegion_##region_name,  \
    450 		                                       offsetof(BeamformerParameterBlock, name),      \
    451 		                                       g_beamformer_library_context.timeout_ms);      \
    452 	} \
    453 	return result; \
    454 }
    455 BEAMFORMER_UPLOAD_FNS
    456 #undef X
    457 
    458 #define X(name, dtype, ...) \
    459 b32 beamformer_push_##name (dtype *data, u32 count) { \
    460 	b32 result = beamformer_push_##name ##_at(data, count, 0); \
    461 	return result; \
    462 }
    463 BEAMFORMER_UPLOAD_FNS
    464 #undef X
    465 
    466 #define BEAMFORMER_REDUCE_A1S2_CONTRAST_FN(name) void name(void *restrict output_v, \
    467                                                            void *restrict input_v, \
    468                                                            u32 sample_count)
    469 typedef BEAMFORMER_REDUCE_A1S2_CONTRAST_FN(beamformer_reduce_a1s2_contrast_fn);
    470 
    471 #define BEAMFORMER_REDUCE_A1S2_CONTRAST_LIST \
    472 	X(i16) \
    473 	X(f32) \
    474 	X(f16) \
    475 
    476 static_assert(BeamformerDataKind_Float16Complex == (BeamformerDataKind_Count - 1), "");
    477 
    478 #define X(type, ...) \
    479 function BEAMFORMER_REDUCE_A1S2_CONTRAST_FN(beamformer_reduce_a1s2_contrast_##type) \
    480 { \
    481 	type *input_a = (type *)input_v + 0 * sample_count; \
    482 	type *input_b = (type *)input_v + 1 * sample_count; \
    483 	type *input_c = (type *)input_v + 2 * sample_count; \
    484 	type *output  = (type *)output_v; \
    485 	for (u32 sample = 0; sample < sample_count; sample++) \
    486 		output[sample] = input_a[sample] - input_b[sample] - input_c[sample]; \
    487 }
    488 BEAMFORMER_REDUCE_A1S2_CONTRAST_LIST
    489 #undef X
    490 
    491 function b32
    492 beamformer_push_data_base(void *data, u32 data_size, i32 timeout_ms, u32 block)
    493 {
    494 	b32 result = 0;
    495 	Arena *scratch = beamformer_shared_memory_scratch_arena(g_beamformer_library_context.bp,
    496 	                                                        g_beamformer_library_context.shared_memory_size);
    497 	BeamformerParameterBlock *b  = beamformer_parameter_block(g_beamformer_library_context.bp, block);
    498 	BeamformerParameters     *bp = &b->parameters;
    499 	BeamformerDataKind     data_kind     = b->pipeline.data_kind;
    500 	BeamformerContrastMode contrast_mode = bp->contrast_mode;
    501 
    502 
    503 	u64 arena_size  = scratch->reserved - scratch->position;
    504 	u64 max_rf_size = g_beamformer_library_context.bp->capabilities.max_rf_data_size;
    505 	u32 rf_size     = bp->acquisition_count * bp->sample_count * bp->channel_count * beamformer_data_kind_byte_size[data_kind];
    506 	u32 raw_size    = bp->raw_data_dimensions.x * bp->raw_data_dimensions.y * beamformer_data_kind_byte_size[data_kind];
    507 
    508 	// TODO(rnp): support multi push upload so that max_rf_size is actual limit
    509 	if (lib_error_check(rf_size <= arena_size, BufferOverflow) &&
    510 	    lib_error_check(rf_size <= max_rf_size, RFDataSizeOverflow) &&
    511 	    lib_error_check(rf_size <= data_size && data_size == raw_size, DataSizeMismatch))
    512 	{
    513 		if (lib_try_lock(BeamformerSharedMemoryLockKind_UploadRF, timeout_ms)) {
    514 			if (lib_try_lock(BeamformerSharedMemoryLockKind_ScratchSpace, 0)) {
    515 				u32 channel_count      = bp->channel_count;
    516 				u32 out_channel_stride = beamformer_data_kind_byte_size[data_kind] * bp->sample_count * bp->acquisition_count;
    517 				u32 in_channel_stride  = beamformer_data_kind_byte_size[data_kind] * bp->raw_data_dimensions.x;
    518 
    519 				for (u32 channel = 0; channel < channel_count; channel++) {
    520 					u16 data_channel = (u16)b->channel_mapping[channel];
    521 					u32 out_off = out_channel_stride * channel;
    522 					u32 in_off  = in_channel_stride  * data_channel;
    523 					u8 *memory  = (u8 *)scratch + scratch->position + out_off;
    524 					switch (contrast_mode) {
    525 					default:{
    526 						/* NOTE(rnp): non temporal copy would be better, but we can't ensure
    527 						 * 64 byte boundaries. */
    528 						memory_copy(memory, (u8 *)data + in_off, out_channel_stride);
    529 					}break;
    530 
    531 					case BeamformerContrastMode_A1S2:{
    532 						read_only local_persist u8 reduce_a1s2_index_map[] = {
    533 							[BeamformerDataKind_Int16]          = 0,
    534 							[BeamformerDataKind_Int16Complex]   = 0,
    535 							[BeamformerDataKind_Float32]        = 1,
    536 							[BeamformerDataKind_Float32Complex] = 1,
    537 							[BeamformerDataKind_Float16]        = 2,
    538 							[BeamformerDataKind_Float16Complex] = 2,
    539 						};
    540 						static_assert(BeamformerDataKind_Float16Complex == (BeamformerDataKind_Count - 1), "");
    541 
    542 						read_only local_persist beamformer_reduce_a1s2_contrast_fn *reduce_a1s2_fn_table[] = {
    543 							#define X(type, ...) beamformer_reduce_a1s2_contrast_##type,
    544 							BEAMFORMER_REDUCE_A1S2_CONTRAST_LIST
    545 							#undef X
    546 						};
    547 
    548 						// TODO(rnp): HACK: for some unknown reason loading contrast data after loading
    549 						// non-contrast data causes the dataset to not be stored correctly (it looks
    550 						// like mix of the old and new dataset). Putting this here fixes the issue.
    551 						// Counter-intuitively this improves throughput on my zen4 test computer,
    552 						// however it obviously should not be needed.
    553 						memory_clear(memory, 0, out_channel_stride);
    554 
    555 						u32 sample_count = bp->sample_count * beamformer_data_kind_element_count[data_kind];
    556 						reduce_a1s2_fn_table[reduce_a1s2_index_map[data_kind]](memory, (u8 *)data + in_off, sample_count);
    557 					}break;
    558 					}
    559 				}
    560 
    561 				lib_release_lock(BeamformerSharedMemoryLockKind_ScratchSpace);
    562 				/* TODO(rnp): need a better way to communicate this */
    563 				u64 rf_block_rf_size = (u64)block << 32ULL | (u64)rf_size;
    564 				atomic_store_u64(&g_beamformer_library_context.bp->rf_block_rf_size, rf_block_rf_size);
    565 				result = 1;
    566 			}
    567 		}
    568 	}
    569 	return result;
    570 }
    571 
    572 b32
    573 beamformer_push_data_with_compute(void *data, u32 data_size, u32 image_plane_tag, u32 parameter_slot)
    574 {
    575 	b32 result = 0;
    576 	if (check_shared_memory()) {
    577 		u32 reserved_blocks = g_beamformer_library_context.bp->reserved_parameter_blocks;
    578 		if (lib_error_check(image_plane_tag < BeamformerViewPlaneTag_Count, InvalidImagePlane) &&
    579 		    lib_error_check(parameter_slot < reserved_blocks, ParameterBlockUnallocated) &&
    580 		    beamformer_push_data_base(data, data_size, g_beamformer_library_context.timeout_ms, parameter_slot))
    581 		{
    582 			BeamformWork *work = try_push_work_queue();
    583 			if (work) {
    584 				work->kind = BeamformerWorkKind_ComputeIndirect;
    585 				work->compute_context.view_plane      = image_plane_tag;
    586 				work->compute_context.parameter_block = parameter_slot;
    587 				beamform_work_queue_push_commit(&g_beamformer_library_context.bp->external_work_queue);
    588 				beamformer_flush_commands();
    589 				result = 1;
    590 			}
    591 		}
    592 	}
    593 	return result;
    594 }
    595 
    596 b32
    597 beamformer_push_parameters_at(BeamformerParameters *bp, u32 block)
    598 {
    599 	b32 result = check_shared_memory() && validate_parameters(bp);
    600 	if (result) {
    601 		result = parameter_block_region_upload(bp, sizeof(*bp), block,
    602 		                                       BeamformerParameterBlockRegion_Parameters,
    603 		                                       offsetof(BeamformerParameterBlock, parameters),
    604 		                                       g_beamformer_library_context.timeout_ms);
    605 		if (result) {
    606 			BeamformerParameterBlock *pb = beamformer_parameter_block(g_beamformer_library_context.bp, block);
    607 			atomic_or_u32(&pb->region_update_flags, 1u << BeamformerParameterRegionFlag_NotifyUI);
    608 		}
    609 	}
    610 	return result;
    611 }
    612 
    613 b32
    614 beamformer_push_parameters(BeamformerParameters *bp)
    615 {
    616 	b32 result = beamformer_push_parameters_at(bp, 0);
    617 	return result;
    618 }
    619 
    620 b32
    621 beamformer_push_simple_parameters_at(BeamformerSimpleParameters *bp, u32 block)
    622 {
    623 	b32 result = check_shared_memory();
    624 	if (result) {
    625 		alignas(64) v2 focal_vectors[countof(bp->steering_angles)];
    626 		for (u32 i = 0; i < countof(bp->steering_angles); i++)
    627 			focal_vectors[i] = (v2){{bp->steering_angles[i], bp->focal_depths[i]}};
    628 
    629 		result &= beamformer_push_parameters_at((BeamformerParameters *)bp, block);
    630 		result &= beamformer_push_pipeline_at(bp->compute_stages, bp->compute_stages_count, (BeamformerDataKind)bp->data_kind, block);
    631 		result &= beamformer_push_channel_mapping_at(bp->channel_mapping, bp->channel_count, block);
    632 		result &= beamformer_push_focal_vectors_at((f32 *)focal_vectors, countof(focal_vectors), block);
    633 		result &= beamformer_push_transmit_receive_orientations_at(bp->transmit_receive_orientations,
    634 		                                                           bp->acquisition_count, block);
    635 
    636 		if (bp->acquisition_kind == BeamformerAcquisitionKind_UFORCES ||
    637 		    bp->acquisition_kind == BeamformerAcquisitionKind_UHERCULES)
    638 		{
    639 			result &= beamformer_push_sparse_elements_at(bp->sparse_elements, bp->acquisition_count, block);
    640 		}
    641 
    642 		for (u32 stage = 0; stage < bp->compute_stages_count; stage++)
    643 			result &= beamformer_set_pipeline_stage_parameters_at(stage, bp->compute_stage_parameters[stage], block);
    644 	}
    645 	return result;
    646 }
    647 
    648 b32
    649 beamformer_push_simple_parameters(BeamformerSimpleParameters *bp)
    650 {
    651 	b32 result = beamformer_push_simple_parameters_at(bp, 0);
    652 	return result;
    653 }
    654 
    655 function b32
    656 beamformer_export_buffer(BeamformerExportContext export_context)
    657 {
    658 	BeamformWork *work = try_push_work_queue();
    659 	b32 result = work && lib_try_lock(BeamformerSharedMemoryLockKind_ExportSync, 0);
    660 	if (result) {
    661 		work->export_context = export_context;
    662 		work->kind = BeamformerWorkKind_ExportBuffer;
    663 		work->lock = BeamformerSharedMemoryLockKind_ScratchSpace;
    664 		beamform_work_queue_push_commit(&g_beamformer_library_context.bp->external_work_queue);
    665 	}
    666 	return result;
    667 }
    668 
    669 function b32
    670 beamformer_export(BeamformerExportContext export, void *out, i32 timeout_ms)
    671 {
    672 	b32 result = 0;
    673 	if (beamformer_export_buffer(export)) {
    674 		/* NOTE(rnp): if this fails it just means that the work from push_data hasn't
    675 		 * started yet. This is here to catch the other case where the work started
    676 		 * and finished before we finished queuing the export work item */
    677 		beamformer_flush_commands();
    678 
    679 		if (lib_try_lock(BeamformerSharedMemoryLockKind_ExportSync, timeout_ms)) {
    680 			if (lib_try_lock(BeamformerSharedMemoryLockKind_ScratchSpace, 0)) {
    681 				void *sm = beamformer_shared_memory_data_pointer(g_beamformer_library_context.bp,
    682 				                                                 g_beamformer_library_context.shared_memory_size);
    683 				memory_copy(out, sm, export.size);
    684 				lib_release_lock(BeamformerSharedMemoryLockKind_ScratchSpace);
    685 				result = 1;
    686 			}
    687 			lib_release_lock(BeamformerSharedMemoryLockKind_ExportSync);
    688 		}
    689 	}
    690 	return result;
    691 }
    692 
    693 BEAMFORMER_LIB_EXPORT b32
    694 beamformer_get_last_frames(void *out_data, u64 out_data_size, u32 count)
    695 {
    696 	BeamformerExportContext export = {0};
    697 	export.kind  = BeamformerExportKind_BeamformedData;
    698 	export.count = count;
    699 	export.size  = out_data_size;
    700 	b32 result = out_data && out_data_size && count && beamformer_export(export, out_data, g_beamformer_library_context.timeout_ms);
    701 	return result;
    702 }
    703 
    704 b32
    705 beamformer_beamform_data(BeamformerSimpleParameters *bp, void *data, uint32_t data_size,
    706                          void *out_data, int32_t timeout_ms)
    707 {
    708 	b32 result = beamformer_push_simple_parameters(bp);
    709 	if (result) {
    710 		iv3 output_points = bp->output_points.xyz;
    711 		output_points.E[0] = Max(1, output_points.E[0]);
    712 		output_points.E[1] = Max(1, output_points.E[1]);
    713 		output_points.E[2] = Max(1, output_points.E[2]);
    714 
    715 		b32 complex = 0;
    716 		for (u32 stage = 0; stage < bp->compute_stages_count; stage++) {
    717 			BeamformerShaderKind shader = (BeamformerShaderKind)bp->compute_stages[stage];
    718 			complex |= shader == BeamformerShaderKind_Demodulate || shader == BeamformerShaderKind_Hilbert;
    719 		}
    720 
    721 		u64 output_size = output_points.x * output_points.y * output_points.z * sizeof(f32);
    722 		if (complex) output_size *= 2;
    723 
    724 		Arena *scratch = beamformer_shared_memory_scratch_arena(g_beamformer_library_context.bp,
    725 		                                                        g_beamformer_library_context.shared_memory_size);
    726 		u64 scratch_size = scratch->reserved - scratch->position;
    727 		if (result && out_data) result &= lib_error_check(output_size <= scratch_size, ExportSpaceOverflow);
    728 
    729 		if (result) {
    730 			result = beamformer_push_data_with_compute(data, data_size, 0, 0);
    731 			if (result && out_data)
    732 				result = beamformer_get_last_frames(out_data, output_size, 1);
    733 		}
    734 	}
    735 	return result;
    736 }
    737 
    738 BEAMFORMER_LIB_EXPORT b32
    739 beamformer_compute_timings(BeamformerComputeStatsTable *output, i32 timeout_ms)
    740 {
    741 	b32 result = 0;
    742 	if (check_shared_memory()) {
    743 		Arena *scratch = beamformer_shared_memory_scratch_arena(g_beamformer_library_context.bp,
    744 		                                                        g_beamformer_library_context.shared_memory_size);
    745 		u64 scratch_size = scratch->reserved - scratch->position;
    746 		if (lib_error_check(sizeof(*output) <= scratch_size, ExportSpaceOverflow)) {
    747 			BeamformerExportContext export = {0};
    748 			export.kind = BeamformerExportKind_Stats;
    749 			export.size = sizeof(*output);
    750 			result = beamformer_export(export, output, timeout_ms);
    751 		}
    752 	}
    753 	return result;
    754 }
    755 
    756 i32
    757 beamformer_live_parameters_get_dirty_flag(void)
    758 {
    759 	i32 result = -1;
    760 	if (check_shared_memory()) {
    761 		u32 flag = ctz_u64(g_beamformer_library_context.bp->live_imaging_dirty_flags);
    762 		if (flag != 64) {
    763 			atomic_and_u32(&g_beamformer_library_context.bp->live_imaging_dirty_flags, ~(1u << flag));
    764 			result = (i32)flag;
    765 		}
    766 	}
    767 	return result;
    768 }
    769 
    770 BeamformerLiveImagingParameters *
    771 beamformer_get_live_parameters(void)
    772 {
    773 	BeamformerLiveImagingParameters *result = 0;
    774 	if (check_shared_memory()) result = &g_beamformer_library_context.bp->live_imaging_parameters;
    775 	return result;
    776 }
    777 
    778 b32
    779 beamformer_set_live_parameters(BeamformerLiveImagingParameters *new)
    780 {
    781 	b32 result = 0;
    782 	if (check_shared_memory()) {
    783 		memory_copy(&g_beamformer_library_context.bp->live_imaging_parameters, new, sizeof(*new));
    784 		store_fence();
    785 		result = 1;
    786 	}
    787 	return result;
    788 }