From da61e202caaa8e14654b05b8b8105a0ae9f8cd73 Mon Sep 17 00:00:00 2001
From: Ben Vanik <benvanik@google.com>
Date: Sun, 18 Apr 2021 20:25:19 -0700
Subject: [PATCH] Adding the IREE_VM_BYTE_BUFFER_ORIGIN_* marker. This enables
 us to know the source of a byte buffer without trusting modules about its
 lifetime guarantees. A byte buffer that has a lifetime of the entire module
 can safely be aliased as by construction it will outlive any use from the
 module while allocated or external buffers may not have the same guarantees.

---
 iree/modules/hal/hal_module.c |  4 ++++
 iree/vm/builtin_types.h       | 15 +++++++++++++++
 iree/vm/bytecode_module.c     |  1 +
 3 files changed, 20 insertions(+)

diff --git a/iree/modules/hal/hal_module.c b/iree/modules/hal/hal_module.c
index a19091a616c3..e0e630d8c596 100644
--- a/iree/modules/hal/hal_module.c
+++ b/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;
diff --git a/iree/vm/builtin_types.h b/iree/vm/builtin_types.h
index c8144b5c6e0e..fe3c92c5eb13 100644
--- a/iree/vm/builtin_types.h
+++ b/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;
diff --git a/iree/vm/bytecode_module.c b/iree/vm/bytecode_module.c
index 9a4b8e3859df..05f4975539b9 100644
--- a/iree/vm/bytecode_module.c
+++ b/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));