diff --git a/runtime/src/iree/hal/buffer.c b/runtime/src/iree/hal/buffer.c index 3d1f68a6b8fc..6b3327618d2a 100644 --- a/runtime/src/iree/hal/buffer.c +++ b/runtime/src/iree/hal/buffer.c @@ -959,6 +959,26 @@ IREE_API_EXPORT iree_status_t iree_hal_buffer_map_range( iree_hal_buffer_mapping_t* out_buffer_mapping) { IREE_ASSERT_ARGUMENT(buffer); IREE_ASSERT_ARGUMENT(out_buffer_mapping); + IREE_RETURN_IF_ERROR(iree_hal_buffer_prepare_map_range( + buffer, mapping_mode, memory_access, byte_offset, byte_length, + out_buffer_mapping)); + iree_status_t status = iree_hal_buffer_commit_map_range( + buffer, mapping_mode, memory_access, out_buffer_mapping); + if (!iree_status_is_ok(status)) { + // Scoped mappings retain the buffer until unmapped. + if (!out_buffer_mapping->impl.is_persistent) iree_hal_buffer_retain(buffer); + memset(out_buffer_mapping, 0, sizeof(*out_buffer_mapping)); + } + return status; +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_prepare_map_range( + iree_hal_buffer_t* buffer, iree_hal_mapping_mode_t mapping_mode, + iree_hal_memory_access_t memory_access, iree_device_size_t byte_offset, + iree_device_size_t byte_length, + iree_hal_buffer_mapping_t* out_buffer_mapping) { + IREE_ASSERT_ARGUMENT(buffer); + IREE_ASSERT_ARGUMENT(out_buffer_mapping); IREE_TRACE_ZONE_BEGIN(z0); memset(out_buffer_mapping, 0, sizeof(*out_buffer_mapping)); @@ -994,20 +1014,24 @@ IREE_API_EXPORT iree_status_t iree_hal_buffer_map_range( out_buffer_mapping->impl.allowed_access = memory_access; out_buffer_mapping->impl.is_persistent = is_persistent ? 1 : 0; out_buffer_mapping->impl.byte_offset = local_byte_offset; + out_buffer_mapping->contents = iree_make_byte_span(NULL, local_byte_length); - iree_status_t status = _VTABLE_DISPATCH(buffer, map_range)( - buffer, mapping_mode, memory_access, out_buffer_mapping->impl.byte_offset, - local_byte_length, out_buffer_mapping); - - if (iree_status_is_ok(status)) { - // Scoped mappings retain the buffer until unmapped. - if (!is_persistent) iree_hal_buffer_retain(buffer); - } else { - memset(out_buffer_mapping, 0, sizeof(*out_buffer_mapping)); - } + // Scoped mappings retain the buffer until unmapped. + if (!is_persistent) iree_hal_buffer_retain(buffer); IREE_TRACE_ZONE_END(z0); - return status; + return iree_ok_status(); +} + +IREE_API_EXPORT iree_status_t iree_hal_buffer_commit_map_range( + iree_hal_buffer_t* buffer, iree_hal_mapping_mode_t mapping_mode, + iree_hal_memory_access_t memory_access, + iree_hal_buffer_mapping_t* buffer_mapping) { + IREE_ASSERT_ARGUMENT(buffer); + IREE_ASSERT_ARGUMENT(buffer_mapping); + return _VTABLE_DISPATCH(buffer, map_range)( + buffer, mapping_mode, memory_access, buffer_mapping->impl.byte_offset, + buffer_mapping->contents.data_length, buffer_mapping); } IREE_API_EXPORT iree_status_t @@ -1017,9 +1041,12 @@ iree_hal_buffer_unmap_range(iree_hal_buffer_mapping_t* buffer_mapping) { if (!buffer) return iree_ok_status(); IREE_TRACE_ZONE_BEGIN(z0); - iree_status_t status = _VTABLE_DISPATCH(buffer, unmap_range)( - buffer, buffer_mapping->impl.byte_offset, - buffer_mapping->contents.data_length, buffer_mapping); + iree_status_t status = iree_ok_status(); + if (buffer_mapping->contents.data != NULL) { + status = _VTABLE_DISPATCH(buffer, unmap_range)( + buffer, buffer_mapping->impl.byte_offset, + buffer_mapping->contents.data_length, buffer_mapping); + } if (!buffer_mapping->impl.is_persistent) { iree_hal_buffer_release(buffer); diff --git a/runtime/src/iree/hal/buffer.h b/runtime/src/iree/hal/buffer.h index c8322289053a..82c1f478df22 100644 --- a/runtime/src/iree/hal/buffer.h +++ b/runtime/src/iree/hal/buffer.h @@ -830,6 +830,45 @@ IREE_API_EXPORT iree_status_t iree_hal_buffer_map_range( iree_device_size_t byte_length, iree_hal_buffer_mapping_t* out_buffer_mapping); +// Prepares for mapping the buffer to be accessed as a host pointer into +// |out_buffer_mapping|. The byte offset and byte length may be adjusted for +// device alignment. The output data pointer will be properly aligned to the +// start of the data. Fails if the memory could not be mapped (invalid access +// type, invalid range, or unsupported memory type). +// +// Requires that the buffer has the IREE_HAL_BUFFER_USAGE_MAPPING bit set. +// If the buffer is not IREE_HAL_MEMORY_TYPE_HOST_COHERENT then the caller must +// invalidate the byte range they want to access to update the visibility of the +// mapped memory. +// +// This is the first part of a paired operation with +// iree_hal_buffer_commit_map_range. This allows callers to prepare for mapping +// (performing all of the validation) without actually resolving the host +// pointer yet. Once prepared the mapping must be unmapped with +// iree_hal_buffer_unmap_range even if it is never committed. +// +// Callers are allowed to prepare mappings prior to the |buffer| having +// allocated storage. Committing the mapping requires that storage has been +// bound for the duration the mapping will be live. +// +// Example usage: +// iree_hal_buffer_prepare_map_range(..., &mapping); +// if (maybe) iree_hal_buffer_commit_map_range(..., &mapping); +// iree_hal_buffer_unmap_range(&mapping); +IREE_API_EXPORT iree_status_t iree_hal_buffer_prepare_map_range( + iree_hal_buffer_t* buffer, iree_hal_mapping_mode_t mapping_mode, + iree_hal_memory_access_t memory_access, iree_device_size_t byte_offset, + iree_device_size_t byte_length, + iree_hal_buffer_mapping_t* out_buffer_mapping); + +// Commits a mapping operation from iree_hal_buffer_prepare_map_range. +// May fail for internal reasons but not any of those previously validated +// during preparation. +IREE_API_EXPORT iree_status_t iree_hal_buffer_commit_map_range( + iree_hal_buffer_t* buffer, iree_hal_mapping_mode_t mapping_mode, + iree_hal_memory_access_t memory_access, + iree_hal_buffer_mapping_t* buffer_mapping); + // Unmaps the buffer as was previously mapped to |buffer_mapping|. // // If the buffer is not IREE_HAL_MEMORY_TYPE_HOST_COHERENT then the caller must