Adding split iree_hal_buffer_prepare_map_range/commit_map_range. (ire…

…e-org#18159)

This allows for the validation/offsetting/etc that acts only on metadata
to occur separate from the actual commit of the mapping where the host
storage is accessed. Code that wants to validate early (such as during
command buffer recording) can now do so on uncommitted buffers and have
a fast-path for when the underlying storage is available.

This will be used for native implementations of reusable command buffers
by allowing them to validate direct and indirect bindings during
recording even if the bindings are of uncommitted buffers (iree-org#18158) that
won't be available until execution-time.

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);

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