Reworking status printing to use the system allocator.
diff --git a/bindings/python/iree/runtime/status_utils.cc b/bindings/python/iree/runtime/status_utils.cc
index 30239fb..b12ccd2 100644
--- a/bindings/python/iree/runtime/status_utils.cc
+++ b/bindings/python/iree/runtime/status_utils.cc
@@ -24,6 +24,21 @@
}
}
+static std::string ApiStatusToString(iree_status_t status) {
+ iree_host_size_t buffer_length = 0;
+ if (IREE_UNLIKELY(!iree_status_format(status, /*buffer_capacity=*/0,
+ /*buffer=*/NULL, &buffer_length))) {
+ return "";
+ }
+ std::string result;
+ result.resize(buffer_length);
+ // NOTE: buffer capacity needs to be +1 for the NUL terminator in snprintf.
+ return iree_status_format(status, result.size() + 1,
+ const_cast<char*>(result.data()), &buffer_length)
+ ? result
+ : "";
+}
+
} // namespace
pybind11::error_already_set ApiStatusToPyExc(iree_status_t status,
@@ -31,15 +46,12 @@
assert(!iree_status_is_ok(status));
std::string full_message;
- char* iree_message;
- size_t iree_message_length;
- if (iree_status_to_string(status, &iree_message, &iree_message_length)) {
- full_message = std::string(message) + ": " +
- std::string(iree_message, iree_message_length);
- iree_allocator_free(iree_allocator_system(), iree_message);
- } else {
+ auto status_str = ApiStatusToString(status);
+ if (status_str.empty()) {
full_message = std::string(message) + ": " +
iree_status_code_string(iree_status_code(status));
+ } else {
+ full_message = std::string(message) + ": " + status_str;
}
PyErr_SetString(ApiStatusToPyExcClass(status), full_message.c_str());
diff --git a/iree/base/internal/flags.c b/iree/base/internal/flags.c
index 05a7b5f..5b47112 100644
--- a/iree/base/internal/flags.c
+++ b/iree/base/internal/flags.c
@@ -392,12 +392,8 @@
if (iree_status_is_ok(status)) return;
fprintf(stderr, "\x1b[31mFLAGS ERROR: (╯°□°)╯︵👻\x1b[0m\n");
- char* buffer = NULL;
- iree_host_size_t buffer_length = 0;
- iree_status_to_string(status, &buffer, &buffer_length);
- fprintf(stderr, "%.*s\n\n", (int)buffer_length, buffer);
+ iree_status_fprint(stderr, status);
fflush(stderr);
- iree_allocator_free(iree_allocator_system(), buffer);
exit(EXIT_FAILURE);
}
diff --git a/iree/base/status.c b/iree/base/status.c
index b3ba831..8f95ec6 100644
--- a/iree/base/status.c
+++ b/iree/base/status.c
@@ -729,27 +729,33 @@
return true;
}
-IREE_API_EXPORT bool iree_status_to_string(
- iree_status_t status, char** out_buffer,
- iree_host_size_t* out_buffer_length) {
+// Converts the status to an allocated string value using the given allocator.
+// The caller must free the buffer with |allocator|.
+static bool iree_status_to_string(iree_status_t status,
+ iree_allocator_t allocator, char** out_buffer,
+ iree_host_size_t* out_buffer_length) {
*out_buffer_length = 0;
iree_host_size_t buffer_length = 0;
if (IREE_UNLIKELY(!iree_status_format(status, /*buffer_capacity=*/0,
/*buffer=*/NULL, &buffer_length))) {
return false;
}
- // Buffer capacity needs to be +1 to account for the terminating null of
- // snprintf.
- buffer_length++;
- char* buffer = (char*)malloc(buffer_length);
- if (IREE_UNLIKELY(!buffer)) return false;
+
+ // Buffer capacity needs to be +1 for the NUL terminator (see snprintf).
+ char* buffer = NULL;
+ iree_status_t malloc_status =
+ iree_allocator_malloc(allocator, buffer_length + 1, (void**)&buffer);
+ if (!iree_status_is_ok(malloc_status)) {
+ iree_status_ignore(malloc_status);
+ return false;
+ }
bool ret =
- iree_status_format(status, buffer_length, buffer, out_buffer_length);
+ iree_status_format(status, buffer_length + 1, buffer, out_buffer_length);
if (ret) {
*out_buffer = buffer;
return true;
} else {
- free(buffer);
+ iree_allocator_free(allocator, buffer);
return false;
}
}
@@ -757,10 +763,15 @@
IREE_API_EXPORT void iree_status_fprint(FILE* file, iree_status_t status) {
// TODO(benvanik): better support for colors/etc - possibly move to logging.
// TODO(benvanik): do this without allocation by streaming the status.
+ iree_allocator_t allocator = iree_allocator_system();
char* status_buffer = NULL;
iree_host_size_t status_buffer_length = 0;
- iree_status_to_string(status, &status_buffer, &status_buffer_length);
- fprintf(file, "%.*s\n", (int)status_buffer_length, status_buffer);
- free(status_buffer);
+ if (iree_status_to_string(status, allocator, &status_buffer,
+ &status_buffer_length)) {
+ fprintf(file, "%.*s\n", (int)status_buffer_length, status_buffer);
+ iree_allocator_free(allocator, status_buffer);
+ } else {
+ fprintf(file, "(failed to format status)\n");
+ }
fflush(file);
}
diff --git a/iree/base/status.h b/iree/base/status.h
index 0ba8fa0..790c297 100644
--- a/iree/base/status.h
+++ b/iree/base/status.h
@@ -407,12 +407,6 @@
char* buffer,
iree_host_size_t* out_buffer_length);
-// Converts the status to an allocated string value.
-// The caller must free the buffer with the system allocator.
-IREE_API_EXPORT bool iree_status_to_string(iree_status_t status,
- char** out_buffer,
- iree_host_size_t* out_buffer_length);
-
// Prints |status| to the given |file| as a string with all available
// annotations. This will produce multiple lines of output and should be used
// only when dumping a status on failure.
