Splitting utils from iree/hal/buffer_view.h. buffer_view.h now has the core type and its accessors while the alloc helpers and parsing/printing are moved to buffer_view_util.h.
diff --git a/iree/hal/BUILD b/iree/hal/BUILD index c87a07d..ece606e 100644 --- a/iree/hal/BUILD +++ b/iree/hal/BUILD
@@ -30,6 +30,8 @@ "buffer_heap_impl.h", "buffer_view.c", "buffer_view.h", + "buffer_view_util.c", + "buffer_view_util.h", "command_buffer.c", "command_buffer.h", "command_buffer_validation.c",
diff --git a/iree/hal/CMakeLists.txt b/iree/hal/CMakeLists.txt index 53528cb..eb0afe7 100644 --- a/iree/hal/CMakeLists.txt +++ b/iree/hal/CMakeLists.txt
@@ -25,6 +25,8 @@ "buffer_heap_impl.h" "buffer_view.c" "buffer_view.h" + "buffer_view_util.c" + "buffer_view_util.h" "command_buffer.c" "command_buffer.h" "command_buffer_validation.c"
diff --git a/iree/hal/api.h b/iree/hal/api.h index 33e2bdd..0ca7171 100644 --- a/iree/hal/api.h +++ b/iree/hal/api.h
@@ -12,6 +12,7 @@ #include "iree/hal/allocator.h" // IWYU pragma: export #include "iree/hal/buffer.h" // IWYU pragma: export #include "iree/hal/buffer_view.h" // IWYU pragma: export +#include "iree/hal/buffer_view_util.h" // IWYU pragma: export #include "iree/hal/command_buffer.h" // IWYU pragma: export #include "iree/hal/descriptor_set.h" // IWYU pragma: export #include "iree/hal/descriptor_set_layout.h" // IWYU pragma: export
diff --git a/iree/hal/buffer_view.c b/iree/hal/buffer_view.c index 17c7bbc..c338235 100644 --- a/iree/hal/buffer_view.c +++ b/iree/hal/buffer_view.c
@@ -6,14 +6,11 @@ #include "iree/hal/buffer_view.h" -#include <inttypes.h> -#include <stdbool.h> - #include "iree/base/api.h" #include "iree/base/tracing.h" #include "iree/hal/allocator.h" +#include "iree/hal/buffer_view_util.h" #include "iree/hal/resource.h" -#include "iree/hal/string_util.h" struct iree_hal_buffer_view_t { iree_atomic_ref_count_t ref_count; @@ -94,216 +91,6 @@ IREE_TRACE_ZONE_END(z0); } -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_allocate_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_buffer_usage_t allowed_usage, iree_const_byte_span_t initial_data, - iree_hal_buffer_view_t** out_buffer_view) { - IREE_ASSERT_ARGUMENT(allocator); - IREE_ASSERT_ARGUMENT(out_buffer_view); - IREE_TRACE_ZONE_BEGIN(z0); - - iree_device_size_t allocation_size = 0; - iree_status_t status = iree_hal_buffer_compute_view_size( - shape, shape_rank, element_type, encoding_type, &allocation_size); - - iree_hal_buffer_t* buffer = NULL; - if (iree_status_is_ok(status)) { - status = iree_hal_allocator_allocate_buffer(allocator, memory_type, - allowed_usage, allocation_size, - initial_data, &buffer); - } - - if (iree_status_is_ok(status)) { - status = iree_hal_buffer_view_create( - buffer, shape, shape_rank, element_type, encoding_type, - iree_hal_allocator_host_allocator(allocator), out_buffer_view); - } - - iree_hal_buffer_release(buffer); - IREE_TRACE_ZONE_END(z0); - return status; -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_heap_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_memory_access_t allowed_access, - iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, - iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view) { - IREE_ASSERT_ARGUMENT(allocator); - IREE_ASSERT_ARGUMENT(out_buffer_view); - IREE_TRACE_ZONE_BEGIN(z0); - - // NOTE: this will fail if the data cannot be imported into the allocator. - iree_hal_buffer_t* buffer = NULL; - iree_status_t status = iree_hal_allocator_wrap_buffer( - allocator, memory_type, allowed_access, allowed_usage, data, - data_allocator, &buffer); - - if (iree_status_is_ok(status)) { - status = iree_hal_buffer_view_create( - buffer, shape, shape_rank, element_type, encoding_type, - iree_hal_allocator_host_allocator(allocator), out_buffer_view); - } - - iree_hal_buffer_release(buffer); - IREE_TRACE_ZONE_END(z0); - return status; -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_or_clone_heap_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_memory_access_t allowed_access, - iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, - iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view) { - IREE_ASSERT_ARGUMENT(allocator); - - // Not all HAL implementations support wrapping buffers, and of those that do - // some may only support it in special situations such as when the buffer is - // not DEVICE_VISIBLE. The user application can query whether the wrapping is - // possible and decide to use alternative means of upload if it is not; we - // make no policy (other than validity) over what's best here. - iree_hal_buffer_compatibility_t compatibility = - iree_hal_allocator_query_buffer_compatibility( - allocator, memory_type, allowed_usage, IREE_HAL_BUFFER_USAGE_MAPPING, - (iree_device_size_t)data.data_length); - bool wrap_allowed = iree_all_bits_set( - compatibility, IREE_HAL_BUFFER_COMPATIBILITY_IMPORTABLE); - if (wrap_allowed) { - return iree_hal_buffer_view_wrap_heap_buffer( - allocator, shape, shape_rank, element_type, encoding_type, memory_type, - allowed_access, allowed_usage, data, data_allocator, out_buffer_view); - } else { - return iree_hal_buffer_view_allocate_buffer( - allocator, shape, shape_rank, element_type, encoding_type, memory_type, - allowed_usage, iree_make_const_byte_span(data.data, data.data_length), - out_buffer_view); - } -} - -static iree_status_t iree_hal_buffer_view_generate_buffer_in_situ( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_buffer_usage_t allowed_usage, - iree_hal_buffer_view_generator_callback_t callback, void* user_data, - iree_hal_buffer_view_t** out_buffer_view) { - // DO NOT SUBMIT - memory_type |= IREE_HAL_MEMORY_TYPE_HOST_VISIBLE; - - // Allocate the buffer view and entire buffer contents with the target memory - // type and the mapping bits. - iree_hal_buffer_view_t* buffer_view = NULL; - IREE_RETURN_IF_ERROR(iree_hal_buffer_view_allocate_buffer( - allocator, shape, shape_rank, element_type, encoding_type, memory_type, - allowed_usage | IREE_HAL_BUFFER_USAGE_MAPPING, - iree_const_byte_span_empty(), &buffer_view)); - - // Map the buffer into host-visible memory. - iree_hal_buffer_mapping_t buffer_mapping = {{0}}; - iree_status_t status = iree_hal_buffer_map_range( - iree_hal_buffer_view_buffer(buffer_view), IREE_HAL_MAPPING_MODE_SCOPED, - IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, 0, IREE_WHOLE_BUFFER, - &buffer_mapping); - - // Generate using the callback directly into the buffer. - if (iree_status_is_ok(status)) { - status = callback(&buffer_mapping, user_data); - } - - status = - iree_status_join(status, iree_hal_buffer_unmap_range(&buffer_mapping)); - if (iree_status_is_ok(status)) { - *out_buffer_view = buffer_view; - } else { - iree_hal_buffer_view_release(buffer_view); - } - return status; -} - -static iree_status_t iree_hal_buffer_view_generate_buffer_on_host( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_buffer_usage_t allowed_usage, iree_device_size_t allocation_size, - iree_hal_buffer_view_generator_callback_t callback, void* user_data, - iree_hal_buffer_view_t** out_buffer_view) { - // Allocate the host memory and generate the contents. - iree_allocator_t host_allocator = - iree_hal_allocator_host_allocator(allocator); - void* host_ptr = NULL; - IREE_RETURN_IF_ERROR( - iree_allocator_malloc(host_allocator, allocation_size, &host_ptr)); - iree_hal_buffer_mapping_t mapping = { - .contents = iree_make_byte_span(host_ptr, allocation_size), - }; - iree_status_t status = callback(&mapping, user_data); - if (!iree_status_is_ok(status)) { - iree_allocator_free(host_allocator, host_ptr); - return status; - } - - // Try to wrap the host allocation to avoid the extra allocation and copy - - // this call will either hang on to the memory or do the copy and immediately - // free it. - return iree_hal_buffer_view_wrap_or_clone_heap_buffer( - allocator, shape, shape_rank, element_type, encoding_type, memory_type, - IREE_HAL_MEMORY_ACCESS_ALL, allowed_usage, - iree_make_byte_span(host_ptr, allocation_size), host_allocator, - out_buffer_view); -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_generate_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_buffer_usage_t allowed_usage, - iree_hal_buffer_view_generator_callback_t callback, void* user_data, - iree_hal_buffer_view_t** out_buffer_view) { - IREE_ASSERT_ARGUMENT(allocator); - IREE_ASSERT_ARGUMENT(callback); - IREE_ASSERT_ARGUMENT(out_buffer_view); - IREE_TRACE_ZONE_BEGIN(z0); - - // Compute how large of an allocation we need to hold the whole view. - iree_device_size_t allocation_size = 0; - IREE_RETURN_AND_END_ZONE_IF_ERROR( - z0, iree_hal_buffer_compute_view_size(shape, shape_rank, element_type, - encoding_type, &allocation_size)); - - // If we can create the requested memory type with mapping then we'll do that - // and avoid needing to allocate the staging memory. If we can't get that - // memory type (or the allocator doesn't want us using it) then we'll fall - // back to allocation -> generation -> copy. - iree_hal_buffer_compatibility_t compatibility = - iree_hal_allocator_query_buffer_compatibility( - allocator, memory_type, allowed_usage, IREE_HAL_BUFFER_USAGE_MAPPING, - allocation_size); - bool is_mappable = iree_all_bits_set( - compatibility, IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE); - - iree_status_t status = iree_ok_status(); - if (is_mappable) { - // Compatible with allocate -> map -> generate. - status = iree_hal_buffer_view_generate_buffer_in_situ( - allocator, shape, shape_rank, element_type, encoding_type, memory_type, - allowed_usage, callback, user_data, out_buffer_view); - } else { - // Allocate host-local memory first and generate into that. - status = iree_hal_buffer_view_generate_buffer_on_host( - allocator, shape, shape_rank, element_type, encoding_type, memory_type, - allowed_usage, allocation_size, callback, user_data, out_buffer_view); - } - - IREE_TRACE_ZONE_END(z0); - return status; -} - IREE_API_EXPORT iree_hal_buffer_t* iree_hal_buffer_view_buffer( const iree_hal_buffer_view_t* buffer_view) { IREE_ASSERT_ARGUMENT(buffer_view); @@ -372,6 +159,10 @@ IREE_ASSERT_ARGUMENT(shape); if (shape_rank != buffer_view->shape_rank) { + // Rank changes require reallocation of the structure as we inline the + // shape dimensions. We could lighten this restriction to allow for rank + // reduction but knowing that rank changes aren't allowed is easier than + // remembering all the conditions in which they may be. return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "buffer view reshapes must have the same rank; " "target=%zu, existing=%zu", @@ -442,394 +233,3 @@ buffer_view->encoding_type, start_indices, indices_count, lengths, lengths_count, out_start_offset, out_length); } - -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_size( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, - iree_device_size_t* out_allocation_size) { - IREE_ASSERT_ARGUMENT(!shape_rank || shape); - IREE_ASSERT_ARGUMENT(out_allocation_size); - *out_allocation_size = 0; - - iree_device_size_t byte_length = 0; - - switch (encoding_type) { - case IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR: { - if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || - IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { - return iree_make_status( - IREE_STATUS_INVALID_ARGUMENT, - "opaque and sub-byte aligned element types cannot be indexed"); - } - byte_length = iree_hal_element_dense_byte_count(element_type); - for (iree_host_size_t i = 0; i < shape_rank; ++i) { - byte_length *= shape[i]; - } - break; - } - default: - return iree_make_status(IREE_STATUS_UNIMPLEMENTED, - "unimplemented encoding type size calculation"); - } - - *out_allocation_size = byte_length; - return iree_ok_status(); -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_offset( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* indices, - iree_host_size_t indices_count, iree_device_size_t* out_offset) { - IREE_ASSERT_ARGUMENT(shape); - IREE_ASSERT_ARGUMENT(indices); - IREE_ASSERT_ARGUMENT(out_offset); - *out_offset = 0; - if (IREE_UNLIKELY(encoding_type != IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR)) { - return iree_make_status( - IREE_STATUS_INVALID_ARGUMENT, - "only dense encodings support view range computation"); - } else if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || - IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { - return iree_make_status( - IREE_STATUS_INVALID_ARGUMENT, - "opaque and sub-byte aligned element types cannot be indexed"); - } else if (IREE_UNLIKELY(shape_rank != indices_count)) { - return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, - "shape rank/indices mismatch: %zu != %zu", - shape_rank, indices_count); - } - - iree_device_size_t offset = 0; - for (iree_host_size_t i = 0; i < indices_count; ++i) { - if (IREE_UNLIKELY(indices[i] >= shape[i])) { - return iree_make_status(IREE_STATUS_OUT_OF_RANGE, - "index[%zu] out of bounds: %d >= %d", i, - indices[i], shape[i]); - } - iree_device_size_t axis_offset = indices[i]; - for (iree_host_size_t j = i + 1; j < shape_rank; ++j) { - axis_offset *= shape[j]; - } - offset += axis_offset; - } - offset *= iree_hal_element_dense_byte_count(element_type); - - *out_offset = offset; - return iree_ok_status(); -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_range( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* start_indices, - iree_host_size_t indices_count, const iree_hal_dim_t* lengths, - iree_host_size_t lengths_count, iree_device_size_t* out_start_offset, - iree_device_size_t* out_length) { - IREE_ASSERT_ARGUMENT(shape); - IREE_ASSERT_ARGUMENT(start_indices); - IREE_ASSERT_ARGUMENT(lengths); - IREE_ASSERT_ARGUMENT(out_start_offset); - IREE_ASSERT_ARGUMENT(out_length); - *out_start_offset = 0; - *out_length = 0; - if (IREE_UNLIKELY(encoding_type != IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR)) { - return iree_make_status( - IREE_STATUS_INVALID_ARGUMENT, - "only dense encodings support view range computation"); - } else if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || - IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { - return iree_make_status( - IREE_STATUS_INVALID_ARGUMENT, - "opaque and sub-byte aligned element types cannot be indexed"); - } else if (IREE_UNLIKELY(indices_count != lengths_count)) { - return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, - "indices/lengths mismatch: %zu != %zu", - indices_count, lengths_count); - } else if (IREE_UNLIKELY(shape_rank != indices_count)) { - return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, - "shape rank/indices mismatch: %zu != %zu", - shape_rank, indices_count); - } - - iree_hal_dim_t* end_indices = - iree_alloca(shape_rank * sizeof(iree_hal_dim_t)); - iree_device_size_t element_size = - iree_hal_element_dense_byte_count(element_type); - iree_device_size_t subspan_length = element_size; - for (iree_host_size_t i = 0; i < lengths_count; ++i) { - subspan_length *= lengths[i]; - end_indices[i] = start_indices[i] + lengths[i] - 1; - } - - iree_device_size_t start_byte_offset = 0; - IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_offset( - shape, shape_rank, element_type, encoding_type, start_indices, - indices_count, &start_byte_offset)); - iree_device_size_t end_byte_offset = 0; - IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_offset( - shape, shape_rank, element_type, encoding_type, end_indices, shape_rank, - &end_byte_offset)); - - // Non-contiguous regions not yet implemented. Will be easier to detect when - // we have strides. - iree_device_size_t offset_length = - end_byte_offset - start_byte_offset + element_size; - if (subspan_length != offset_length) { - return iree_make_status( - IREE_STATUS_UNIMPLEMENTED, - "non-contiguous range region computation not implemented"); - } - - *out_start_offset = start_byte_offset; - *out_length = subspan_length; - return iree_ok_status(); -} - -typedef struct iree_hal_buffer_view_parse_params_t { - iree_string_view_t data_str; - iree_hal_element_type_t element_type; -} iree_hal_buffer_view_parse_params_t; -static iree_status_t iree_hal_buffer_view_parse_into( - iree_hal_buffer_mapping_t* mapping, void* user_data) { - iree_hal_buffer_view_parse_params_t* params = - (iree_hal_buffer_view_parse_params_t*)user_data; - return iree_hal_parse_buffer_elements(params->data_str, params->element_type, - mapping->contents); -} - -static iree_status_t iree_hal_buffer_view_parse_impl( - iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, - iree_hal_buffer_view_t** out_buffer_view) { - // Strip whitespace that may come along (linefeeds/etc). - value = iree_string_view_trim(value); - value = iree_string_view_strip_prefix(value, IREE_SV("\"")); - value = iree_string_view_strip_suffix(value, IREE_SV("\"")); - if (iree_string_view_is_empty(value)) { - // Empty lines are invalid; need at least the shape/type information. - *out_buffer_view = NULL; - return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "empty string input"); - } - - // The part of the string corresponding to the shape, e.g. 1x2x3. - iree_string_view_t shape_str = iree_string_view_empty(); - // The part of the string corresponding to the type, e.g. f32 - iree_string_view_t type_str = iree_string_view_empty(); - // The part of the string corresponding to the buffer data, e.g. 1 2 3 4 5 6 - iree_string_view_t data_str = iree_string_view_empty(); - - iree_string_view_t shape_and_type_str = value; - iree_string_view_split(value, '=', &shape_and_type_str, &data_str); - iree_host_size_t last_x_index = iree_string_view_find_last_of( - shape_and_type_str, IREE_SV("x"), IREE_STRING_VIEW_NPOS); - if (last_x_index == IREE_STRING_VIEW_NPOS) { - // Scalar. - type_str = shape_and_type_str; - } else { - // Has a shape. - shape_str = iree_string_view_substr(shape_and_type_str, 0, last_x_index); - type_str = iree_string_view_substr(shape_and_type_str, last_x_index + 1, - IREE_STRING_VIEW_NPOS); - } - - // AxBxC... - iree_host_size_t shape_rank = 0; - iree_status_t shape_result = - iree_hal_parse_shape(shape_str, 0, NULL, &shape_rank); - if (!iree_status_is_ok(shape_result) && - !iree_status_is_out_of_range(shape_result)) { - return shape_result; - } else if (shape_rank > 128) { - return iree_make_status( - IREE_STATUS_RESOURCE_EXHAUSTED, - "a shape rank of %zu is just a little bit excessive, eh?", shape_rank); - } - shape_result = iree_status_ignore(shape_result); - iree_hal_dim_t* shape = - (iree_hal_dim_t*)iree_alloca(shape_rank * sizeof(iree_hal_dim_t)); - IREE_RETURN_IF_ERROR( - iree_hal_parse_shape(shape_str, shape_rank, shape, &shape_rank)); - - // f32, i32, etc - iree_hal_element_type_t element_type = IREE_HAL_ELEMENT_TYPE_NONE; - IREE_RETURN_IF_ERROR(iree_hal_parse_element_type(type_str, &element_type)); - - // TODO(benvanik): allow specifying the encoding. - iree_hal_encoding_type_t encoding_type = - IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR; - - // Allocate the buffer from the provided allocator and parse directly into it. - iree_hal_buffer_view_parse_params_t params = { - .data_str = data_str, - .element_type = element_type, - }; - return iree_hal_buffer_view_generate_buffer( - buffer_allocator, shape, shape_rank, element_type, encoding_type, - IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL, - IREE_HAL_BUFFER_USAGE_DISPATCH | IREE_HAL_BUFFER_USAGE_TRANSFER, - iree_hal_buffer_view_parse_into, ¶ms, out_buffer_view); -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_parse( - iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, - iree_hal_buffer_view_t** out_buffer_view) { - IREE_ASSERT_ARGUMENT(buffer_allocator); - IREE_ASSERT_ARGUMENT(out_buffer_view); - *out_buffer_view = NULL; - IREE_TRACE_ZONE_BEGIN(z0); - iree_status_t status = - iree_hal_buffer_view_parse_impl(value, buffer_allocator, out_buffer_view); - IREE_TRACE_ZONE_END(z0); - return status; -} - -#define APPEND_CHAR(c) \ - { \ - if (buffer) { \ - if (buffer_length < buffer_capacity - 1) { \ - buffer[buffer_length] = c; \ - buffer[buffer_length + 1] = '\0'; \ - } else { \ - buffer = NULL; \ - } \ - } \ - ++buffer_length; \ - } - -static iree_status_t iree_hal_buffer_view_format_impl( - const iree_hal_buffer_view_t* buffer_view, - iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, - char* buffer, iree_host_size_t* out_buffer_length) { - if (out_buffer_length) { - *out_buffer_length = 0; - } - if (buffer && buffer_capacity) { - buffer[0] = 0; - } - - iree_host_size_t buffer_length = 0; - if (iree_hal_buffer_view_shape_rank(buffer_view) > 0) { - // Shape: 1x2x3 - iree_host_size_t shape_length = 0; - iree_status_t status = iree_hal_format_shape( - iree_hal_buffer_view_shape_dims(buffer_view), - iree_hal_buffer_view_shape_rank(buffer_view), - buffer ? buffer_capacity - buffer_length : 0, - buffer ? buffer + buffer_length : NULL, &shape_length); - buffer_length += shape_length; - if (iree_status_is_out_of_range(status)) { - status = iree_status_ignore(status); - buffer = NULL; - } else if (!iree_status_is_ok(status)) { - return status; - } - - // Separator: <shape>x<format> - APPEND_CHAR('x'); - } - - // Element type: f32 - iree_host_size_t element_type_length = 0; - iree_status_t status = iree_hal_format_element_type( - iree_hal_buffer_view_element_type(buffer_view), - buffer ? buffer_capacity - buffer_length : 0, - buffer ? buffer + buffer_length : NULL, &element_type_length); - buffer_length += element_type_length; - if (iree_status_is_out_of_range(status)) { - status = iree_status_ignore(status); - buffer = NULL; - } else if (!iree_status_is_ok(status)) { - return status; - } - - // TODO(benvanik): allow printing the encoding. - - // Separator: <meta>=<value> - APPEND_CHAR('='); - - // Buffer contents: 0 1 2 3 ... - iree_hal_buffer_mapping_t buffer_mapping = {{0}}; - IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( - iree_hal_buffer_view_buffer(buffer_view), IREE_HAL_MAPPING_MODE_SCOPED, - IREE_HAL_MEMORY_ACCESS_READ, 0, IREE_WHOLE_BUFFER, &buffer_mapping)); - iree_host_size_t elements_length = 0; - status = iree_hal_format_buffer_elements( - iree_make_const_byte_span(buffer_mapping.contents.data, - buffer_mapping.contents.data_length), - iree_hal_buffer_view_shape_dims(buffer_view), - iree_hal_buffer_view_shape_rank(buffer_view), - iree_hal_buffer_view_element_type(buffer_view), max_element_count, - buffer ? buffer_capacity - buffer_length : 0, - buffer ? buffer + buffer_length : NULL, &elements_length); - buffer_length += elements_length; - status = - iree_status_join(status, iree_hal_buffer_unmap_range(&buffer_mapping)); - if (iree_status_is_out_of_range(status)) { - status = iree_status_ignore(status); - buffer = NULL; - } else if (!iree_status_is_ok(status)) { - return status; - } - - if (out_buffer_length) { - *out_buffer_length = buffer_length; - } - return buffer ? iree_ok_status() - : iree_status_from_code(IREE_STATUS_OUT_OF_RANGE); -} - -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_format( - const iree_hal_buffer_view_t* buffer_view, - iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, - char* buffer, iree_host_size_t* out_buffer_length) { - IREE_ASSERT_ARGUMENT(buffer_view); - IREE_TRACE_ZONE_BEGIN(z0); - iree_status_t status = iree_hal_buffer_view_format_impl( - buffer_view, max_element_count, buffer_capacity, buffer, - out_buffer_length); - IREE_TRACE_ZONE_END(z0); - return status; -} - -// TODO(benvanik): streaming all the way down (needs string_util updates). -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_fprint( - FILE* file, const iree_hal_buffer_view_t* buffer_view, - iree_host_size_t max_element_count) { - IREE_ASSERT_ARGUMENT(file); - IREE_ASSERT_ARGUMENT(buffer_view); - IREE_TRACE_ZONE_BEGIN(z0); - - // Query the string length (in characters). - iree_host_size_t buffer_length = 0; - iree_status_t status = iree_hal_buffer_view_format( - buffer_view, max_element_count, 0, NULL, &buffer_length); - if (!iree_status_is_out_of_range(status)) { - IREE_TRACE_ZONE_END(z0); - return status; - } - - // Allocate scratch space to format in to. - // We should be streaming. - iree_allocator_t host_allocator = buffer_view->host_allocator; - iree_host_size_t buffer_capacity = buffer_length + 1; // NUL - char* buffer = NULL; - status = - iree_allocator_malloc(host_allocator, buffer_capacity, (void**)&buffer); - - // Format the buffer into the string storage. - if (iree_status_is_ok(status)) { - status = - iree_hal_buffer_view_format(buffer_view, max_element_count, - buffer_capacity, buffer, &buffer_length); - } - - // Dump to the file. - if (iree_status_is_ok(status)) { - fprintf(file, "%.*s", (int)buffer_length, buffer); - } - - iree_allocator_free(host_allocator, buffer); - IREE_TRACE_ZONE_END(z0); - return status; -}
diff --git a/iree/hal/buffer_view.h b/iree/hal/buffer_view.h index da1c10b..a6ee1d6 100644 --- a/iree/hal/buffer_view.h +++ b/iree/hal/buffer_view.h
@@ -7,9 +7,7 @@ #ifndef IREE_HAL_BUFFER_VIEW_H_ #define IREE_HAL_BUFFER_VIEW_H_ -#include <stdbool.h> #include <stdint.h> -#include <stdio.h> #include "iree/base/api.h" #include "iree/hal/buffer.h" @@ -156,35 +154,6 @@ typedef int32_t iree_hal_dim_t; //===----------------------------------------------------------------------===// -// Buffer view math -//===----------------------------------------------------------------------===// - -// Calculates the allocation size of a buffer view. -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_size( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, - iree_device_size_t* out_allocation_size); - -// Calculates a byte offset into a buffer at the given indices. -// Only works with densely-packed representations. -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_offset( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* indices, - size_t indices_count, iree_device_size_t* out_offset); - -// Calculates a byte range into a buffer of the given contiguous range. -// Only works with densely-packed representations. -IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_range( - const iree_hal_dim_t* shape, iree_host_size_t shape_rank, - iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* start_indices, - iree_host_size_t indices_count, const iree_hal_dim_t* lengths, - iree_host_size_t lengths_count, iree_device_size_t* out_start_offset, - iree_device_size_t* out_length); - -//===----------------------------------------------------------------------===// // iree_hal_buffer_view_t //===----------------------------------------------------------------------===// @@ -204,89 +173,6 @@ iree_hal_encoding_type_t encoding_type, iree_allocator_t host_allocator, iree_hal_buffer_view_t** out_buffer_view); -// Allocates a buffer from |allocator| and wraps it in a buffer view. -// -// This is equivalent to: -// 1. iree_hal_buffer_compute_view_size -// 2. iree_hal_allocator_allocate_buffer -// 3. iree_hal_buffer_view_create -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_allocate_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_buffer_usage_t allowed_usage, iree_const_byte_span_t initial_data, - iree_hal_buffer_view_t** out_buffer_view); - -// Imports a host buffer using |allocator| and wraps it in a buffer view. -// -// This is equivalent to: -// 1. iree_hal_allocator_wrap_buffer -// 2. iree_hal_buffer_view_create -// -// NOTE: not all buffers can be imported and not all allocators support -// importing. See iree_hal_allocator_wrap_buffer for more information. -// Fails if the buffer cannot be imported. -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_heap_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_memory_access_t allowed_access, - iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, - iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view); - -// Tries to import a host buffer using |allocator| and wrap it in a buffer view. -// If the buffer cannot be imported then a new buffer will be allocated and the -// source data will be copied into it. -// -// This is equivalent to: -// if iree_hal_allocator_query_buffer_compatibility ok: -// 1. iree_hal_allocator_wrap_buffer -// 2. iree_hal_buffer_view_create -// else: -// 1. iree_hal_allocator_allocate_buffer -// 2. iree_hal_buffer_write_data -// 3. iree_hal_buffer_view_create -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_or_clone_heap_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_memory_access_t allowed_access, - iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, - iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view); - -typedef iree_status_t(IREE_API_PTR* iree_hal_buffer_view_generator_callback_t)( - iree_hal_buffer_mapping_t* mapping, void* user_data); - -// Generates a buffer view with its initial contents produced by a callback. -// When host and device memory are shared this allows direct generation into the -// target device buffer. If not shared this can avoid expensive transfer mapping -// operations at the cost of a transient host memory allocation. The mapped host -// pointer passed to the callback is only valid within the callback. -// -// Buffers allocated like this do not need the IREE_HAL_BUFFER_USAGE_MAPPING bit -// set; it will be added automatically if the allocator needs it and otherwise -// the memory can remain unmappable (and thus fully device isolated). -// -// As this _may_ require allocation of the entire buffer content in host memory -// it is always preferable to stage and issue copy commands via the device -// queue. Even better is to do all generation on-device via dispatches without -// the need to ever transfer. Usage of this method should be limited to times -// where device-side generation isn't possible or memory consumption is not a -// concern. -// -// This is equivalent to: -// 1. iree_hal_buffer_compute_view_size -// 2. iree_hal_allocator_allocate_buffer -// 3. iree_hal_buffer_map_range + callback + iree_hal_buffer_unmap_range -// 4. iree_hal_buffer_view_create -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_generate_buffer( - iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, - iree_host_size_t shape_rank, iree_hal_element_type_t element_type, - iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, - iree_hal_buffer_usage_t allowed_usage, - iree_hal_buffer_view_generator_callback_t callback, void* user_data, - iree_hal_buffer_view_t** out_buffer_view); - // Retains the given |buffer_view| for the caller. IREE_API_EXPORT void iree_hal_buffer_view_retain( iree_hal_buffer_view_t* buffer_view); @@ -371,42 +257,6 @@ const iree_hal_dim_t* lengths, iree_host_size_t lengths_count, iree_device_size_t* out_start_offset, iree_device_size_t* out_length); -// Parses a serialized set of buffer elements in the canonical tensor format -// (the same as produced by iree_hal_buffer_view_format). The underlying buffer -// will be allocated with |buffer_allocator| as a host-local/device-visible -// buffer. -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_parse( - iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, - iree_hal_buffer_view_t** out_buffer_view); - -// TODO(#5413): enum for printing mode (include shape, precision). - -// Converts buffer view elements into a fully-specified string-form format like -// `2x4xi16=[[1 2][3 4]]`. -// -// |max_element_count| can be used to limit the total number of elements printed -// when the count may be large. Elided elements will be replaced with `...`. -// -// |buffer_capacity| defines the size of |buffer| in bytes and -// |out_buffer_length| will return the string length in characters. Returns -// IREE_STATUS_OUT_OF_RANGE if the buffer capacity is insufficient to hold the -// formatted elements and |out_buffer_length| will contain the required size. -// -// Follows the standard API string formatting rules. See iree/base/api.h. -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_format( - const iree_hal_buffer_view_t* buffer_view, - iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, - char* buffer, iree_host_size_t* out_buffer_length); - -// Prints buffer view elements into a fully-specified string-form format like -// `2x4xi16=[[1 2][3 4]]`. -// -// |max_element_count| can be used to limit the total number of elements printed -// when the count may be large. Elided elements will be replaced with `...`. -IREE_API_EXPORT iree_status_t iree_hal_buffer_view_fprint( - FILE* file, const iree_hal_buffer_view_t* buffer_view, - iree_host_size_t max_element_count); - //===----------------------------------------------------------------------===// // iree_hal_buffer_view_t implementation details //===----------------------------------------------------------------------===//
diff --git a/iree/hal/buffer_view_util.c b/iree/hal/buffer_view_util.c new file mode 100644 index 0000000..ae2cdac --- /dev/null +++ b/iree/hal/buffer_view_util.c
@@ -0,0 +1,626 @@ +// Copyright 2020 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 "iree/hal/buffer_view_util.h" + +#include <inttypes.h> +#include <stdbool.h> + +#include "iree/base/api.h" +#include "iree/base/tracing.h" +#include "iree/hal/allocator.h" +#include "iree/hal/resource.h" +#include "iree/hal/string_util.h" + +//===----------------------------------------------------------------------===// +// Buffer view math +//===----------------------------------------------------------------------===// + +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_size( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, + iree_device_size_t* out_allocation_size) { + IREE_ASSERT_ARGUMENT(!shape_rank || shape); + IREE_ASSERT_ARGUMENT(out_allocation_size); + *out_allocation_size = 0; + + iree_device_size_t byte_length = 0; + + switch (encoding_type) { + case IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR: { + if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || + IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "opaque and sub-byte aligned element types cannot be indexed"); + } + byte_length = iree_hal_element_dense_byte_count(element_type); + for (iree_host_size_t i = 0; i < shape_rank; ++i) { + byte_length *= shape[i]; + } + break; + } + default: + return iree_make_status(IREE_STATUS_UNIMPLEMENTED, + "unimplemented encoding type size calculation"); + } + + *out_allocation_size = byte_length; + return iree_ok_status(); +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_offset( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* indices, + iree_host_size_t indices_count, iree_device_size_t* out_offset) { + IREE_ASSERT_ARGUMENT(shape); + IREE_ASSERT_ARGUMENT(indices); + IREE_ASSERT_ARGUMENT(out_offset); + *out_offset = 0; + if (IREE_UNLIKELY(encoding_type != IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR)) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "only dense encodings support view range computation"); + } else if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || + IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "opaque and sub-byte aligned element types cannot be indexed"); + } else if (IREE_UNLIKELY(shape_rank != indices_count)) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "shape rank/indices mismatch: %zu != %zu", + shape_rank, indices_count); + } + + iree_device_size_t offset = 0; + for (iree_host_size_t i = 0; i < indices_count; ++i) { + if (IREE_UNLIKELY(indices[i] >= shape[i])) { + return iree_make_status(IREE_STATUS_OUT_OF_RANGE, + "index[%zu] out of bounds: %d >= %d", i, + indices[i], shape[i]); + } + iree_device_size_t axis_offset = indices[i]; + for (iree_host_size_t j = i + 1; j < shape_rank; ++j) { + axis_offset *= shape[j]; + } + offset += axis_offset; + } + offset *= iree_hal_element_dense_byte_count(element_type); + + *out_offset = offset; + return iree_ok_status(); +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_range( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* start_indices, + iree_host_size_t indices_count, const iree_hal_dim_t* lengths, + iree_host_size_t lengths_count, iree_device_size_t* out_start_offset, + iree_device_size_t* out_length) { + IREE_ASSERT_ARGUMENT(shape); + IREE_ASSERT_ARGUMENT(start_indices); + IREE_ASSERT_ARGUMENT(lengths); + IREE_ASSERT_ARGUMENT(out_start_offset); + IREE_ASSERT_ARGUMENT(out_length); + *out_start_offset = 0; + *out_length = 0; + if (IREE_UNLIKELY(encoding_type != IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR)) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "only dense encodings support view range computation"); + } else if (IREE_UNLIKELY(iree_hal_element_bit_count(element_type) == 0) || + IREE_UNLIKELY(!iree_hal_element_is_byte_aligned(element_type))) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "opaque and sub-byte aligned element types cannot be indexed"); + } else if (IREE_UNLIKELY(indices_count != lengths_count)) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "indices/lengths mismatch: %zu != %zu", + indices_count, lengths_count); + } else if (IREE_UNLIKELY(shape_rank != indices_count)) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "shape rank/indices mismatch: %zu != %zu", + shape_rank, indices_count); + } + + iree_hal_dim_t* end_indices = + iree_alloca(shape_rank * sizeof(iree_hal_dim_t)); + iree_device_size_t element_size = + iree_hal_element_dense_byte_count(element_type); + iree_device_size_t subspan_length = element_size; + for (iree_host_size_t i = 0; i < lengths_count; ++i) { + subspan_length *= lengths[i]; + end_indices[i] = start_indices[i] + lengths[i] - 1; + } + + iree_device_size_t start_byte_offset = 0; + IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_offset( + shape, shape_rank, element_type, encoding_type, start_indices, + indices_count, &start_byte_offset)); + iree_device_size_t end_byte_offset = 0; + IREE_RETURN_IF_ERROR(iree_hal_buffer_compute_view_offset( + shape, shape_rank, element_type, encoding_type, end_indices, shape_rank, + &end_byte_offset)); + + // Non-contiguous regions not yet implemented. Will be easier to detect when + // we have strides. + iree_device_size_t offset_length = + end_byte_offset - start_byte_offset + element_size; + if (subspan_length != offset_length) { + return iree_make_status( + IREE_STATUS_UNIMPLEMENTED, + "non-contiguous range region computation not implemented"); + } + + *out_start_offset = start_byte_offset; + *out_length = subspan_length; + return iree_ok_status(); +} + +//===----------------------------------------------------------------------===// +// Buffer view allocation and generation +//===----------------------------------------------------------------------===// + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_allocate_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, iree_const_byte_span_t initial_data, + iree_hal_buffer_view_t** out_buffer_view) { + IREE_ASSERT_ARGUMENT(allocator); + IREE_ASSERT_ARGUMENT(out_buffer_view); + IREE_TRACE_ZONE_BEGIN(z0); + + iree_device_size_t allocation_size = 0; + iree_status_t status = iree_hal_buffer_compute_view_size( + shape, shape_rank, element_type, encoding_type, &allocation_size); + + iree_hal_buffer_t* buffer = NULL; + if (iree_status_is_ok(status)) { + status = iree_hal_allocator_allocate_buffer(allocator, memory_type, + allowed_usage, allocation_size, + initial_data, &buffer); + } + + if (iree_status_is_ok(status)) { + status = iree_hal_buffer_view_create( + buffer, shape, shape_rank, element_type, encoding_type, + iree_hal_allocator_host_allocator(allocator), out_buffer_view); + } + + iree_hal_buffer_release(buffer); + IREE_TRACE_ZONE_END(z0); + return status; +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_heap_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_memory_access_t allowed_access, + iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, + iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view) { + IREE_ASSERT_ARGUMENT(allocator); + IREE_ASSERT_ARGUMENT(out_buffer_view); + IREE_TRACE_ZONE_BEGIN(z0); + + // NOTE: this will fail if the data cannot be imported into the allocator. + iree_hal_buffer_t* buffer = NULL; + iree_status_t status = iree_hal_allocator_wrap_buffer( + allocator, memory_type, allowed_access, allowed_usage, data, + data_allocator, &buffer); + + if (iree_status_is_ok(status)) { + status = iree_hal_buffer_view_create( + buffer, shape, shape_rank, element_type, encoding_type, + iree_hal_allocator_host_allocator(allocator), out_buffer_view); + } + + iree_hal_buffer_release(buffer); + IREE_TRACE_ZONE_END(z0); + return status; +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_or_clone_heap_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_memory_access_t allowed_access, + iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, + iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view) { + IREE_ASSERT_ARGUMENT(allocator); + + // Not all HAL implementations support wrapping buffers, and of those that do + // some may only support it in special situations such as when the buffer is + // not DEVICE_VISIBLE. The user application can query whether the wrapping is + // possible and decide to use alternative means of upload if it is not; we + // make no policy (other than validity) over what's best here. + iree_hal_buffer_compatibility_t compatibility = + iree_hal_allocator_query_buffer_compatibility( + allocator, memory_type, allowed_usage, IREE_HAL_BUFFER_USAGE_MAPPING, + (iree_device_size_t)data.data_length); + bool wrap_allowed = iree_all_bits_set( + compatibility, IREE_HAL_BUFFER_COMPATIBILITY_IMPORTABLE); + if (wrap_allowed) { + return iree_hal_buffer_view_wrap_heap_buffer( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + allowed_access, allowed_usage, data, data_allocator, out_buffer_view); + } else { + return iree_hal_buffer_view_allocate_buffer( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + allowed_usage, iree_make_const_byte_span(data.data, data.data_length), + out_buffer_view); + } +} + +static iree_status_t iree_hal_buffer_view_generate_buffer_in_situ( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, + iree_hal_buffer_view_generator_callback_t callback, void* user_data, + iree_hal_buffer_view_t** out_buffer_view) { + // Allocate the buffer view and entire buffer contents with the target memory + // type and the mapping bits. + iree_hal_buffer_view_t* buffer_view = NULL; + IREE_RETURN_IF_ERROR(iree_hal_buffer_view_allocate_buffer( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + allowed_usage | IREE_HAL_BUFFER_USAGE_MAPPING, + iree_const_byte_span_empty(), &buffer_view)); + + // Map the buffer into host-visible memory. + iree_hal_buffer_mapping_t buffer_mapping = {{0}}; + iree_status_t status = iree_hal_buffer_map_range( + iree_hal_buffer_view_buffer(buffer_view), IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE, 0, IREE_WHOLE_BUFFER, + &buffer_mapping); + + // Generate using the callback directly into the buffer. + if (iree_status_is_ok(status)) { + status = callback(&buffer_mapping, user_data); + } + + status = + iree_status_join(status, iree_hal_buffer_unmap_range(&buffer_mapping)); + if (iree_status_is_ok(status)) { + *out_buffer_view = buffer_view; + } else { + iree_hal_buffer_view_release(buffer_view); + } + return status; +} + +static iree_status_t iree_hal_buffer_view_generate_buffer_on_host( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, iree_device_size_t allocation_size, + iree_hal_buffer_view_generator_callback_t callback, void* user_data, + iree_hal_buffer_view_t** out_buffer_view) { + // Allocate the host memory and generate the contents. + iree_allocator_t host_allocator = + iree_hal_allocator_host_allocator(allocator); + void* host_ptr = NULL; + IREE_RETURN_IF_ERROR( + iree_allocator_malloc(host_allocator, allocation_size, &host_ptr)); + iree_hal_buffer_mapping_t mapping = { + .contents = iree_make_byte_span(host_ptr, allocation_size), + }; + iree_status_t status = callback(&mapping, user_data); + if (!iree_status_is_ok(status)) { + iree_allocator_free(host_allocator, host_ptr); + return status; + } + + // Try to wrap the host allocation to avoid the extra allocation and copy - + // this call will either hang on to the memory or do the copy and immediately + // free it. + return iree_hal_buffer_view_wrap_or_clone_heap_buffer( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + IREE_HAL_MEMORY_ACCESS_ALL, allowed_usage, + iree_make_byte_span(host_ptr, allocation_size), host_allocator, + out_buffer_view); +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_generate_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, + iree_hal_buffer_view_generator_callback_t callback, void* user_data, + iree_hal_buffer_view_t** out_buffer_view) { + IREE_ASSERT_ARGUMENT(allocator); + IREE_ASSERT_ARGUMENT(callback); + IREE_ASSERT_ARGUMENT(out_buffer_view); + IREE_TRACE_ZONE_BEGIN(z0); + + // Compute how large of an allocation we need to hold the whole view. + iree_device_size_t allocation_size = 0; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_hal_buffer_compute_view_size(shape, shape_rank, element_type, + encoding_type, &allocation_size)); + + // If we can create the requested memory type with mapping then we'll do that + // and avoid needing to allocate the staging memory. If we can't get that + // memory type (or the allocator doesn't want us using it) then we'll fall + // back to allocation -> generation -> copy. + iree_hal_buffer_compatibility_t compatibility = + iree_hal_allocator_query_buffer_compatibility( + allocator, memory_type, allowed_usage, IREE_HAL_BUFFER_USAGE_MAPPING, + allocation_size); + bool is_mappable = iree_all_bits_set( + compatibility, IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE); + + iree_status_t status = iree_ok_status(); + if (is_mappable) { + // Compatible with allocate -> map -> generate. + status = iree_hal_buffer_view_generate_buffer_in_situ( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + allowed_usage, callback, user_data, out_buffer_view); + } else { + // Allocate host-local memory first and generate into that. + status = iree_hal_buffer_view_generate_buffer_on_host( + allocator, shape, shape_rank, element_type, encoding_type, memory_type, + allowed_usage, allocation_size, callback, user_data, out_buffer_view); + } + + IREE_TRACE_ZONE_END(z0); + return status; +} + +//===----------------------------------------------------------------------===// +// Buffer view parsing and printing +//===----------------------------------------------------------------------===// + +typedef struct iree_hal_buffer_view_parse_params_t { + iree_string_view_t data_str; + iree_hal_element_type_t element_type; +} iree_hal_buffer_view_parse_params_t; +static iree_status_t iree_hal_buffer_view_parse_into( + iree_hal_buffer_mapping_t* mapping, void* user_data) { + iree_hal_buffer_view_parse_params_t* params = + (iree_hal_buffer_view_parse_params_t*)user_data; + return iree_hal_parse_buffer_elements(params->data_str, params->element_type, + mapping->contents); +} + +static iree_status_t iree_hal_buffer_view_parse_impl( + iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, + iree_hal_buffer_view_t** out_buffer_view) { + // Strip whitespace that may come along (linefeeds/etc). + value = iree_string_view_trim(value); + value = iree_string_view_strip_prefix(value, IREE_SV("\"")); + value = iree_string_view_strip_suffix(value, IREE_SV("\"")); + if (iree_string_view_is_empty(value)) { + // Empty lines are invalid; need at least the shape/type information. + *out_buffer_view = NULL; + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "empty string input"); + } + + // The part of the string corresponding to the shape, e.g. 1x2x3. + iree_string_view_t shape_str = iree_string_view_empty(); + // The part of the string corresponding to the type, e.g. f32 + iree_string_view_t type_str = iree_string_view_empty(); + // The part of the string corresponding to the buffer data, e.g. 1 2 3 4 5 6 + iree_string_view_t data_str = iree_string_view_empty(); + + iree_string_view_t shape_and_type_str = value; + iree_string_view_split(value, '=', &shape_and_type_str, &data_str); + iree_host_size_t last_x_index = iree_string_view_find_last_of( + shape_and_type_str, IREE_SV("x"), IREE_STRING_VIEW_NPOS); + if (last_x_index == IREE_STRING_VIEW_NPOS) { + // Scalar. + type_str = shape_and_type_str; + } else { + // Has a shape. + shape_str = iree_string_view_substr(shape_and_type_str, 0, last_x_index); + type_str = iree_string_view_substr(shape_and_type_str, last_x_index + 1, + IREE_STRING_VIEW_NPOS); + } + + // AxBxC... + iree_host_size_t shape_rank = 0; + iree_status_t shape_result = + iree_hal_parse_shape(shape_str, 0, NULL, &shape_rank); + if (!iree_status_is_ok(shape_result) && + !iree_status_is_out_of_range(shape_result)) { + return shape_result; + } else if (shape_rank > 128) { + return iree_make_status( + IREE_STATUS_RESOURCE_EXHAUSTED, + "a shape rank of %zu is just a little bit excessive, eh?", shape_rank); + } + shape_result = iree_status_ignore(shape_result); + iree_hal_dim_t* shape = + (iree_hal_dim_t*)iree_alloca(shape_rank * sizeof(iree_hal_dim_t)); + IREE_RETURN_IF_ERROR( + iree_hal_parse_shape(shape_str, shape_rank, shape, &shape_rank)); + + // f32, i32, etc + iree_hal_element_type_t element_type = IREE_HAL_ELEMENT_TYPE_NONE; + IREE_RETURN_IF_ERROR(iree_hal_parse_element_type(type_str, &element_type)); + + // TODO(benvanik): allow specifying the encoding. + iree_hal_encoding_type_t encoding_type = + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR; + + // Allocate the buffer from the provided allocator and parse directly into it. + iree_hal_buffer_view_parse_params_t params = { + .data_str = data_str, + .element_type = element_type, + }; + return iree_hal_buffer_view_generate_buffer( + buffer_allocator, shape, shape_rank, element_type, encoding_type, + IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE, + IREE_HAL_BUFFER_USAGE_DISPATCH | IREE_HAL_BUFFER_USAGE_TRANSFER | + IREE_HAL_BUFFER_USAGE_MAPPING, + iree_hal_buffer_view_parse_into, ¶ms, out_buffer_view); +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_parse( + iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, + iree_hal_buffer_view_t** out_buffer_view) { + IREE_ASSERT_ARGUMENT(buffer_allocator); + IREE_ASSERT_ARGUMENT(out_buffer_view); + *out_buffer_view = NULL; + IREE_TRACE_ZONE_BEGIN(z0); + iree_status_t status = + iree_hal_buffer_view_parse_impl(value, buffer_allocator, out_buffer_view); + IREE_TRACE_ZONE_END(z0); + return status; +} + +#define APPEND_CHAR(c) \ + { \ + if (buffer) { \ + if (buffer_length < buffer_capacity - 1) { \ + buffer[buffer_length] = c; \ + buffer[buffer_length + 1] = '\0'; \ + } else { \ + buffer = NULL; \ + } \ + } \ + ++buffer_length; \ + } + +static iree_status_t iree_hal_buffer_view_format_impl( + const iree_hal_buffer_view_t* buffer_view, + iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, + char* buffer, iree_host_size_t* out_buffer_length) { + if (out_buffer_length) { + *out_buffer_length = 0; + } + if (buffer && buffer_capacity) { + buffer[0] = 0; + } + + iree_host_size_t buffer_length = 0; + if (iree_hal_buffer_view_shape_rank(buffer_view) > 0) { + // Shape: 1x2x3 + iree_host_size_t shape_length = 0; + iree_status_t status = iree_hal_format_shape( + iree_hal_buffer_view_shape_dims(buffer_view), + iree_hal_buffer_view_shape_rank(buffer_view), + buffer ? buffer_capacity - buffer_length : 0, + buffer ? buffer + buffer_length : NULL, &shape_length); + buffer_length += shape_length; + if (iree_status_is_out_of_range(status)) { + status = iree_status_ignore(status); + buffer = NULL; + } else if (!iree_status_is_ok(status)) { + return status; + } + + // Separator: <shape>x<format> + APPEND_CHAR('x'); + } + + // Element type: f32 + iree_host_size_t element_type_length = 0; + iree_status_t status = iree_hal_format_element_type( + iree_hal_buffer_view_element_type(buffer_view), + buffer ? buffer_capacity - buffer_length : 0, + buffer ? buffer + buffer_length : NULL, &element_type_length); + buffer_length += element_type_length; + if (iree_status_is_out_of_range(status)) { + status = iree_status_ignore(status); + buffer = NULL; + } else if (!iree_status_is_ok(status)) { + return status; + } + + // TODO(benvanik): allow printing the encoding. + + // Separator: <meta>=<value> + APPEND_CHAR('='); + + // Buffer contents: 0 1 2 3 ... + iree_hal_buffer_mapping_t buffer_mapping = {{0}}; + IREE_RETURN_IF_ERROR(iree_hal_buffer_map_range( + iree_hal_buffer_view_buffer(buffer_view), IREE_HAL_MAPPING_MODE_SCOPED, + IREE_HAL_MEMORY_ACCESS_READ, 0, IREE_WHOLE_BUFFER, &buffer_mapping)); + iree_host_size_t elements_length = 0; + status = iree_hal_format_buffer_elements( + iree_make_const_byte_span(buffer_mapping.contents.data, + buffer_mapping.contents.data_length), + iree_hal_buffer_view_shape_dims(buffer_view), + iree_hal_buffer_view_shape_rank(buffer_view), + iree_hal_buffer_view_element_type(buffer_view), max_element_count, + buffer ? buffer_capacity - buffer_length : 0, + buffer ? buffer + buffer_length : NULL, &elements_length); + buffer_length += elements_length; + status = + iree_status_join(status, iree_hal_buffer_unmap_range(&buffer_mapping)); + if (iree_status_is_out_of_range(status)) { + status = iree_status_ignore(status); + buffer = NULL; + } else if (!iree_status_is_ok(status)) { + return status; + } + + if (out_buffer_length) { + *out_buffer_length = buffer_length; + } + return buffer ? iree_ok_status() + : iree_status_from_code(IREE_STATUS_OUT_OF_RANGE); +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_format( + const iree_hal_buffer_view_t* buffer_view, + iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, + char* buffer, iree_host_size_t* out_buffer_length) { + IREE_ASSERT_ARGUMENT(buffer_view); + IREE_TRACE_ZONE_BEGIN(z0); + iree_status_t status = iree_hal_buffer_view_format_impl( + buffer_view, max_element_count, buffer_capacity, buffer, + out_buffer_length); + IREE_TRACE_ZONE_END(z0); + return status; +} + +// TODO(benvanik): streaming all the way down (needs string_util updates). +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_fprint( + FILE* file, const iree_hal_buffer_view_t* buffer_view, + iree_host_size_t max_element_count, iree_allocator_t host_allocator) { + IREE_ASSERT_ARGUMENT(file); + IREE_ASSERT_ARGUMENT(buffer_view); + IREE_TRACE_ZONE_BEGIN(z0); + + // Query the string length (in characters). + iree_host_size_t buffer_length = 0; + iree_status_t status = iree_hal_buffer_view_format( + buffer_view, max_element_count, 0, NULL, &buffer_length); + if (!iree_status_is_out_of_range(status)) { + IREE_TRACE_ZONE_END(z0); + return status; + } + + // Allocate scratch space to format in to. + // We should be streaming. + iree_host_size_t buffer_capacity = buffer_length + 1; // NUL + char* buffer = NULL; + status = + iree_allocator_malloc(host_allocator, buffer_capacity, (void**)&buffer); + + // Format the buffer into the string storage. + if (iree_status_is_ok(status)) { + status = + iree_hal_buffer_view_format(buffer_view, max_element_count, + buffer_capacity, buffer, &buffer_length); + } + + // Dump to the file. + if (iree_status_is_ok(status)) { + fprintf(file, "%.*s", (int)buffer_length, buffer); + } + + iree_allocator_free(host_allocator, buffer); + IREE_TRACE_ZONE_END(z0); + return status; +}
diff --git a/iree/hal/buffer_view_util.h b/iree/hal/buffer_view_util.h new file mode 100644 index 0000000..08bf361 --- /dev/null +++ b/iree/hal/buffer_view_util.h
@@ -0,0 +1,184 @@ +// 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 + +#ifndef IREE_HAL_BUFFER_VIEW_UTIL_H_ +#define IREE_HAL_BUFFER_VIEW_UTIL_H_ + +#include <stdbool.h> +#include <stdint.h> +#include <stdio.h> + +#include "iree/base/api.h" +#include "iree/hal/buffer_view.h" + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +//===----------------------------------------------------------------------===// +// Buffer view math +//===----------------------------------------------------------------------===// + +// Calculates the allocation size of a buffer view. +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_size( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, + iree_device_size_t* out_allocation_size); + +// Calculates a byte offset into a buffer at the given indices. +// Only works with densely-packed representations. +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_offset( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* indices, + size_t indices_count, iree_device_size_t* out_offset); + +// Calculates a byte range into a buffer of the given contiguous range. +// Only works with densely-packed representations. +IREE_API_EXPORT iree_status_t iree_hal_buffer_compute_view_range( + const iree_hal_dim_t* shape, iree_host_size_t shape_rank, + iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, const iree_hal_dim_t* start_indices, + iree_host_size_t indices_count, const iree_hal_dim_t* lengths, + iree_host_size_t lengths_count, iree_device_size_t* out_start_offset, + iree_device_size_t* out_length); + +//===----------------------------------------------------------------------===// +// Buffer view allocation and generation +//===----------------------------------------------------------------------===// + +// Allocates a buffer from |allocator| and wraps it in a buffer view. +// +// This is equivalent to: +// 1. iree_hal_buffer_compute_view_size +// 2. iree_hal_allocator_allocate_buffer +// 3. iree_hal_buffer_view_create +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_allocate_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, iree_const_byte_span_t initial_data, + iree_hal_buffer_view_t** out_buffer_view); + +// Imports a host buffer using |allocator| and wraps it in a buffer view. +// +// This is equivalent to: +// 1. iree_hal_allocator_wrap_buffer +// 2. iree_hal_buffer_view_create +// +// NOTE: not all buffers can be imported and not all allocators support +// importing. See iree_hal_allocator_wrap_buffer for more information. +// Fails if the buffer cannot be imported. +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_heap_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_memory_access_t allowed_access, + iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, + iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view); + +// Tries to import a host buffer using |allocator| and wrap it in a buffer view. +// If the buffer cannot be imported then a new buffer will be allocated and the +// source data will be copied into it. +// +// This is equivalent to: +// if iree_hal_allocator_query_buffer_compatibility ok: +// 1. iree_hal_allocator_wrap_buffer +// 2. iree_hal_buffer_view_create +// else: +// 1. iree_hal_allocator_allocate_buffer +// 2. iree_hal_buffer_write_data +// 3. iree_hal_buffer_view_create +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_wrap_or_clone_heap_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_memory_access_t allowed_access, + iree_hal_buffer_usage_t allowed_usage, iree_byte_span_t data, + iree_allocator_t data_allocator, iree_hal_buffer_view_t** out_buffer_view); + +typedef iree_status_t(IREE_API_PTR* iree_hal_buffer_view_generator_callback_t)( + iree_hal_buffer_mapping_t* mapping, void* user_data); + +// Generates a buffer view with its initial contents produced by a callback. +// When host and device memory are shared this allows direct generation into the +// target device buffer. If not shared this can avoid expensive transfer mapping +// operations at the cost of a transient host memory allocation. The mapped host +// pointer passed to the callback is only valid within the callback. +// +// Buffers allocated like this do not need the IREE_HAL_BUFFER_USAGE_MAPPING bit +// set; it will be added automatically if the allocator needs it and otherwise +// the memory can remain unmappable (and thus fully device isolated). +// +// As this _may_ require allocation of the entire buffer content in host memory +// it is always preferable to stage and issue copy commands via the device +// queue. Even better is to do all generation on-device via dispatches without +// the need to ever transfer. Usage of this method should be limited to times +// where device-side generation isn't possible or memory consumption is not a +// concern. +// +// This is equivalent to: +// 1. iree_hal_buffer_compute_view_size +// 2. iree_hal_allocator_allocate_buffer +// 3. iree_hal_buffer_map_range + callback + iree_hal_buffer_unmap_range +// 4. iree_hal_buffer_view_create +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_generate_buffer( + iree_hal_allocator_t* allocator, const iree_hal_dim_t* shape, + iree_host_size_t shape_rank, iree_hal_element_type_t element_type, + iree_hal_encoding_type_t encoding_type, iree_hal_memory_type_t memory_type, + iree_hal_buffer_usage_t allowed_usage, + iree_hal_buffer_view_generator_callback_t callback, void* user_data, + iree_hal_buffer_view_t** out_buffer_view); + +//===----------------------------------------------------------------------===// +// Buffer view parsing and printing +//===----------------------------------------------------------------------===// + +// Parses a serialized set of buffer elements in the canonical tensor format +// (the same as produced by iree_hal_buffer_view_format). The underlying buffer +// will be allocated with |buffer_allocator| as a host-local/device-visible +// buffer. +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_parse( + iree_string_view_t value, iree_hal_allocator_t* buffer_allocator, + iree_hal_buffer_view_t** out_buffer_view); + +// TODO(#5413): enum for printing mode (include shape, precision). + +// Converts buffer view elements into a fully-specified string-form format like +// `2x4xi16=[[1 2][3 4]]`. +// +// |max_element_count| can be used to limit the total number of elements printed +// when the count may be large. Elided elements will be replaced with `...`. +// +// |buffer_capacity| defines the size of |buffer| in bytes and +// |out_buffer_length| will return the string length in characters. Returns +// IREE_STATUS_OUT_OF_RANGE if the buffer capacity is insufficient to hold the +// formatted elements and |out_buffer_length| will contain the required size. +// +// Follows the standard API string formatting rules. See iree/base/api.h. +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_format( + const iree_hal_buffer_view_t* buffer_view, + iree_host_size_t max_element_count, iree_host_size_t buffer_capacity, + char* buffer, iree_host_size_t* out_buffer_length); + +// Prints buffer view elements into a fully-specified string-form format like +// `2x4xi16=[[1 2][3 4]]`. +// +// |max_element_count| can be used to limit the total number of elements printed +// when the count may be large. Elided elements will be replaced with `...`. +// +// |host_allocator| will be used for any transient allocations required while +// printing. +IREE_API_EXPORT iree_status_t iree_hal_buffer_view_fprint( + FILE* file, const iree_hal_buffer_view_t* buffer_view, + iree_host_size_t max_element_count, iree_allocator_t host_allocator); + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus + +#endif // IREE_HAL_BUFFER_VIEW_UTIL_H_
diff --git a/iree/runtime/demo/hello_world_explained.c b/iree/runtime/demo/hello_world_explained.c index cbdfb09..589a1d1 100644 --- a/iree/runtime/demo/hello_world_explained.c +++ b/iree/runtime/demo/hello_world_explained.c
@@ -183,6 +183,8 @@ // in other sessions depending on whether they share a compatible device. iree_hal_allocator_t* device_allocator = iree_runtime_session_device_allocator(session); + iree_allocator_t host_allocator = + iree_runtime_session_host_allocator(session); iree_status_t status = iree_ok_status(); { // %arg0: tensor<4xf32> @@ -212,7 +214,7 @@ } if (iree_status_is_ok(status)) { IREE_IGNORE_ERROR(iree_hal_buffer_view_fprint( - stdout, arg0, /*max_element_count=*/4096)); + stdout, arg0, /*max_element_count=*/4096, host_allocator)); // Add to the call inputs list (which retains the buffer view). status = iree_runtime_call_inputs_push_back_buffer_view(&call, arg0); } @@ -237,7 +239,7 @@ } if (iree_status_is_ok(status)) { IREE_IGNORE_ERROR(iree_hal_buffer_view_fprint( - stdout, arg1, /*max_element_count=*/4096)); + stdout, arg1, /*max_element_count=*/4096, host_allocator)); status = iree_runtime_call_inputs_push_back_buffer_view(&call, arg1); } iree_hal_buffer_view_release(arg1); @@ -259,8 +261,8 @@ if (iree_status_is_ok(status)) { // This prints the buffer view out but an application could read its // contents, pass it to another call, etc. - status = - iree_hal_buffer_view_fprint(stdout, ret0, /*max_element_count=*/4096); + status = iree_hal_buffer_view_fprint( + stdout, ret0, /*max_element_count=*/4096, host_allocator); } iree_hal_buffer_view_release(ret0);
diff --git a/iree/runtime/demo/hello_world_terse.c b/iree/runtime/demo/hello_world_terse.c index 3a10d32..b333e60 100644 --- a/iree/runtime/demo/hello_world_terse.c +++ b/iree/runtime/demo/hello_world_terse.c
@@ -87,8 +87,9 @@ IREE_HAL_MEMORY_ACCESS_READ, IREE_HAL_BUFFER_USAGE_ALL, iree_make_byte_span((void*)arg0_data, sizeof(arg0_data)), iree_allocator_null(), &arg0)); - IREE_CHECK_OK( - iree_hal_buffer_view_fprint(stdout, arg0, /*max_element_count=*/4096)); + IREE_CHECK_OK(iree_hal_buffer_view_fprint( + stdout, arg0, /*max_element_count=*/4096, + iree_runtime_session_host_allocator(session))); IREE_CHECK_OK(iree_runtime_call_inputs_push_back_buffer_view(&call, arg0)); iree_hal_buffer_view_release(arg0); @@ -106,8 +107,9 @@ IREE_HAL_MEMORY_ACCESS_READ, IREE_HAL_BUFFER_USAGE_ALL, iree_make_byte_span((void*)arg1_data, sizeof(arg1_data)), iree_allocator_null(), &arg1)); - IREE_CHECK_OK( - iree_hal_buffer_view_fprint(stdout, arg1, /*max_element_count=*/4096)); + IREE_CHECK_OK(iree_hal_buffer_view_fprint( + stdout, arg1, /*max_element_count=*/4096, + iree_runtime_session_host_allocator(session))); IREE_CHECK_OK(iree_runtime_call_inputs_push_back_buffer_view(&call, arg1)); iree_hal_buffer_view_release(arg1); @@ -118,8 +120,9 @@ // -> tensor<4xf32> iree_hal_buffer_view_t* ret0 = NULL; IREE_CHECK_OK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &ret0)); - IREE_CHECK_OK( - iree_hal_buffer_view_fprint(stdout, ret0, /*max_element_count=*/4096)); + IREE_CHECK_OK(iree_hal_buffer_view_fprint( + stdout, ret0, /*max_element_count=*/4096, + iree_runtime_session_host_allocator(session))); iree_hal_buffer_view_release(ret0); iree_runtime_call_deinitialize(&call);
diff --git a/iree/samples/dynamic_shapes/main.c b/iree/samples/dynamic_shapes/main.c index 4f2594a..3808a33 100644 --- a/iree/samples/dynamic_shapes/main.c +++ b/iree/samples/dynamic_shapes/main.c
@@ -198,8 +198,10 @@ status = reduce_sum_2d(session, input, 6, &result_buffer_view); if (iree_status_is_ok(status)) { fprintf(stdout, "reduce_sum_2d([[1, 2, 3], [10, 20, 30]]): "); - status = iree_hal_buffer_view_fprint(stdout, result_buffer_view, - /*max_element_count=*/4096); + status = iree_hal_buffer_view_fprint( + stdout, result_buffer_view, + /*max_element_count=*/4096, + iree_runtime_session_host_allocator(session)); fprintf(stdout, "\n"); } iree_hal_buffer_view_release(result_buffer_view); @@ -213,8 +215,10 @@ if (iree_status_is_ok(status)) { fprintf(stdout, "reduce_sum_2d([[1, 2, 3], [10, 20, 30], [100, 200, 300]]): "); - status = iree_hal_buffer_view_fprint(stdout, result_buffer_view, - /*max_element_count=*/4096); + status = iree_hal_buffer_view_fprint( + stdout, result_buffer_view, + /*max_element_count=*/4096, + iree_runtime_session_host_allocator(session)); fprintf(stdout, "\n"); } iree_hal_buffer_view_release(result_buffer_view); @@ -227,8 +231,10 @@ status = add_one(session, input, 3, &result_buffer_view); if (iree_status_is_ok(status)) { fprintf(stdout, "add_one([1, 10, 100]): "); - status = iree_hal_buffer_view_fprint(stdout, result_buffer_view, - /*max_element_count=*/64); + status = iree_hal_buffer_view_fprint( + stdout, result_buffer_view, + /*max_element_count=*/64, + iree_runtime_session_host_allocator(session)); fprintf(stdout, "\n"); } iree_hal_buffer_view_release(result_buffer_view);
diff --git a/iree/tools/utils/trace_replay.c b/iree/tools/utils/trace_replay.c index c3cc197..fc1a185 100644 --- a/iree/tools/utils/trace_replay.c +++ b/iree/tools/utils/trace_replay.c
@@ -824,7 +824,8 @@ return iree_ok_status(); } -static iree_status_t iree_trace_replay_print_item(iree_vm_variant_t* value); +static iree_status_t iree_trace_replay_print_item( + iree_vm_variant_t* value, iree_allocator_t host_allocator); static iree_status_t iree_trace_replay_print_scalar(iree_vm_variant_t* value) { switch (value->type.value_type) { @@ -853,35 +854,34 @@ return iree_ok_status(); } -static iree_status_t iree_trace_replay_print_vm_list(iree_vm_list_t* list) { +static iree_status_t iree_trace_replay_print_vm_list( + iree_vm_list_t* list, iree_allocator_t host_allocator) { for (iree_host_size_t i = 0; i < iree_vm_list_size(list); ++i) { iree_vm_variant_t variant = iree_vm_variant_empty(); IREE_RETURN_IF_ERROR(iree_vm_list_get_variant(list, i, &variant), "variant %zu not present", i); - IREE_RETURN_IF_ERROR(iree_trace_replay_print_item(&variant)); + IREE_RETURN_IF_ERROR( + iree_trace_replay_print_item(&variant, host_allocator)); fprintf(stdout, "\n"); } return iree_ok_status(); } -static iree_status_t iree_trace_replay_print_hal_buffer_view( - iree_hal_buffer_view_t* buffer_view) { - return iree_hal_buffer_view_fprint(stdout, buffer_view, - /*max_element_count=*/1024); -} - -static iree_status_t iree_trace_replay_print_item(iree_vm_variant_t* value) { +static iree_status_t iree_trace_replay_print_item( + iree_vm_variant_t* value, iree_allocator_t host_allocator) { if (iree_vm_variant_is_value(*value)) { IREE_RETURN_IF_ERROR(iree_trace_replay_print_scalar(value)); } else if (iree_vm_variant_is_ref(*value)) { if (iree_hal_buffer_view_isa(value->ref)) { iree_hal_buffer_view_t* buffer_view = iree_hal_buffer_view_deref(value->ref); - IREE_RETURN_IF_ERROR( - iree_trace_replay_print_hal_buffer_view(buffer_view)); + IREE_RETURN_IF_ERROR(iree_hal_buffer_view_fprint( + stdout, buffer_view, + /*max_element_count=*/1024, host_allocator)); } else if (iree_vm_list_isa(value->ref)) { iree_vm_list_t* list = iree_vm_list_deref(value->ref); - IREE_RETURN_IF_ERROR(iree_trace_replay_print_vm_list(list)); + IREE_RETURN_IF_ERROR( + iree_trace_replay_print_vm_list(list, host_allocator)); } else { // TODO(benvanik): a way for ref types to describe themselves. fprintf(stdout, "(no printer)"); @@ -990,7 +990,8 @@ // Print the outputs. if (iree_status_is_ok(status)) { - status = iree_trace_replay_print_vm_list(output_list); + status = + iree_trace_replay_print_vm_list(output_list, replay->host_allocator); } iree_vm_list_release(output_list);