Replacing iree_hal_command_buffer_discard_buffer with advise_buffer.

This is similar to madvise and cudaMemAdvise but within a command
buffer. Stream operations in the compiler can lower to this to request
cache management or lifetime to allow ASAN and other implementation
features to know more about our aliased/slab buffer contents.
There is likely a need for an iree_hal_buffer_advise host operation.

Author: benvanik

experimental/webgpu/command_buffer.c

@@ -575,9 +575,10 @@ static iree_status_t iree_hal_webgpu_command_buffer_wait_events(
   return iree_ok_status();
 }
 
-static iree_status_t iree_hal_webgpu_command_buffer_discard_buffer(
+static iree_status_t iree_hal_webgpu_command_buffer_advise_buffer(
     iree_hal_command_buffer_t* base_command_buffer,
-    iree_hal_buffer_ref_t buffer_ref) {
+    iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags,
+    uint64_t arg0, uint64_t arg1) {
   // No-op: though maybe it'd be a useful addition to the spec as otherwise
   // false dependencies can creep in.
   return iree_ok_status();
@@ -1043,7 +1044,7 @@ const iree_hal_command_buffer_vtable_t iree_hal_webgpu_command_buffer_vtable = {
     .signal_event = iree_hal_webgpu_command_buffer_signal_event,
     .reset_event = iree_hal_webgpu_command_buffer_reset_event,
     .wait_events = iree_hal_webgpu_command_buffer_wait_events,
-    .discard_buffer = iree_hal_webgpu_command_buffer_discard_buffer,
+    .advise_buffer = iree_hal_webgpu_command_buffer_advise_buffer,
     .fill_buffer = iree_hal_webgpu_command_buffer_fill_buffer,
     .update_buffer = iree_hal_webgpu_command_buffer_update_buffer,
     .copy_buffer = iree_hal_webgpu_command_buffer_copy_buffer,

runtime/src/iree/hal/command_buffer.c

@@ -405,18 +405,19 @@ IREE_API_EXPORT iree_status_t iree_hal_command_buffer_wait_events(
   return status;
 }
 
-IREE_API_EXPORT iree_status_t iree_hal_command_buffer_discard_buffer(
-    iree_hal_command_buffer_t* command_buffer,
-    iree_hal_buffer_ref_t buffer_ref) {
+IREE_API_EXPORT iree_status_t iree_hal_command_buffer_advise_buffer(
+    iree_hal_command_buffer_t* command_buffer, iree_hal_buffer_ref_t buffer_ref,
+    iree_hal_memory_advise_flags_t flags, uint64_t arg0, uint64_t arg1) {
   IREE_ASSERT_ARGUMENT(command_buffer);
   IREE_TRACE_ZONE_BEGIN(z0);
   IF_VALIDATING(command_buffer, {
     IREE_RETURN_AND_END_ZONE_IF_ERROR(
-        z0, iree_hal_command_buffer_discard_buffer_validation(
-                command_buffer, VALIDATION_STATE(command_buffer), buffer_ref));
+        z0, iree_hal_command_buffer_advise_buffer_validation(
+                command_buffer, VALIDATION_STATE(command_buffer), buffer_ref,
+                flags, arg0, arg1));
   });
-  iree_status_t status = _VTABLE_DISPATCH(command_buffer, discard_buffer)(
-      command_buffer, buffer_ref);
+  iree_status_t status = _VTABLE_DISPATCH(command_buffer, advise_buffer)(
+      command_buffer, buffer_ref, flags, arg0, arg1);
   IREE_TRACE_ZONE_END(z0);
   return status;
 }

runtime/src/iree/hal/command_buffer.h

@@ -214,6 +214,16 @@ typedef struct iree_hal_buffer_barrier_t {
   iree_hal_buffer_ref_t buffer_ref;
 } iree_hal_buffer_barrier_t;
 
+// Bitfield indicating advice for implementations managing a buffer.
+typedef uint64_t iree_hal_memory_advise_flags_t;
+enum iree_hal_memory_advise_flag_bits_t {
+  IREE_HAL_MEMORY_ADVISE_FLAG_NONE = 0,
+  // TODO(benvanik): cache control operations (invalidate/flush). arg0/arg1
+  // could source/target queue affinities.
+  // TODO(benvanik): prefetch and access type hints.
+  // TODO(benvanik): ASAN hints (protect/unprotect).
+};
+
 // Bitfield specifying flags controlling a fill operation.
 typedef uint64_t iree_hal_fill_flags_t;
 enum iree_hal_fill_flag_bits_t {
@@ -687,13 +697,12 @@ IREE_API_EXPORT iree_status_t iree_hal_command_buffer_wait_events(
     iree_host_size_t buffer_barrier_count,
     const iree_hal_buffer_barrier_t* buffer_barriers);
 
-// Hints to the device queue that the given buffer will not be used again.
-// After encoding a discard the buffer contents will be considered undefined.
-// This is because the discard may be used to elide write backs to host memory
-// or aggressively reuse the allocation for other purposes.
-IREE_API_EXPORT iree_status_t iree_hal_command_buffer_discard_buffer(
-    iree_hal_command_buffer_t* command_buffer,
-    iree_hal_buffer_ref_t buffer_ref);
+// Advises the device about the usage of the given buffer.
+// The device may use this information to perform cache management or ignore it
+// entirely.
+IREE_API_EXPORT iree_status_t iree_hal_command_buffer_advise_buffer(
+    iree_hal_command_buffer_t* command_buffer, iree_hal_buffer_ref_t buffer_ref,
+    iree_hal_memory_advise_flags_t flags, uint64_t arg0, uint64_t arg1);
 
 // Fills the target buffer with the given repeating value.
 // Expects that |pattern_length| is one of 1, 2, or 4 and that the offset and
@@ -892,9 +901,10 @@ typedef struct iree_hal_command_buffer_vtable_t {
       iree_host_size_t buffer_barrier_count,
       const iree_hal_buffer_barrier_t* buffer_barriers);
 
-  iree_status_t(IREE_API_PTR* discard_buffer)(
+  iree_status_t(IREE_API_PTR* advise_buffer)(
       iree_hal_command_buffer_t* command_buffer,
-      iree_hal_buffer_ref_t buffer_ref);
+      iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags,
+      uint64_t arg0, uint64_t arg1);
 
   iree_status_t(IREE_API_PTR* fill_buffer)(
       iree_hal_command_buffer_t* command_buffer,

runtime/src/iree/hal/command_buffer_validation.c

@@ -331,10 +331,11 @@ iree_status_t iree_hal_command_buffer_wait_events_validation(
   return iree_ok_status();
 }
 
-iree_status_t iree_hal_command_buffer_discard_buffer_validation(
+iree_status_t iree_hal_command_buffer_advise_buffer_validation(
     iree_hal_command_buffer_t* command_buffer,
     iree_hal_command_buffer_validation_state_t* validation_state,
-    iree_hal_buffer_ref_t buffer_ref) {
+    iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags,
+    uint64_t arg0, uint64_t arg1) {
   IREE_RETURN_IF_ERROR(iree_hal_command_buffer_validate_categories(
       command_buffer, validation_state, IREE_HAL_COMMAND_CATEGORY_TRANSFER));
 

runtime/src/iree/hal/command_buffer_validation.h

@@ -99,10 +99,11 @@ iree_status_t iree_hal_command_buffer_wait_events_validation(
     iree_host_size_t buffer_barrier_count,
     const iree_hal_buffer_barrier_t* buffer_barriers);
 
-iree_status_t iree_hal_command_buffer_discard_buffer_validation(
+iree_status_t iree_hal_command_buffer_advise_buffer_validation(
     iree_hal_command_buffer_t* command_buffer,
     iree_hal_command_buffer_validation_state_t* validation_state,
-    iree_hal_buffer_ref_t buffer_ref);
+    iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags,
+    uint64_t arg0, uint64_t arg1);
 
 iree_status_t iree_hal_command_buffer_fill_buffer_validation(
     iree_hal_command_buffer_t* command_buffer,

runtime/src/iree/hal/drivers/cuda/graph_command_buffer.c

@@ -477,9 +477,10 @@ static iree_status_t iree_hal_cuda_graph_command_buffer_wait_events(
   return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "event not yet supported");
 }
 
-static iree_status_t iree_hal_cuda_graph_command_buffer_discard_buffer(
+static iree_status_t iree_hal_cuda_graph_command_buffer_advise_buffer(
     iree_hal_command_buffer_t* base_command_buffer,
-    iree_hal_buffer_ref_t buffer_ref) {
+    iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags,
+    uint64_t arg0, uint64_t arg1) {
   // We could mark the memory as invalidated so that if this is a managed buffer
   // CUDA does not try to copy it back to the host.
   return iree_ok_status();
@@ -849,7 +850,7 @@ static const iree_hal_command_buffer_vtable_t
         .signal_event = iree_hal_cuda_graph_command_buffer_signal_event,
         .reset_event = iree_hal_cuda_graph_command_buffer_reset_event,
         .wait_events = iree_hal_cuda_graph_command_buffer_wait_events,
-        .discard_buffer = iree_hal_cuda_graph_command_buffer_discard_buffer,
+        .advise_buffer = iree_hal_cuda_graph_command_buffer_advise_buffer,
         .fill_buffer = iree_hal_cuda_graph_command_buffer_fill_buffer,
         .update_buffer = iree_hal_cuda_graph_command_buffer_update_buffer,
         .copy_buffer = iree_hal_cuda_graph_command_buffer_copy_buffer,

runtime/src/iree/hal/drivers/cuda/stream_command_buffer.c

@@ -308,9 +308,10 @@ static iree_status_t iree_hal_cuda_stream_command_buffer_wait_events(
   return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "event not yet supported");
 }
 
-static iree_status_t iree_hal_cuda_stream_command_buffer_discard_buffer(
+static iree_status_t iree_hal_cuda_stream_command_buffer_advise_buffer(
     iree_hal_command_buffer_t* base_command_buffer,
-    iree_hal_buffer_ref_t buffer_ref) {
+    iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags,
+    uint64_t arg0, uint64_t arg1) {
   // We could mark the memory as invalidated so that if managed CUDA does not
   // try to copy it back to the host.
   return iree_ok_status();
@@ -601,7 +602,7 @@ static const iree_hal_command_buffer_vtable_t
         .signal_event = iree_hal_cuda_stream_command_buffer_signal_event,
         .reset_event = iree_hal_cuda_stream_command_buffer_reset_event,
         .wait_events = iree_hal_cuda_stream_command_buffer_wait_events,
-        .discard_buffer = iree_hal_cuda_stream_command_buffer_discard_buffer,
+        .advise_buffer = iree_hal_cuda_stream_command_buffer_advise_buffer,
         .fill_buffer = iree_hal_cuda_stream_command_buffer_fill_buffer,
         .update_buffer = iree_hal_cuda_stream_command_buffer_update_buffer,
         .copy_buffer = iree_hal_cuda_stream_command_buffer_copy_buffer,

runtime/src/iree/hal/drivers/hip/graph_command_buffer.c

@@ -486,9 +486,10 @@ static iree_status_t iree_hal_hip_graph_command_buffer_wait_events(
   return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "event not yet supported");
 }
 
-static iree_status_t iree_hal_hip_graph_command_buffer_discard_buffer(
+static iree_status_t iree_hal_hip_graph_command_buffer_advise_buffer(
     iree_hal_command_buffer_t* base_command_buffer,
-    iree_hal_buffer_ref_t buffer_ref) {
+    iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags,
+    uint64_t arg0, uint64_t arg1) {
   // We could mark the memory as invalidated so that if this is a managed buffer
   // HIP does not try to copy it back to the host.
   return iree_ok_status();
@@ -858,7 +859,7 @@ static const iree_hal_command_buffer_vtable_t
         .signal_event = iree_hal_hip_graph_command_buffer_signal_event,
         .reset_event = iree_hal_hip_graph_command_buffer_reset_event,
         .wait_events = iree_hal_hip_graph_command_buffer_wait_events,
-        .discard_buffer = iree_hal_hip_graph_command_buffer_discard_buffer,
+        .advise_buffer = iree_hal_hip_graph_command_buffer_advise_buffer,
         .fill_buffer = iree_hal_hip_graph_command_buffer_fill_buffer,
         .update_buffer = iree_hal_hip_graph_command_buffer_update_buffer,
         .copy_buffer = iree_hal_hip_graph_command_buffer_copy_buffer,

runtime/src/iree/hal/drivers/hip/stream_command_buffer.c

@@ -299,9 +299,10 @@ static iree_status_t iree_hal_hip_stream_command_buffer_wait_events(
   return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "event not yet supported");
 }
 
-static iree_status_t iree_hal_hip_stream_command_buffer_discard_buffer(
+static iree_status_t iree_hal_hip_stream_command_buffer_advise_buffer(
     iree_hal_command_buffer_t* base_command_buffer,
-    iree_hal_buffer_ref_t buffer_ref) {
+    iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags,
+    uint64_t arg0, uint64_t arg1) {
   // We could mark the memory as invalidated so that if managed HIP does not
   // try to copy it back to the host.
   return iree_ok_status();
@@ -590,7 +591,7 @@ static const iree_hal_command_buffer_vtable_t
         .signal_event = iree_hal_hip_stream_command_buffer_signal_event,
         .reset_event = iree_hal_hip_stream_command_buffer_reset_event,
         .wait_events = iree_hal_hip_stream_command_buffer_wait_events,
-        .discard_buffer = iree_hal_hip_stream_command_buffer_discard_buffer,
+        .advise_buffer = iree_hal_hip_stream_command_buffer_advise_buffer,
         .fill_buffer = iree_hal_hip_stream_command_buffer_fill_buffer,
         .update_buffer = iree_hal_hip_stream_command_buffer_update_buffer,
         .copy_buffer = iree_hal_hip_stream_command_buffer_copy_buffer,

runtime/src/iree/hal/drivers/local_task/task_command_buffer.c

@@ -455,12 +455,13 @@ static iree_status_t iree_hal_task_command_buffer_wait_events(
 }
 
 //===----------------------------------------------------------------------===//
-// iree_hal_command_buffer_discard_buffer
+// iree_hal_command_buffer_advise_buffer
 //===----------------------------------------------------------------------===//
 
-static iree_status_t iree_hal_task_command_buffer_discard_buffer(
+static iree_status_t iree_hal_task_command_buffer_advise_buffer(
     iree_hal_command_buffer_t* base_command_buffer,
-    iree_hal_buffer_ref_t buffer_ref) {
+    iree_hal_buffer_ref_t buffer_ref, iree_hal_memory_advise_flags_t flags,
+    uint64_t arg0, uint64_t arg1) {
   return iree_ok_status();
 }
 
@@ -948,7 +949,7 @@ static const iree_hal_command_buffer_vtable_t
         .signal_event = iree_hal_task_command_buffer_signal_event,
         .reset_event = iree_hal_task_command_buffer_reset_event,
         .wait_events = iree_hal_task_command_buffer_wait_events,
-        .discard_buffer = iree_hal_task_command_buffer_discard_buffer,
+        .advise_buffer = iree_hal_task_command_buffer_advise_buffer,
         .fill_buffer = iree_hal_task_command_buffer_fill_buffer,
         .update_buffer = iree_hal_task_command_buffer_update_buffer,
         .copy_buffer = iree_hal_task_command_buffer_copy_buffer,