| // Copyright 2021 The IREE Authors |
| // |
| // Licensed under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| |
| #include "experimental/rocm/native_executable.h" |
| |
| #include <stddef.h> |
| |
| #include "experimental/rocm/dynamic_symbols.h" |
| #include "experimental/rocm/pipeline_layout.h" |
| #include "experimental/rocm/status_util.h" |
| #include "iree/base/api.h" |
| |
| // flatcc schemas: |
| #include "iree/base/internal/flatcc/parsing.h" |
| #include "iree/schemas/rocm_executable_def_reader.h" |
| #include "iree/schemas/rocm_executable_def_verifier.h" |
| |
| typedef struct iree_hal_rocm_entry_point_t { |
| iree_hal_rocm_kernel_params_t kernel_params; |
| iree_string_view_t name; |
| IREE_TRACE(iree_hal_rocm_FileLineLocDef_table_t source_location;) |
| IREE_TRACE(iree_hal_rocm_StageLocationDef_vec_t stage_locations;) |
| } iree_hal_rocm_entry_point_t; |
| |
| typedef struct iree_hal_rocm_native_executable_t { |
| iree_hal_resource_t resource; |
| iree_hal_rocm_context_wrapper_t* context; |
| iree_hal_pipeline_layout_t** pipeline_layouts; |
| iree_host_size_t entry_count; |
| hipModule_t module; |
| iree_host_size_t entry_point_count; |
| iree_hal_rocm_entry_point_t entry_points[]; |
| } iree_hal_rocm_native_executable_t; |
| |
| static const iree_hal_executable_vtable_t |
| iree_hal_rocm_native_executable_vtable; |
| |
| static iree_hal_rocm_native_executable_t* iree_hal_rocm_native_executable_cast( |
| iree_hal_executable_t* base_value) { |
| IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_rocm_native_executable_vtable); |
| return (iree_hal_rocm_native_executable_t*)base_value; |
| } |
| |
| iree_status_t iree_hal_rocm_native_executable_create( |
| iree_hal_rocm_context_wrapper_t* context, |
| const iree_hal_executable_params_t* executable_params, |
| iree_hal_executable_t** out_executable) { |
| IREE_ASSERT_ARGUMENT(context); |
| IREE_ASSERT_ARGUMENT(executable_params); |
| IREE_ASSERT_ARGUMENT(out_executable); |
| *out_executable = NULL; |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| iree_hal_rocm_native_executable_t* executable = NULL; |
| |
| // TODO: Verify the flat buffer. |
| iree_hal_rocm_ExecutableDef_table_t executable_def = |
| iree_hal_rocm_ExecutableDef_as_root( |
| executable_params->executable_data.data); |
| |
| // Create the kernel module. |
| flatbuffers_string_t hsaco_image = |
| iree_hal_rocm_ExecutableDef_hsaco_image_get(executable_def); |
| flatbuffers_string_vec_t entry_points_vec = |
| iree_hal_rocm_ExecutableDef_entry_points_get(executable_def); |
| iree_hal_rocm_BlockSizeDef_vec_t block_sizes_vec = |
| iree_hal_rocm_ExecutableDef_block_sizes_get(executable_def); |
| flatbuffers_uint32_vec_t shared_memory_sizes = |
| iree_hal_rocm_ExecutableDef_shared_memory_sizes_get(executable_def); |
| iree_host_size_t entry_count = flatbuffers_string_vec_len(entry_points_vec); |
| |
| iree_host_size_t total_size = |
| sizeof(*executable) + entry_count * sizeof(executable->entry_points[0]); |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_allocator_malloc(context->host_allocator, total_size, |
| (void**)&executable)); |
| |
| iree_hal_resource_initialize(&iree_hal_rocm_native_executable_vtable, |
| &executable->resource); |
| |
| executable->context = context; |
| executable->entry_point_count = entry_count; |
| iree_status_t status = ROCM_RESULT_TO_STATUS( |
| context->syms, |
| hipModuleLoadDataEx(&executable->module, hsaco_image, 0, NULL, NULL), |
| "hipModuleLoadDataEx"); |
| if (!iree_status_is_ok(status)) { |
| status = iree_status_annotate( |
| status, |
| IREE_SV("mismatched target chip? missing/wrong bitcode directory?")); |
| } |
| |
| // Query allowed max shared memory. |
| int32_t max_shared_mem = 0; |
| if (iree_status_is_ok(status)) { |
| status = ROCM_RESULT_TO_STATUS( |
| context->syms, |
| hipDeviceGetAttribute(&max_shared_mem, |
| hipDeviceAttributeMaxSharedMemoryPerBlock, |
| context->rocm_device), |
| "hipDeviceGetAttribute"); |
| } |
| |
| if (iree_status_is_ok(status)) { |
| executable->entry_count = entry_count; |
| for (iree_host_size_t i = 0; i < entry_count; i++) { |
| if (iree_status_is_ok(status)) { |
| hipFunction_t function = NULL; |
| flatbuffers_string_t entry_name = |
| flatbuffers_string_vec_at(entry_points_vec, i); |
| status = ROCM_RESULT_TO_STATUS( |
| context->syms, |
| hipModuleGetFunction(&function, executable->module, entry_name), |
| "hipModuleGetFunction"); |
| if (!iree_status_is_ok(status)) break; |
| if (!function) { |
| status = iree_make_status(IREE_STATUS_NOT_FOUND, |
| "exported module function %s not found", |
| entry_name); |
| break; |
| } |
| if (shared_memory_sizes[i] > max_shared_mem) { |
| status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT, |
| "function '%s' requested shared memory " |
| "size of %d larger than allowed size of %d", |
| entry_name, shared_memory_sizes[i], |
| max_shared_mem); |
| } else if (shared_memory_sizes[i] != 0) { |
| status = ROCM_RESULT_TO_STATUS( |
| context->syms, |
| hipFuncSetAttribute( |
| function, |
| (hipFuncAttribute) |
| HIP_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, |
| shared_memory_sizes[i]), |
| "hipFuncSetAttribute"); |
| } |
| // Package required parameters for kernel launches for each entry point. |
| iree_hal_rocm_entry_point_t* entry_point_info = |
| &executable->entry_points[i]; |
| iree_hal_rocm_kernel_params_t* params = |
| &entry_point_info->kernel_params; |
| params->layout = executable_params->pipeline_layouts[i]; |
| iree_hal_pipeline_layout_retain(params->layout); |
| params->function = function; |
| params->block_size[0] = block_sizes_vec[i].x; |
| params->block_size[1] = block_sizes_vec[i].y; |
| params->block_size[2] = block_sizes_vec[i].z; |
| params->shared_memory_size = shared_memory_sizes[i]; |
| entry_point_info->name = iree_make_string_view( |
| entry_name, flatbuffers_string_len(entry_name)); |
| } |
| } |
| |
| #if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION |
| if (iree_status_is_ok(status)) { |
| if (iree_hal_rocm_ExecutableDef_source_locations_is_present( |
| executable_def)) { |
| iree_hal_rocm_FileLineLocDef_vec_t source_locations_vec = |
| iree_hal_rocm_ExecutableDef_source_locations_get(executable_def); |
| for (iree_host_size_t i = 0; i < entry_count; ++i) { |
| executable->entry_points[i].source_location = |
| iree_hal_rocm_FileLineLocDef_vec_at(source_locations_vec, i); |
| } |
| } |
| if (iree_hal_rocm_ExecutableDef_stage_locations_is_present( |
| executable_def)) { |
| iree_hal_rocm_StageLocationsDef_vec_t stage_locations_vec = |
| iree_hal_rocm_ExecutableDef_stage_locations_get(executable_def); |
| for (iree_host_size_t i = 0; i < entry_count; ++i) { |
| iree_hal_rocm_StageLocationsDef_table_t stage_locations = |
| iree_hal_rocm_StageLocationsDef_vec_at(stage_locations_vec, i); |
| executable->entry_points[i].stage_locations = |
| iree_hal_rocm_StageLocationsDef_locations_get(stage_locations); |
| } |
| } |
| |
| // Publish any embedded source files to the tracing infrastructure. |
| if (iree_hal_rocm_ExecutableDef_source_files_is_present(executable_def)) { |
| iree_hal_rocm_SourceFileDef_vec_t source_files_vec = |
| iree_hal_rocm_ExecutableDef_source_files_get(executable_def); |
| for (iree_host_size_t i = 0; |
| i < iree_hal_rocm_SourceFileDef_vec_len(source_files_vec); ++i) { |
| iree_hal_rocm_SourceFileDef_table_t source_file = |
| iree_hal_rocm_SourceFileDef_vec_at(source_files_vec, i); |
| flatbuffers_string_t path = |
| iree_hal_rocm_SourceFileDef_path_get(source_file); |
| flatbuffers_uint8_vec_t content = |
| iree_hal_rocm_SourceFileDef_content_get(source_file); |
| IREE_TRACE_PUBLISH_SOURCE_FILE(path, flatbuffers_string_len(path), |
| content, |
| flatbuffers_uint8_vec_len(content)); |
| } |
| } |
| } |
| #endif // IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION |
| } |
| |
| if (iree_status_is_ok(status)) { |
| *out_executable = (iree_hal_executable_t*)executable; |
| } else { |
| if (executable) { |
| iree_hal_executable_destroy((iree_hal_executable_t*)executable); |
| } |
| } |
| |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| hipFunction_t iree_hal_rocm_native_executable_for_entry_point( |
| iree_hal_executable_t* base_executable, int32_t entry_point) { |
| iree_hal_rocm_native_executable_t* executable = |
| iree_hal_rocm_native_executable_cast(base_executable); |
| return executable->entry_points[entry_point].kernel_params.function; |
| } |
| |
| static void iree_hal_rocm_native_executable_destroy( |
| iree_hal_executable_t* base_executable) { |
| iree_hal_rocm_native_executable_t* executable = |
| iree_hal_rocm_native_executable_cast(base_executable); |
| iree_allocator_t host_allocator = executable->context->host_allocator; |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| if (executable->module) { |
| iree_status_t status = ROCM_RESULT_TO_STATUS( |
| executable->context->syms, hipModuleUnload(executable->module), |
| "hipModuleUnload"); |
| if (!iree_status_is_ok(status)) { |
| fprintf(stderr, "Failed unloading ROCm module: "); |
| iree_status_fprint(stderr, status); |
| iree_status_free(status); |
| } |
| } |
| |
| if (executable->pipeline_layouts) { |
| for (iree_host_size_t i = 0; i < executable->entry_count; ++i) { |
| if (executable->pipeline_layouts[i]) { |
| iree_hal_pipeline_layout_release(executable->pipeline_layouts[i]); |
| } |
| } |
| } |
| |
| iree_allocator_free(host_allocator, executable); |
| |
| IREE_TRACE_ZONE_END(z0); |
| } |
| |
| iree_status_t iree_hal_rocm_native_executable_entry_point_kernel_params( |
| iree_hal_executable_t* base_executable, int32_t entry_point, |
| iree_hal_rocm_kernel_params_t* out_params) { |
| iree_hal_rocm_native_executable_t* executable = |
| iree_hal_rocm_native_executable_cast(base_executable); |
| if (entry_point >= executable->entry_count) { |
| return iree_make_status(IREE_STATUS_OUT_OF_RANGE, |
| "invalid entry point ordinal %d", entry_point); |
| } |
| memcpy(out_params, &executable->entry_points[entry_point], |
| sizeof(*out_params)); |
| return iree_ok_status(); |
| } |
| |
| void iree_hal_rocm_native_executable_entry_point_source_location( |
| iree_hal_executable_t* base_executable, iree_host_size_t entry_ordinal, |
| iree_hal_rocm_source_location_t* out_source_location) { |
| iree_hal_rocm_native_executable_t* executable = |
| iree_hal_rocm_native_executable_cast(base_executable); |
| memset(out_source_location, 0, sizeof(*out_source_location)); |
| if (entry_ordinal >= executable->entry_point_count) { |
| return; |
| } |
| const iree_hal_rocm_entry_point_t* entry_point = |
| &executable->entry_points[entry_ordinal]; |
| |
| out_source_location->func_name = entry_point->name; |
| |
| #if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION |
| iree_hal_rocm_FileLineLocDef_table_t source_location = |
| entry_point->source_location; |
| if (entry_point->stage_locations) { |
| for (size_t i = 0; i < iree_hal_rocm_StageLocationDef_vec_len( |
| entry_point->stage_locations); |
| ++i) { |
| iree_hal_rocm_StageLocationDef_table_t stage_location = |
| iree_hal_rocm_StageLocationDef_vec_at(entry_point->stage_locations, |
| i); |
| // TODO(benvanik): a way to select what location is chosen. For now we |
| // just pick the first one. |
| source_location = |
| iree_hal_rocm_StageLocationDef_location_get(stage_location); |
| break; |
| } |
| } |
| if (source_location) { |
| flatbuffers_string_t filename = |
| iree_hal_rocm_FileLineLocDef_filename_get(source_location); |
| out_source_location->file_name = |
| iree_make_string_view(filename, flatbuffers_string_len(filename)); |
| out_source_location->line = |
| iree_hal_rocm_FileLineLocDef_line_get(source_location); |
| } else { |
| out_source_location->file_name = out_source_location->func_name; |
| out_source_location->line = 0; |
| } |
| #else |
| out_source_location->file_name = out_source_location->func_name; |
| out_source_location->line = 0; |
| #endif // IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION |
| } |
| |
| static const iree_hal_executable_vtable_t |
| iree_hal_rocm_native_executable_vtable = { |
| .destroy = iree_hal_rocm_native_executable_destroy, |
| }; |
