Adding iree_hal_deferred_buffer_t helper. (#18158)
This is effectively just a subspan buffer but with the ability to
commit/decommit the backing storage. It's got a lot of footguns but
should mostly be used very carefully by binding layers or HAL
implementations that manage them.
This will be used to implement async allocations on CPU and any other
backend that does not natively support virtual memory. The intended
usage is for `iree_hal_device_queue_alloca` to reserve a deferred buffer
and then enqueue a commit in the queue timeline while
`iree_hal_device_queue_dealloca` will enqueue a decommit.
diff --git a/runtime/src/iree/hal/allocator.h b/runtime/src/iree/hal/allocator.h
index 254bedb..de0ec30 100644
--- a/runtime/src/iree/hal/allocator.h
+++ b/runtime/src/iree/hal/allocator.h
@@ -64,89 +64,6 @@
iree_device_size_t min_alignment;
} iree_hal_allocator_memory_heap_t;
-// Parameters defining how a buffer should be allocated.
-//
-// Designed to be zero-initialized: any field with a 0 value will be assigned
-// a default as indicated in the field description.
-//
-// For ergonomics when used from C++ w/o named initializers the first field is
-// the most commonly used so that it can be initialized by location:
-// some_fn(..., {IREE_HAL_BUFFER_USAGE_FOO}, ...)
-typedef struct iree_hal_buffer_params_t {
- // Specifies the usage allowed by HAL APIs and aids in memory placement.
- // Devices may have different memory types for different usage and require
- // the intended usage to be declared upon allocation. It's always best to
- // limit the allowed usage bits to precisely what the actual usage will be to
- // avoid additional copies, synchronization, and expensive emulation.
- //
- // If 0 then the usage will default to all usage modes.
- iree_hal_buffer_usage_t usage;
-
- // Specifies the access allowed to the memory via the HAL APIs.
- // For example, if the IREE_HAL_MEMORY_ACCESS_WRITE bit is not set then any
- // API call that would write to the memory will fail (such as
- // iree_hal_command_buffer_update_buffer). This does not limit any untrusted
- // dispatch or external use of the buffer and should not be treated as a
- // memory protection mechanism.
- //
- // If 0 then the access will be set as IREE_HAL_MEMORY_ACCESS_ALL.
- iree_hal_memory_access_t access;
-
- // Specifies the memory type properties used for selecting a memory space.
- // This should often be IREE_HAL_MEMORY_TYPE_OPTIMAL to allow the allocator
- // to place the allocation based on usage bits but can be specified if the
- // exact memory type must be used for compatibility with external code.
- //
- // If 0 then the type will be set as IREE_HAL_MEMORY_TYPE_OPTIMAL.
- iree_hal_memory_type_t type;
-
- // Queue affinity bitmap indicating which queues may access this buffer.
- // For NUMA devices this can be used to more tightly scope the allocation to
- // particular device memory and provide better pool placement. When a device
- // supports peering or replication the affinity bitmap will be used to choose
- // which subdevices require configuration.
- //
- // If 0 then the buffer will be available on any queue as if
- // IREE_HAL_QUEUE_AFFINITY_ANY was specified.
- iree_hal_queue_affinity_t queue_affinity;
-
- // Minimum alignment, in bytes, of the resulting allocation.
- // The actual alignment may be any value greater-than-or-equal-to this value.
- //
- // If 0 then the alignment will be decided by the allocator based on optimal
- // device parameters.
- iree_device_size_t min_alignment;
-} iree_hal_buffer_params_t;
-
-// Canonicalizes |params| fields when zero initialization is used.
-static inline void iree_hal_buffer_params_canonicalize(
- iree_hal_buffer_params_t* params) {
- if (!params->usage) {
- params->usage = IREE_HAL_BUFFER_USAGE_DEFAULT;
- }
- if (!params->access) {
- params->access = IREE_HAL_MEMORY_ACCESS_ALL;
- }
- if (!params->type) {
- params->type = IREE_HAL_MEMORY_TYPE_OPTIMAL;
- }
- if (!params->queue_affinity) {
- params->queue_affinity = IREE_HAL_QUEUE_AFFINITY_ANY;
- }
-}
-
-// Returns |params| with the given |usage| bits OR'ed in.
-static inline iree_hal_buffer_params_t iree_hal_buffer_params_with_usage(
- const iree_hal_buffer_params_t params, iree_hal_buffer_usage_t usage) {
- iree_hal_buffer_params_t result = params;
- if (!result.usage) {
- result.usage =
- IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE | IREE_HAL_BUFFER_USAGE_TRANSFER;
- }
- result.usage |= usage;
- return result;
-}
-
// A bitfield indicating compatible behavior for buffers in an allocator.
enum iree_hal_buffer_compatibility_bits_t {
// Indicates (in the absence of other bits) the buffer is not compatible with
diff --git a/runtime/src/iree/hal/buffer.c b/runtime/src/iree/hal/buffer.c
index c66835b..3d1f68a 100644
--- a/runtime/src/iree/hal/buffer.c
+++ b/runtime/src/iree/hal/buffer.c
@@ -151,6 +151,7 @@
iree_allocator_t host_allocator, iree_hal_buffer_t** out_buffer) {
IREE_ASSERT_ARGUMENT(allocated_buffer);
IREE_ASSERT_ARGUMENT(out_buffer);
+ *out_buffer = NULL;
IREE_TRACE_ZONE_BEGIN(z0);
iree_hal_buffer_t* buffer = NULL;
@@ -193,7 +194,6 @@
static iree_status_t iree_hal_subspan_buffer_unmap_range(
iree_hal_buffer_t* buffer, iree_device_size_t local_byte_offset,
iree_device_size_t local_byte_length, iree_hal_buffer_mapping_t* mapping) {
- if (!buffer->allocated_buffer) return iree_ok_status();
return _VTABLE_DISPATCH(buffer->allocated_buffer, unmap_range)(
buffer->allocated_buffer, local_byte_offset, local_byte_length, mapping);
}
@@ -222,6 +222,148 @@
};
//===----------------------------------------------------------------------===//
+// iree_hal_deferred_buffer_t
+//===----------------------------------------------------------------------===//
+
+typedef struct iree_hal_deferred_buffer_t {
+ iree_hal_buffer_t base;
+ iree_hal_queue_affinity_t queue_affinity;
+ iree_device_size_t min_alignment;
+} iree_hal_deferred_buffer_t;
+
+static const iree_hal_buffer_vtable_t iree_hal_deferred_buffer_vtable;
+
+IREE_API_EXPORT iree_status_t iree_hal_deferred_buffer_create_reserved(
+ iree_hal_allocator_t* device_allocator, iree_device_size_t allocation_size,
+ iree_device_size_t byte_offset, iree_device_size_t byte_length,
+ iree_hal_buffer_params_t params, iree_allocator_t host_allocator,
+ iree_hal_buffer_t** out_buffer) {
+ IREE_ASSERT_ARGUMENT(out_buffer);
+ *out_buffer = NULL;
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ iree_hal_deferred_buffer_t* buffer = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0,
+ iree_allocator_malloc(host_allocator, sizeof(*buffer), (void**)&buffer));
+ iree_hal_buffer_initialize(host_allocator, device_allocator, NULL,
+ allocation_size, byte_offset, byte_length,
+ params.type, params.access, params.usage,
+ &iree_hal_deferred_buffer_vtable, &buffer->base);
+ buffer->queue_affinity = params.queue_affinity;
+ buffer->min_alignment = params.min_alignment;
+ *out_buffer = &buffer->base;
+
+ IREE_TRACE_ZONE_END(z0);
+ return iree_ok_status();
+}
+
+static void iree_hal_deferred_buffer_destroy(iree_hal_buffer_t* base_buffer) {
+ iree_allocator_t host_allocator = base_buffer->host_allocator;
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ if (base_buffer->allocated_buffer) {
+ iree_hal_buffer_release(base_buffer->allocated_buffer);
+ }
+ iree_allocator_free(host_allocator, base_buffer);
+
+ IREE_TRACE_ZONE_END(z0);
+}
+
+IREE_API_EXPORT iree_status_t
+iree_hal_deferred_buffer_commit(iree_hal_buffer_t* base_buffer) {
+ iree_hal_deferred_buffer_t* buffer = (iree_hal_deferred_buffer_t*)base_buffer;
+ if (IREE_UNLIKELY(base_buffer->allocated_buffer)) {
+ // Already committed - no-op.
+ return iree_ok_status();
+ }
+ IREE_TRACE_ZONE_BEGIN(z0);
+ iree_hal_buffer_params_t params = {
+ .usage = base_buffer->allowed_usage,
+ .access = base_buffer->allowed_access,
+ .type = base_buffer->memory_type,
+ .queue_affinity = buffer->queue_affinity,
+ .min_alignment = buffer->min_alignment,
+ };
+ iree_status_t status = iree_hal_allocator_allocate_buffer(
+ base_buffer->device_allocator, params, base_buffer->allocation_size,
+ &base_buffer->allocated_buffer);
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
+
+IREE_API_EXPORT iree_status_t
+iree_hal_deferred_buffer_decommit(iree_hal_buffer_t* buffer) {
+ IREE_TRACE_ZONE_BEGIN(z0);
+ if (IREE_LIKELY(buffer->allocated_buffer)) {
+ iree_hal_buffer_release(buffer->allocated_buffer);
+ buffer->allocated_buffer = NULL;
+ }
+ IREE_TRACE_ZONE_END(z0);
+ return iree_ok_status();
+}
+
+static iree_status_t iree_hal_deferred_buffer_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 local_byte_offset, iree_device_size_t local_byte_length,
+ iree_hal_buffer_mapping_t* mapping) {
+ if (IREE_UNLIKELY(!buffer->allocated_buffer)) {
+ // Performance warning: this is likely to be happening synchronously in the
+ // caller in an unexpected way. We could FAILED_PRECONDITION if we wanted
+ // to be strict but by doing this on-demand we allow deferred buffers to be
+ // used with callers that may not know that this is a reserved deferred
+ // buffer (particularly useful for outputs/copy targets).
+ IREE_RETURN_IF_ERROR(iree_hal_deferred_buffer_commit(buffer));
+ }
+ return _VTABLE_DISPATCH(buffer->allocated_buffer, map_range)(
+ buffer->allocated_buffer, mapping_mode, memory_access, local_byte_offset,
+ local_byte_length, mapping);
+}
+
+static iree_status_t iree_hal_deferred_buffer_unmap_range(
+ iree_hal_buffer_t* buffer, iree_device_size_t local_byte_offset,
+ iree_device_size_t local_byte_length, iree_hal_buffer_mapping_t* mapping) {
+ if (IREE_UNLIKELY(!buffer->allocated_buffer)) {
+ return iree_make_status(IREE_STATUS_FAILED_PRECONDITION,
+ "buffer does not have committed storage");
+ }
+ return _VTABLE_DISPATCH(buffer->allocated_buffer, unmap_range)(
+ buffer->allocated_buffer, local_byte_offset, local_byte_length, mapping);
+}
+
+static iree_status_t iree_hal_deferred_buffer_invalidate_range(
+ iree_hal_buffer_t* buffer, iree_device_size_t local_byte_offset,
+ iree_device_size_t local_byte_length) {
+ if (IREE_UNLIKELY(!buffer->allocated_buffer)) {
+ return iree_make_status(IREE_STATUS_FAILED_PRECONDITION,
+ "buffer does not have committed storage");
+ }
+ return _VTABLE_DISPATCH(buffer->allocated_buffer, invalidate_range)(
+ buffer->allocated_buffer, local_byte_offset, local_byte_length);
+}
+
+static iree_status_t iree_hal_deferred_buffer_flush_range(
+ iree_hal_buffer_t* buffer, iree_device_size_t local_byte_offset,
+ iree_device_size_t local_byte_length) {
+ if (IREE_UNLIKELY(!buffer->allocated_buffer)) {
+ return iree_make_status(IREE_STATUS_FAILED_PRECONDITION,
+ "buffer does not have committed storage");
+ }
+ return _VTABLE_DISPATCH(buffer->allocated_buffer, flush_range)(
+ buffer->allocated_buffer, local_byte_offset, local_byte_length);
+}
+
+static const iree_hal_buffer_vtable_t iree_hal_deferred_buffer_vtable = {
+ .recycle = iree_hal_buffer_recycle,
+ .destroy = iree_hal_deferred_buffer_destroy,
+ .map_range = iree_hal_deferred_buffer_map_range,
+ .unmap_range = iree_hal_deferred_buffer_unmap_range,
+ .invalidate_range = iree_hal_deferred_buffer_invalidate_range,
+ .flush_range = iree_hal_deferred_buffer_flush_range,
+};
+
+//===----------------------------------------------------------------------===//
// iree_hal_buffer_t
//===----------------------------------------------------------------------===//
@@ -513,7 +655,7 @@
// the super deep indirection that could arise.
iree_hal_buffer_t* allocated_buffer =
iree_hal_buffer_allocated_buffer(buffer);
- if (allocated_buffer != buffer) {
+ if (allocated_buffer && allocated_buffer != buffer) {
return iree_hal_buffer_subspan(allocated_buffer, byte_offset, byte_length,
out_buffer);
}
diff --git a/runtime/src/iree/hal/buffer.h b/runtime/src/iree/hal/buffer.h
index 586cf77..c832228 100644
--- a/runtime/src/iree/hal/buffer.h
+++ b/runtime/src/iree/hal/buffer.h
@@ -11,6 +11,7 @@
#include <stdint.h>
#include "iree/base/api.h"
+#include "iree/hal/queue.h"
#include "iree/hal/resource.h"
#ifdef __cplusplus
@@ -109,6 +110,16 @@
};
typedef uint32_t iree_hal_memory_type_t;
+// Parses a memory type bitfield from a string.
+// See iree_bitfield_parse for usage.
+IREE_API_EXPORT iree_status_t iree_hal_memory_type_parse(
+ iree_string_view_t value, iree_hal_memory_type_t* out_value);
+
+// Formats a memory type bitfield as a string.
+// See iree_bitfield_format for usage.
+IREE_API_EXPORT iree_string_view_t iree_hal_memory_type_format(
+ iree_hal_memory_type_t value, iree_bitfield_string_temp_t* out_temp);
+
// A bitfield specifying how memory will be accessed in a mapped memory region.
enum iree_hal_memory_access_bits_t {
// Memory is not mapped.
@@ -153,7 +164,15 @@
};
typedef uint16_t iree_hal_memory_access_t;
-typedef uint32_t iree_hal_buffer_compatibility_t;
+// Parses a memory access bitfield from a string.
+// See iree_bitfield_parse for usage.
+IREE_API_EXPORT iree_status_t iree_hal_memory_access_parse(
+ iree_string_view_t value, iree_hal_memory_access_t* out_value);
+
+// Formats a memory access bitfield as a string.
+// See iree_bitfield_format for usage.
+IREE_API_EXPORT iree_string_view_t iree_hal_memory_access_format(
+ iree_hal_memory_access_t value, iree_bitfield_string_temp_t* out_temp);
// Bitfield that defines how a buffer is intended to be used.
// Usage allows the driver to appropriately place the buffer for more
@@ -390,6 +409,16 @@
};
typedef uint32_t iree_hal_buffer_usage_t;
+// Parses a buffer usage bitfield from a string.
+// See iree_bitfield_parse for usage.
+IREE_API_EXPORT iree_status_t iree_hal_buffer_usage_parse(
+ iree_string_view_t value, iree_hal_buffer_usage_t* out_value);
+
+// Formats a buffer usage bitfield as a string.
+// See iree_bitfield_format for usage.
+IREE_API_EXPORT iree_string_view_t iree_hal_buffer_usage_format(
+ iree_hal_buffer_usage_t value, iree_bitfield_string_temp_t* out_temp);
+
// Buffer overlap testing results.
typedef enum iree_hal_buffer_overlap_e {
// No overlap between the two buffers.
@@ -400,6 +429,8 @@
IREE_HAL_BUFFER_OVERLAP_COMPLETE,
} iree_hal_buffer_overlap_t;
+typedef uint32_t iree_hal_buffer_compatibility_t;
+
// A bitfield specifying buffer transfer behavior.
enum iree_hal_transfer_buffer_flag_bits_t {
// TODO(benvanik): flags controlling blocking, flushing, invalidation, and
@@ -426,7 +457,98 @@
};
typedef uint32_t iree_hal_mapping_mode_t;
-// Implementation-specific mapping data.
+//===----------------------------------------------------------------------===//
+// iree_hal_buffer_params_t
+//===----------------------------------------------------------------------===//
+
+// Parameters defining how a buffer should be allocated.
+//
+// Designed to be zero-initialized: any field with a 0 value will be assigned
+// a default as indicated in the field description.
+//
+// For ergonomics when used from C++ w/o named initializers the first field is
+// the most commonly used so that it can be initialized by location:
+// some_fn(..., {IREE_HAL_BUFFER_USAGE_FOO}, ...)
+typedef struct iree_hal_buffer_params_t {
+ // Specifies the usage allowed by HAL APIs and aids in memory placement.
+ // Devices may have different memory types for different usage and require
+ // the intended usage to be declared upon allocation. It's always best to
+ // limit the allowed usage bits to precisely what the actual usage will be to
+ // avoid additional copies, synchronization, and expensive emulation.
+ //
+ // If 0 then the usage will default to all usage modes.
+ iree_hal_buffer_usage_t usage;
+
+ // Specifies the access allowed to the memory via the HAL APIs.
+ // For example, if the IREE_HAL_MEMORY_ACCESS_WRITE bit is not set then any
+ // API call that would write to the memory will fail (such as
+ // iree_hal_command_buffer_update_buffer). This does not limit any untrusted
+ // dispatch or external use of the buffer and should not be treated as a
+ // memory protection mechanism.
+ //
+ // If 0 then the access will be set as IREE_HAL_MEMORY_ACCESS_ALL.
+ iree_hal_memory_access_t access;
+
+ // Specifies the memory type properties used for selecting a memory space.
+ // This should often be IREE_HAL_MEMORY_TYPE_OPTIMAL to allow the allocator
+ // to place the allocation based on usage bits but can be specified if the
+ // exact memory type must be used for compatibility with external code.
+ //
+ // If 0 then the type will be set as IREE_HAL_MEMORY_TYPE_OPTIMAL.
+ iree_hal_memory_type_t type;
+
+ // Queue affinity bitmap indicating which queues may access this buffer.
+ // For NUMA devices this can be used to more tightly scope the allocation to
+ // particular device memory and provide better pool placement. When a device
+ // supports peering or replication the affinity bitmap will be used to choose
+ // which subdevices require configuration.
+ //
+ // If 0 then the buffer will be available on any queue as if
+ // IREE_HAL_QUEUE_AFFINITY_ANY was specified.
+ iree_hal_queue_affinity_t queue_affinity;
+
+ // Minimum alignment, in bytes, of the resulting allocation.
+ // The actual alignment may be any value greater-than-or-equal-to this value.
+ //
+ // If 0 then the alignment will be decided by the allocator based on optimal
+ // device parameters.
+ iree_device_size_t min_alignment;
+} iree_hal_buffer_params_t;
+
+// Canonicalizes |params| fields when zero initialization is used.
+static inline void iree_hal_buffer_params_canonicalize(
+ iree_hal_buffer_params_t* params) {
+ if (!params->usage) {
+ params->usage = IREE_HAL_BUFFER_USAGE_DEFAULT;
+ }
+ if (!params->access) {
+ params->access = IREE_HAL_MEMORY_ACCESS_ALL;
+ }
+ if (!params->type) {
+ params->type = IREE_HAL_MEMORY_TYPE_OPTIMAL;
+ }
+ if (!params->queue_affinity) {
+ params->queue_affinity = IREE_HAL_QUEUE_AFFINITY_ANY;
+ }
+}
+
+// Returns |params| with the given |usage| bits OR'ed in.
+static inline iree_hal_buffer_params_t iree_hal_buffer_params_with_usage(
+ const iree_hal_buffer_params_t params, iree_hal_buffer_usage_t usage) {
+ iree_hal_buffer_params_t result = params;
+ if (!result.usage) {
+ result.usage =
+ IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE | IREE_HAL_BUFFER_USAGE_TRANSFER;
+ }
+ result.usage |= usage;
+ return result;
+}
+
+//===----------------------------------------------------------------------===//
+// iree_hal_buffer_mapping_t
+//===----------------------------------------------------------------------===//
+
+// Implementation-specific private mapping data.
typedef struct iree_hal_buffer_mapping_impl_t {
// Byte offset within the buffer where the mapped data begins.
iree_device_size_t byte_offset;
@@ -464,35 +586,9 @@
iree_hal_buffer_mapping_impl_t impl;
} iree_hal_buffer_mapping_t;
-// Parses a memory type bitfield from a string.
-// See iree_bitfield_parse for usage.
-IREE_API_EXPORT iree_status_t iree_hal_memory_type_parse(
- iree_string_view_t value, iree_hal_memory_type_t* out_value);
-
-// Formats a memory type bitfield as a string.
-// See iree_bitfield_format for usage.
-IREE_API_EXPORT iree_string_view_t iree_hal_memory_type_format(
- iree_hal_memory_type_t value, iree_bitfield_string_temp_t* out_temp);
-
-// Parses a memory access bitfield from a string.
-// See iree_bitfield_parse for usage.
-IREE_API_EXPORT iree_status_t iree_hal_memory_access_parse(
- iree_string_view_t value, iree_hal_memory_access_t* out_value);
-
-// Formats a memory access bitfield as a string.
-// See iree_bitfield_format for usage.
-IREE_API_EXPORT iree_string_view_t iree_hal_memory_access_format(
- iree_hal_memory_access_t value, iree_bitfield_string_temp_t* out_temp);
-
-// Parses a buffer usage bitfield from a string.
-// See iree_bitfield_parse for usage.
-IREE_API_EXPORT iree_status_t iree_hal_buffer_usage_parse(
- iree_string_view_t value, iree_hal_buffer_usage_t* out_value);
-
-// Formats a buffer usage bitfield as a string.
-// See iree_bitfield_format for usage.
-IREE_API_EXPORT iree_string_view_t iree_hal_buffer_usage_format(
- iree_hal_buffer_usage_t value, iree_bitfield_string_temp_t* out_temp);
+//===----------------------------------------------------------------------===//
+// iree_hal_buffer_release_callback_t
+//===----------------------------------------------------------------------===//
typedef void(IREE_API_PTR* iree_hal_buffer_release_fn_t)(
void* user_data, struct iree_hal_buffer_t* buffer);
@@ -809,6 +905,35 @@
iree_allocator_t host_allocator, iree_hal_buffer_t** out_buffer);
//===----------------------------------------------------------------------===//
+// iree_hal_deferred_buffer_t
+//===----------------------------------------------------------------------===//
+
+// Creates a buffer with the given properties that has no backing storage.
+// The buffer can be passed around/retained/etc with just the reservation and
+// committed/decommitted on demand. All usage of the buffer beyond metadata
+// queries requires that it be committed.
+//
+// WARNING: commit/decommit are thread-compatible. Callers must ensure that no
+// threads try to use the buffer contents before a commit has completed and that
+// no threads still have access to the buffer contents prior to a decommit.
+IREE_API_EXPORT iree_status_t iree_hal_deferred_buffer_create_reserved(
+ iree_hal_allocator_t* device_allocator, iree_device_size_t allocation_size,
+ iree_device_size_t byte_offset, iree_device_size_t byte_length,
+ iree_hal_buffer_params_t params, iree_allocator_t host_allocator,
+ iree_hal_buffer_t** out_buffer);
+
+// Commits the backing storage of the |buffer| from its device allocator.
+// Ignored if the buffer is already committed.
+IREE_API_EXPORT iree_status_t
+iree_hal_deferred_buffer_commit(iree_hal_buffer_t* buffer);
+
+// Decommits the backing storage of the |buffer| and returns it to a
+// metadata-only state. No other threads must still have access to the buffer
+// contents.
+IREE_API_EXPORT iree_status_t
+iree_hal_deferred_buffer_decommit(iree_hal_buffer_t* buffer);
+
+//===----------------------------------------------------------------------===//
// iree_hal_heap_buffer_t
//===----------------------------------------------------------------------===//
