Freeze all statuses emitted by calls into dynamic modules (#16066)
When unloading a dynamic module we might inadvertently destroy any
status annotations it added, causing pitfalls/segfaults around status
usage. This changes all statuses emitted at the boundary of dynamic modules
to transfer ownership to the hosting context.
Fixes https://github.com/openxla/iree/issues/16046
diff --git a/runtime/src/iree/base/status.c b/runtime/src/iree/base/status.c
index 3c27d4c..937f058 100644
--- a/runtime/src/iree/base/status.c
+++ b/runtime/src/iree/base/status.c
@@ -6,6 +6,8 @@
#include "iree/base/status.h"
+#include "iree/base/attributes.h"
+
#if defined(IREE_PLATFORM_APPLE)
#include <dlfcn.h>
#include <execinfo.h>
@@ -1058,49 +1060,31 @@
#endif // has IREE_STATUS_FEATURE_ANNOTATIONS
-IREE_API_EXPORT bool iree_status_format(iree_status_t status,
- iree_host_size_t buffer_capacity,
- char* buffer,
- iree_host_size_t* out_buffer_length) {
+static bool iree_status_format_message(iree_status_t status,
+ iree_host_size_t buffer_capacity,
+ char* buffer,
+ iree_host_size_t* out_buffer_length,
+ bool has_prefix) {
*out_buffer_length = 0;
// Grab storage which may have a message and zero or more payloads.
- iree_status_storage_t* storage IREE_ATTRIBUTE_UNUSED =
- iree_status_storage(status);
+ iree_status_storage_t* storage = iree_status_storage(status);
+ // If no storage, nothing to do.
+ if (!storage) {
+ return true;
+ }
// Prefix with source location and status code string (may be 'OK').
iree_host_size_t buffer_length = 0;
- iree_status_code_t status_code = iree_status_code(status);
int n = 0;
-#if (IREE_STATUS_FEATURES & IREE_STATUS_FEATURE_SOURCE_LOCATION) != 0
- if (storage && storage->file) {
- n = snprintf(buffer ? buffer + buffer_length : NULL,
- buffer ? buffer_capacity - buffer_length : 0, "%s:%d: %s",
- storage->file, storage->line,
- iree_status_code_string(status_code));
- } else {
- n = snprintf(buffer ? buffer + buffer_length : NULL,
- buffer ? buffer_capacity - buffer_length : 0, "%s",
- iree_status_code_string(status_code));
- }
-#else
- n = snprintf(buffer ? buffer + buffer_length : NULL,
- buffer ? buffer_capacity - buffer_length : 0, "%s",
- iree_status_code_string(status_code));
-#endif // has IREE_STATUS_FEATURE_SOURCE_LOCATION
- if (IREE_UNLIKELY(n < 0)) {
- return false;
- } else if (buffer && n >= buffer_capacity - buffer_length) {
- buffer = NULL;
- }
- buffer_length += n;
#if (IREE_STATUS_FEATURES & IREE_STATUS_FEATURE_ANNOTATIONS) != 0
// Append base storage message.
if (storage && !iree_string_view_is_empty(storage->message)) {
n = snprintf(buffer ? buffer + buffer_length : NULL,
- buffer ? buffer_capacity - buffer_length : 0, "; %.*s",
- (int)storage->message.size, storage->message.data);
+ buffer ? buffer_capacity - buffer_length : 0,
+ has_prefix ? "; %.*s" : "%.*s", (int)storage->message.size,
+ storage->message.data);
if (IREE_UNLIKELY(n < 0)) {
return false;
} else if (buffer && n >= buffer_capacity - buffer_length) {
@@ -1109,6 +1093,11 @@
buffer_length += n;
}
#endif // has IREE_STATUS_FEATURE_ANNOTATIONS
+ if (IREE_UNLIKELY(n < 0)) {
+ return false;
+ } else if (buffer && n >= buffer_capacity) {
+ buffer = NULL;
+ }
#if IREE_STATUS_FEATURES != 0
// Append each payload separated by a newline.
@@ -1150,6 +1139,142 @@
return true;
}
+IREE_API_EXPORT bool iree_status_format(iree_status_t status,
+ iree_host_size_t buffer_capacity,
+ char* buffer,
+ iree_host_size_t* out_buffer_length) {
+ *out_buffer_length = 0;
+
+ // Grab storage which may have a message and zero or more payloads.
+ iree_status_storage_t* storage IREE_ATTRIBUTE_UNUSED =
+ iree_status_storage(status);
+
+ // Prefix with source location and status code string (may be 'OK').
+ iree_host_size_t prefix_buffer_length = 0;
+ iree_status_code_t status_code = iree_status_code(status);
+ int n = 0;
+#if (IREE_STATUS_FEATURES & IREE_STATUS_FEATURE_SOURCE_LOCATION) != 0
+ if (storage && storage->file) {
+ n = snprintf(buffer ? buffer : NULL, buffer ? buffer_capacity : 0,
+ "%s:%d: %s", storage->file, storage->line,
+ iree_status_code_string(status_code));
+ } else {
+ n = snprintf(buffer ? buffer : NULL, buffer ? buffer_capacity : 0, "%s",
+ iree_status_code_string(status_code));
+ }
+#else
+ n = snprintf(buffer ? buffer + prefix_buffer_length : NULL,
+ buffer ? buffer_capacity - prefix_buffer_length : 0, "%s",
+ iree_status_code_string(status_code));
+#endif // has IREE_STATUS_FEATURE_SOURCE_LOCATION
+ if (IREE_UNLIKELY(n < 0)) {
+ return false;
+ } else if (buffer && n >= buffer_capacity) {
+ buffer = NULL;
+ }
+ prefix_buffer_length += n;
+
+ iree_host_size_t message_buffer_length = 0;
+ bool ret = iree_status_format_message(
+ status, buffer_capacity, buffer ? buffer + prefix_buffer_length : NULL,
+ &message_buffer_length, /*has_prefix=*/true);
+ if (!ret) {
+ return false;
+ }
+ *out_buffer_length = message_buffer_length + prefix_buffer_length;
+ return true;
+}
+
+#if IREE_STATUS_FEATURES == 0
+IREE_API_EXPORT IREE_MUST_USE_RESULT iree_status_t
+iree_status_freeze(iree_status_t status) {
+ // Statuses are just codes; nothing to do.
+ return status;
+}
+#else
+IREE_API_EXPORT IREE_MUST_USE_RESULT iree_status_t
+iree_status_freeze(iree_status_t status) {
+ iree_status_code_t code = iree_status_code(status);
+ if (code == IREE_STATUS_OK) {
+ return iree_ok_status();
+ }
+
+ // Get the size of the formatted message alone. Source file annotations are
+ // handled separately.
+ iree_host_size_t message_buffer_size = 0;
+ if (IREE_UNLIKELY(!iree_status_format_message(
+ status, /*buffer_capacity=*/0,
+ /*buffer=*/NULL, &message_buffer_size, /*has_prefix=*/false))) {
+ iree_status_free(status);
+ return iree_status_from_code(code);
+ }
+ message_buffer_size++; // NUL
+
+ // Compute the storage size for the status with additional room to store the
+ // formatted message and source file location if present.
+ size_t unaligned_storage_size =
+ sizeof(iree_status_storage_t) + message_buffer_size;
+
+#if (IREE_STATUS_FEATURES & IREE_STATUS_FEATURE_SOURCE_LOCATION) != 0
+ // Grab storage for the source file location.
+ iree_status_storage_t* storage = iree_status_storage(status);
+ const char* file = NULL;
+ uint32_t line = 0;
+ if (storage) {
+ file = storage->file;
+ line = storage->line;
+ }
+ size_t file_storage_size = file ? strlen(file) + 1 : 0;
+ unaligned_storage_size += file_storage_size;
+#endif // has IREE_STATUS_FEATURE_SOURCE_LOCATION
+
+ size_t storage_alignment = (IREE_STATUS_CODE_MASK + 1);
+ size_t storage_size =
+ iree_host_align(unaligned_storage_size, storage_alignment);
+ iree_status_storage_t* new_storage =
+ (iree_status_storage_t*)iree_aligned_alloc(storage_alignment,
+ storage_size);
+ if (IREE_UNLIKELY(!new_storage)) {
+ iree_status_free(status);
+ return iree_status_from_code(code);
+ }
+ memset(new_storage, 0, sizeof(*new_storage));
+
+ char* message_data = (char*)new_storage + sizeof(iree_status_storage_t);
+ size_t res_length;
+ // Format the status message directly into the region allocated for it.
+ bool ret =
+ iree_status_format_message(status, message_buffer_size, message_data,
+ &res_length, /*has_prefix=*/false);
+ new_storage->message.size = message_buffer_size - 1;
+ new_storage->message.data =
+ (const char*)new_storage + sizeof(iree_status_storage_t);
+ iree_status_t new_status =
+ (iree_status_t)((uintptr_t)new_storage | (code & IREE_STATUS_CODE_MASK));
+
+ if (IREE_UNLIKELY(!ret)) {
+ iree_status_free(new_status);
+ iree_status_free(status);
+ return iree_status_from_code(code);
+ }
+
+#if (IREE_STATUS_FEATURES & IREE_STATUS_FEATURE_SOURCE_LOCATION) != 0
+ if (file) {
+ new_storage->file = storage->file;
+ char* storage_file = (char*)new_storage + sizeof(iree_status_storage_t) +
+ message_buffer_size;
+ // Copy the file into the storage allocated for it.
+ memcpy(storage_file, file, file_storage_size);
+ new_storage->file = (const char*)storage_file;
+ }
+ new_storage->line = line;
+#endif // has IREE_STATUS_FEATURE_SOURCE_LOCATION
+
+ iree_status_free(status);
+ return new_status;
+}
+#endif // has any IREE_STATUS_FEATURES
+
IREE_API_EXPORT bool iree_status_to_string(
iree_status_t status, const iree_allocator_t* allocator, char** out_buffer,
iree_host_size_t* out_buffer_length) {
diff --git a/runtime/src/iree/base/status.h b/runtime/src/iree/base/status.h
index f00635c..5cfe2eb 100644
--- a/runtime/src/iree/base/status.h
+++ b/runtime/src/iree/base/status.h
@@ -442,6 +442,12 @@
IREE_API_EXPORT IREE_MUST_USE_RESULT iree_status_t
iree_status_clone(iree_status_t status);
+// Freezes |status| into a new status instance. Formats all attached
+// annotations, payloads included, into a single message string allocated to
+// the new status. This always frees the original status.
+IREE_API_EXPORT IREE_MUST_USE_RESULT iree_status_t
+iree_status_freeze(iree_status_t status);
+
// Frees |status| if it has any associated storage.
IREE_API_EXPORT void iree_status_free(iree_status_t status);
diff --git a/runtime/src/iree/vm/dynamic/module.c b/runtime/src/iree/vm/dynamic/module.c
index 621b11f..a10677c 100644
--- a/runtime/src/iree/vm/dynamic/module.c
+++ b/runtime/src/iree/vm/dynamic/module.c
@@ -49,9 +49,9 @@
// Try to create the module, which may fail if the version is incompatible or
// the parameters are invalid.
IREE_RETURN_AND_END_ZONE_IF_ERROR(
- z0,
- create_fn(IREE_VM_DYNAMIC_MODULE_VERSION_LATEST, instance, param_count,
- params, module->allocator, &module->user_module));
+ z0, iree_status_freeze(create_fn(
+ IREE_VM_DYNAMIC_MODULE_VERSION_LATEST, instance, param_count,
+ params, module->allocator, &module->user_module)));
IREE_TRACE_ZONE_END(z0);
return iree_ok_status();
@@ -91,16 +91,16 @@
static iree_status_t IREE_API_PTR iree_vm_dynamic_module_get_module_attr(
void* self, iree_host_size_t index, iree_string_pair_t* out_attr) {
iree_vm_dynamic_module_t* module = (iree_vm_dynamic_module_t*)self;
- return module->user_module->get_module_attr(module->user_module->self, index,
- out_attr);
+ return iree_status_freeze(module->user_module->get_module_attr(
+ module->user_module->self, index, out_attr));
}
static iree_status_t iree_vm_dynamic_module_enumerate_dependencies(
void* self, iree_vm_module_dependency_callback_t callback,
void* user_data) {
iree_vm_dynamic_module_t* module = (iree_vm_dynamic_module_t*)self;
- return module->user_module->enumerate_dependencies(module->user_module->self,
- callback, user_data);
+ return iree_status_freeze(module->user_module->enumerate_dependencies(
+ module->user_module->self, callback, user_data));
}
static iree_status_t IREE_API_PTR iree_vm_dynamic_module_get_function(
@@ -108,9 +108,9 @@
iree_vm_function_t* out_function, iree_string_view_t* out_name,
iree_vm_function_signature_t* out_signature) {
iree_vm_dynamic_module_t* module = (iree_vm_dynamic_module_t*)self;
- IREE_RETURN_IF_ERROR(module->user_module->get_function(
+ IREE_RETURN_IF_ERROR(iree_status_freeze(module->user_module->get_function(
module->user_module->self, linkage, ordinal, out_function, out_name,
- out_signature));
+ out_signature)));
if (out_function) out_function->module = (iree_vm_module_t*)self;
return iree_ok_status();
}
@@ -119,8 +119,8 @@
void* self, iree_vm_function_linkage_t linkage, iree_host_size_t ordinal,
iree_host_size_t index, iree_string_pair_t* out_attr) {
iree_vm_dynamic_module_t* module = (iree_vm_dynamic_module_t*)self;
- return module->user_module->get_function_attr(
- module->user_module->self, linkage, ordinal, index, out_attr);
+ return iree_status_freeze(module->user_module->get_function_attr(
+ module->user_module->self, linkage, ordinal, index, out_attr));
}
static iree_status_t IREE_API_PTR iree_vm_dynamic_module_lookup_function(
@@ -128,9 +128,9 @@
const iree_vm_function_signature_t* expected_signature,
iree_vm_function_t* out_function) {
iree_vm_dynamic_module_t* module = (iree_vm_dynamic_module_t*)self;
- IREE_RETURN_IF_ERROR(module->user_module->lookup_function(
+ IREE_RETURN_IF_ERROR(iree_status_freeze(module->user_module->lookup_function(
module->user_module->self, linkage, name, expected_signature,
- out_function));
+ out_function)));
out_function->module = (iree_vm_module_t*)self;
return iree_ok_status();
}
@@ -141,8 +141,8 @@
iree_vm_source_location_t* out_source_location) {
iree_vm_dynamic_module_t* module = (iree_vm_dynamic_module_t*)self;
if (module->user_module->resolve_source_location) {
- return module->user_module->resolve_source_location(
- module->user_module->self, function, pc, out_source_location);
+ return iree_status_freeze(module->user_module->resolve_source_location(
+ module->user_module->self, function, pc, out_source_location));
}
return iree_status_from_code(IREE_STATUS_UNAVAILABLE);
}
@@ -151,8 +151,8 @@
iree_vm_dynamic_module_alloc_state(void* self, iree_allocator_t allocator,
iree_vm_module_state_t** out_module_state) {
iree_vm_dynamic_module_t* module = (iree_vm_dynamic_module_t*)self;
- return module->user_module->alloc_state(module->user_module->self, allocator,
- out_module_state);
+ return iree_status_freeze(module->user_module->alloc_state(
+ module->user_module->self, allocator, out_module_state));
}
static void IREE_API_PTR iree_vm_dynamic_module_free_state(
@@ -166,29 +166,29 @@
const iree_vm_function_t* function,
const iree_vm_function_signature_t* signature) {
iree_vm_dynamic_module_t* module = (iree_vm_dynamic_module_t*)self;
- return module->user_module->resolve_import(
- module->user_module->self, module_state, ordinal, function, signature);
+ return iree_status_freeze(module->user_module->resolve_import(
+ module->user_module->self, module_state, ordinal, function, signature));
}
static iree_status_t IREE_API_PTR iree_vm_dynamic_module_notify(
void* self, iree_vm_module_state_t* module_state, iree_vm_signal_t signal) {
iree_vm_dynamic_module_t* module = (iree_vm_dynamic_module_t*)self;
- return module->user_module->notify(module->user_module->self, module_state,
- signal);
+ return iree_status_freeze(module->user_module->notify(
+ module->user_module->self, module_state, signal));
}
static iree_status_t IREE_API_PTR iree_vm_dynamic_module_begin_call(
void* self, iree_vm_stack_t* stack, iree_vm_function_call_t call) {
iree_vm_dynamic_module_t* module = (iree_vm_dynamic_module_t*)self;
- return module->user_module->begin_call(module->user_module->self, stack,
- call);
+ return iree_status_freeze(
+ module->user_module->begin_call(module->user_module->self, stack, call));
}
static iree_status_t IREE_API_PTR iree_vm_dynamic_module_resume_call(
void* self, iree_vm_stack_t* stack, iree_byte_span_t call_results) {
iree_vm_dynamic_module_t* module = (iree_vm_dynamic_module_t*)self;
- return module->user_module->resume_call(module->user_module->self, stack,
- call_results);
+ return iree_status_freeze(module->user_module->resume_call(
+ module->user_module->self, stack, call_results));
}
static iree_status_t iree_vm_dynamic_module_create(
diff --git a/samples/custom_module/dynamic/module.cc b/samples/custom_module/dynamic/module.cc
index ac7687e..ad5ffff 100644
--- a/samples/custom_module/dynamic/module.cc
+++ b/samples/custom_module/dynamic/module.cc
@@ -148,6 +148,11 @@
return OkStatus();
}
+ // Prints the contents of the string to stdout.
+ Status ThrowError() {
+ return iree_make_status(IREE_STATUS_UNKNOWN, "Forced failure");
+ }
+
private:
// Allocator that the caller requested we use for any allocations we need to
// perform during operation.
@@ -159,6 +164,7 @@
vm::MakeNativeFunction("string.from_tensor",
&CustomModuleState::StringFromTensor),
vm::MakeNativeFunction("string.print", &CustomModuleState::StringPrint),
+ vm::MakeNativeFunction("error", &CustomModuleState::ThrowError),
};
// The module instance that will be allocated and reused across contexts.
diff --git a/samples/custom_module/dynamic/test/example.mlir b/samples/custom_module/dynamic/test/example.mlir
index 5b9db14..a0437d6 100644
--- a/samples/custom_module/dynamic/test/example.mlir
+++ b/samples/custom_module/dynamic/test/example.mlir
@@ -6,6 +6,14 @@
// RUN: --function=main | \
// RUN: FileCheck %s
+// RUN: ( iree-compile %s --iree-hal-target-backends=vmvx | \
+// RUN: iree-run-module \
+// RUN: --device=local-sync \
+// RUN: --module=$IREE_BINARY_DIR/samples/custom_module/dynamic/module$IREE_DYLIB_EXT@create_custom_module \
+// RUN: --module=- \
+// RUN: --function=error 2>&1 || [[ $? == 1 ]] ) | \
+// RUN: FileCheck %s --check-prefix=CERROR
+
module @example {
//===--------------------------------------------------------------------===//
// Imports
@@ -21,6 +29,9 @@
// Prints the contents of the string to stdout.
func.func private @custom.string.print(!custom.string)
+ // Always returns unknown status with a custom annotation.
+ func.func private @custom.error()
+
//===--------------------------------------------------------------------===//
// Sample methods
//===--------------------------------------------------------------------===//
@@ -41,4 +52,13 @@
return
}
+
+ // CERROR-LABEL: EXEC @error
+ func.func @error() {
+ // Show an example of dumping a custom error message.
+ // CERROR-NEXT: samples/custom_module/dynamic/module.cc:153: UNKNOWN; Forced failure; while invoking C++ function custom.error;
+ call @custom.error() : () -> ()
+ return
+ }
+
}
