Merge pull request #5511 from google/benvanik-buffer-improvements

Small improvements to enable buffer aliasing/importing.

iree/hal/allocator.h

@@ -44,14 +44,20 @@ enum iree_hal_buffer_compatibility_e {
   // valid.
   IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE = 1u << 0,
 
+  // Indicates that the allocator could import external buffers of this type and
+  // usage natively. Imports may fail due to runtime conditions (out of handles,
+  // invalid pointer address spaces/page parameters, etc) but are otherwise
+  // valid.
+  IREE_HAL_BUFFER_COMPATIBILITY_IMPORTABLE = 1u << 1,
+
   // Indicates that the buffer can be used as a transfer source or target on the
   // a device queue (such as being the source or target of a DMA operation,
   // etc). If not set then the buffer may still be usable for
   // iree_hal_buffer_copy_data but not with queued operations.
-  IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_TRANSFER = 1u << 1,
+  IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_TRANSFER = 1u << 10,
 
   // Indicates that the buffer can be used as an input/output to a dispatch.
-  IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_DISPATCH = 1u << 2,
+  IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_DISPATCH = 1u << 11,
 };
 typedef uint32_t iree_hal_buffer_compatibility_t;
 
@@ -110,6 +116,13 @@ IREE_API_EXPORT iree_status_t IREE_API_CALL iree_hal_allocator_allocate_buffer(
     iree_hal_buffer_t** out_buffer);
 
 // Wraps an existing host allocation in a buffer.
+//
+// iree_hal_allocator_query_buffer_compatibility can be used to query whether a
+// buffer can be wrapped when using the given memory type and usage. A
+// compatibility result containing IREE_HAL_BUFFER_COMPATIBILITY_IMPORTABLE
+// means the wrap may succeed however if the pointer/page range is not in a
+// supported mode (no read access, etc) this call may still fail.
+//
 // |data_allocator| will be used to free the memory when the buffer is
 // destroyed. iree_allocator_null() can be passed to indicate the buffer does
 // not own the data.

iree/hal/allocator_heap.c

@@ -78,9 +78,11 @@ iree_hal_heap_allocator_query_buffer_compatibility(
   // based on what's both allowed and intended.
   intended_usage &= allowed_usage;
 
-  // All buffers can be allocated on the heap.
+  // All buffers can be allocated on the heap and all heap-accessible buffers
+  // can be imported.
   iree_hal_buffer_compatibility_t compatibility =
-      IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE;
+      IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE |
+      IREE_HAL_BUFFER_COMPATIBILITY_IMPORTABLE;
 
   // Buffers can only be used on the queue if they are device visible.
   // This is not a strict requirement of heap buffers but matches devices that

iree/modules/hal/hal_module.c

@@ -833,6 +833,10 @@ IREE_VM_ABI_EXPORT(iree_hal_module_executable_create, rrrCrD, r) {
   if (iree_status_is_ok(status)) {
     iree_hal_executable_spec_t spec;
     iree_hal_executable_spec_initialize(&spec);
+    spec.caching_mode |=
+        executable_data->origin == IREE_VM_BYTE_BUFFER_ORIGIN_MODULE
+            ? IREE_HAL_EXECUTABLE_CACHING_MODE_ALIAS_PROVIDED_DATA
+            : 0;
     spec.executable_format = executable_format_str;
     spec.executable_data = executable_data->data;
     spec.executable_layout_count = executable_layout_count;

iree/vm/builtin_types.h

@@ -22,11 +22,25 @@
 extern "C" {
 #endif  // __cplusplus
 
+// Describes where a byte buffer originates from and what guarantees can be made
+// about its lifetime and ownership.
+enum iree_vm_byte_buffer_origin_e {
+  // Buffer references memory in the module space (rodata or rwdata) that is
+  // guaranteed to be live for the lifetime of the module.
+  IREE_VM_BYTE_BUFFER_ORIGIN_MODULE = 0,
+  // Buffer references memory created by the guest module code. It has a
+  // lifetime less than that of the module but is always tracked with proper
+  // references (a handle existing to the memory implies it is valid).
+  IREE_VM_BYTE_BUFFER_ORIGIN_GUEST = 1,
+};
+typedef uint32_t iree_vm_byte_buffer_origin_t;
+
 // The built-in constant buffer type.
 // This simply points at a span of memory. The memory could be owned (in which
 // case a destroy function must be provided) or unowned (NULL destroy function).
 typedef struct {
   iree_vm_ref_object_t ref_object;
+  iree_vm_byte_buffer_origin_t origin;
   iree_const_byte_span_t data;
   iree_vm_ref_destroy_t destroy;
 } iree_vm_ro_byte_buffer_t;
@@ -44,6 +58,7 @@ static inline iree_string_view_t iree_vm_ro_byte_buffer_as_string(
 // case a destroy function must be provided) or unowned (NULL destroy function).
 typedef struct {
   iree_vm_ref_object_t ref_object;
+  iree_vm_byte_buffer_origin_t origin;
   iree_byte_span_t data;
   iree_vm_ref_destroy_t destroy;
 } iree_vm_rw_byte_buffer_t;

iree/vm/bytecode_module.c

@@ -577,6 +577,7 @@ static iree_status_t iree_vm_bytecode_module_alloc_state(
         iree_vm_RodataSegmentDef_vec_at(rodata_segments, i);
     iree_vm_ro_byte_buffer_t* ref = &state->rodata_ref_table[i];
     iree_atomic_ref_count_init(&ref->ref_object.counter);
+    ref->origin = IREE_VM_BYTE_BUFFER_ORIGIN_MODULE;
     ref->data.data = iree_vm_RodataSegmentDef_data(segment);
     ref->data.data_length =
         flatbuffers_uint8_vec_len(iree_vm_RodataSegmentDef_data(segment));