[rocm] Fix crash when executable source information is missing (#16805)
This commit fixes crashes when optional executable
source stage locations and/or files are not present:
```
runtime/src/iree/schemas/rocm_executable_def_reader.h:167:
iree_hal_rocm_StageLocationsDef_table_t iree_hal_rocm_StageLocationsDef_vec_at
(iree_hal_rocm_StageLocationsDef_vec_t, size_t):
Assertion `flatbuffers_vec_len(vec) > (i) && "index out of range"' failed.
```
Now we more closely mirror how Vulkan side is done for it.
diff --git a/experimental/rocm/direct_command_buffer.c b/experimental/rocm/direct_command_buffer.c
index 8986878..dff9cbd 100644
--- a/experimental/rocm/direct_command_buffer.c
+++ b/experimental/rocm/direct_command_buffer.c
@@ -419,14 +419,17 @@
iree_hal_rocm_native_executable_entry_point_kernel_params(
executable, entry_point, &kernel_params));
- IREE_ROCM_TRACE_ZONE_BEGIN_EXTERNAL(
- command_buffer->tracing_context, 0,
- kernel_params.source_location.file_name.data,
- kernel_params.source_location.file_name.size,
- kernel_params.source_location.line,
- kernel_params.source_location.func_name.data,
- kernel_params.source_location.func_name.size,
- /*name=*/NULL, 0);
+ IREE_TRACE({
+ iree_hal_rocm_source_location_t source_location;
+ iree_hal_rocm_native_executable_entry_point_source_location(
+ executable, entry_point, &source_location);
+ IREE_ROCM_TRACE_ZONE_BEGIN_EXTERNAL(
+ command_buffer->tracing_context, /*stream=*/0,
+ source_location.file_name.data, source_location.file_name.size,
+ source_location.line, source_location.func_name.data,
+ source_location.func_name.size,
+ /*name=*/NULL, 0);
+ });
// Patch the push constants in the kernel arguments.
iree_host_size_t num_constants =
diff --git a/experimental/rocm/native_executable.c b/experimental/rocm/native_executable.c
index 217e32e..35d33f5 100644
--- a/experimental/rocm/native_executable.c
+++ b/experimental/rocm/native_executable.c
@@ -18,6 +18,13 @@
#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;
@@ -25,7 +32,7 @@
iree_host_size_t entry_count;
hipModule_t module;
iree_host_size_t entry_point_count;
- iree_hal_rocm_kernel_params_t entry_points[];
+ iree_hal_rocm_entry_point_t entry_points[];
} iree_hal_rocm_native_executable_t;
static const iree_hal_executable_vtable_t
@@ -65,32 +72,17 @@
iree_hal_rocm_ExecutableDef_shared_memory_sizes_get(executable_def);
iree_host_size_t entry_count = flatbuffers_string_vec_len(entry_points_vec);
- // Calculate the total number of characters across all entry point names. This
- // is only required when tracing so that we can store copies of the names as
- // the flatbuffer storing the strings may be released while the executable is
- // still live.
- iree_host_size_t total_entry_point_name_chars = 0;
- IREE_TRACE({
- for (iree_host_size_t i = 0; i < entry_count; i++) {
- const char* entry_name = flatbuffers_string_vec_at(entry_points_vec, i);
- total_entry_point_name_chars += flatbuffers_string_len(entry_name);
- }
- });
-
iree_host_size_t total_size =
- sizeof(*executable) + entry_count * sizeof(executable->entry_points[0]) +
- total_entry_point_name_chars;
+ 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_TRACE(char* string_table_buffer =
- (char*)((char*)executable + sizeof(*executable) +
- entry_count * sizeof(executable->entry_points[0])));
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),
@@ -117,7 +109,8 @@
for (iree_host_size_t i = 0; i < entry_count; i++) {
if (iree_status_is_ok(status)) {
hipFunction_t function = NULL;
- const char* entry_name = flatbuffers_string_vec_at(entry_points_vec, i);
+ 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),
@@ -146,7 +139,10 @@
"hipFuncSetAttribute");
}
// Package required parameters for kernel launches for each entry point.
- iree_hal_rocm_kernel_params_t* params = &executable->entry_points[i];
+ 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;
@@ -154,66 +150,50 @@
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];
-
- // Stash the entry point name in the string table for use when tracing.
- IREE_TRACE({
- iree_host_size_t entry_name_length =
- flatbuffers_string_len(entry_name);
- memcpy(string_table_buffer, entry_name, entry_name_length);
- params->source_location.func_name =
- iree_make_string_view(string_table_buffer, entry_name_length);
- string_table_buffer += entry_name_length;
- });
+ entry_point_info->name = iree_make_string_view(
+ entry_name, flatbuffers_string_len(entry_name));
}
}
#if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
- iree_hal_rocm_FileLineLocDef_vec_t source_locations_vec =
- iree_hal_rocm_ExecutableDef_source_locations_get(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_StageLocationDef_vec_t stage_locations =
- iree_hal_rocm_StageLocationsDef_locations_get(
- iree_hal_rocm_StageLocationsDef_vec_at(stage_locations_vec, i));
- iree_hal_rocm_FileLineLocDef_table_t source_location =
- iree_hal_rocm_FileLineLocDef_vec_at(source_locations_vec, i);
- if (stage_locations) {
- for (size_t j = 0;
- j < iree_hal_rocm_StageLocationDef_vec_len(stage_locations); ++j) {
- iree_hal_rocm_StageLocationDef_table_t stage_location =
- iree_hal_rocm_StageLocationDef_vec_at(stage_locations, j);
- // 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 (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);
}
}
- iree_hal_rocm_kernel_params_t* params = &executable->entry_points[i];
- flatbuffers_string_t filename =
- iree_hal_rocm_FileLineLocDef_filename_get(source_location);
- params->source_location.file_name =
- iree_make_string_view(filename, flatbuffers_string_len(filename));
- params->source_location.line =
- iree_hal_rocm_FileLineLocDef_line_get(source_location);
- }
+ 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 (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));
+ // 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
@@ -235,7 +215,7 @@
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].function;
+ return executable->entry_points[entry_point].kernel_params.function;
}
static void iree_hal_rocm_native_executable_destroy(
@@ -283,6 +263,54 @@
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,
diff --git a/experimental/rocm/native_executable.h b/experimental/rocm/native_executable.h
index b23ab51..997d385 100644
--- a/experimental/rocm/native_executable.h
+++ b/experimental/rocm/native_executable.h
@@ -29,7 +29,6 @@
hipFunction_t function;
uint32_t block_size[3];
uint32_t shared_memory_size;
- IREE_TRACE(iree_hal_rocm_source_location_t source_location;)
} iree_hal_rocm_kernel_params_t;
// Creates an executable from a HSACO module. The module may contain several
@@ -47,6 +46,12 @@
hipFunction_t iree_hal_rocm_native_executable_for_entry_point(
iree_hal_executable_t* executable, int32_t entry_point);
+// Returns the source location for the given entry point. May be empty if not
+// available.
+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);
+
#ifdef __cplusplus
} // extern "C"
#endif // __cplusplus
