Plumbing arena block pools through CUDA/Vulkan.

experimental/rocm/rocm_device.c

diff --git a/experimental/rocm/rocm_device.c b/experimental/rocm/rocm_device.c
index 026650f..6666bd4 100644
--- a/experimental/rocm/rocm_device.c
+++ b/experimental/rocm/rocm_device.c
@@ -172,7 +172,7 @@
 
 static iree_status_t iree_hal_rocm_device_trim(iree_hal_device_t* base_device) {
   iree_hal_rocm_device_t* device = iree_hal_rocm_device_cast(base_device);
-  // TODO(benvanik): trim of ROCM resources, whenever we care.
+  iree_arena_block_pool_trim(&device->block_pool);
   return iree_hal_allocator_trim(device->device_allocator);
 }
 

iree/hal/cuda/cuda_device.c

diff --git a/iree/hal/cuda/cuda_device.c b/iree/hal/cuda/cuda_device.c
index eab84ee..677fb5b 100644
--- a/iree/hal/cuda/cuda_device.c
+++ b/iree/hal/cuda/cuda_device.c
@@ -202,7 +202,7 @@
 
 static iree_status_t iree_hal_cuda_device_trim(iree_hal_device_t* base_device) {
   iree_hal_cuda_device_t* device = iree_hal_cuda_device_cast(base_device);
-  // TODO(benvanik): trim of CUDA resources, whenever we care.
+  iree_arena_block_pool_trim(&device->block_pool);
   return iree_hal_allocator_trim(device->device_allocator);
 }
 
@@ -237,7 +237,7 @@
     case IREE_HAL_CUDA_COMMAND_BUFFER_MODE_GRAPH:
       return iree_hal_cuda_graph_command_buffer_create(
           base_device, &device->context_wrapper, mode, command_categories,
-          queue_affinity, out_command_buffer);
+          queue_affinity, &device->block_pool, out_command_buffer);
     case IREE_HAL_CUDA_COMMAND_BUFFER_MODE_STREAM:
       return iree_hal_deferred_command_buffer_create(
           base_device, mode, command_categories, &device->block_pool,
@@ -324,8 +324,8 @@
       }
     }
   }
-  // TODO(thomasraoux): Conservatively syncronize after every submit until we
-  // support semaphores.
+  // TODO(thomasraoux): implement semaphores - for now this conservatively
+  // synchronizes after every submit.
   CUDA_RETURN_IF_ERROR(device->context_wrapper.syms,
                        cuStreamSynchronize(device->stream),
                        "cuStreamSynchronize");

iree/hal/cuda/graph_command_buffer.c

diff --git a/iree/hal/cuda/graph_command_buffer.c b/iree/hal/cuda/graph_command_buffer.c
index d55ac91..466da1c 100644
--- a/iree/hal/cuda/graph_command_buffer.c
+++ b/iree/hal/cuda/graph_command_buffer.c
@@ -28,6 +28,7 @@
 typedef struct iree_hal_cuda_graph_command_buffer_t {
   iree_hal_command_buffer_t base;
   iree_hal_cuda_context_wrapper_t* context;
+  iree_arena_block_pool_t* block_pool;
 
   CUgraph graph;
   CUgraphExec exec;
@@ -54,8 +55,10 @@
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
     iree_hal_queue_affinity_t queue_affinity,
+    iree_arena_block_pool_t* block_pool,
     iree_hal_command_buffer_t** out_command_buffer) {
   IREE_ASSERT_ARGUMENT(context);
+  IREE_ASSERT_ARGUMENT(block_pool);
   IREE_ASSERT_ARGUMENT(out_command_buffer);
   IREE_TRACE_ZONE_BEGIN(z0);
 
@@ -73,6 +76,7 @@
         device, mode, command_categories, queue_affinity,
         &iree_hal_cuda_graph_command_buffer_vtable, &command_buffer->base);
     command_buffer->context = context;
+    command_buffer->block_pool = block_pool;
     command_buffer->graph = graph;
     command_buffer->exec = NULL;
     command_buffer->last_node = NULL;

iree/hal/cuda/graph_command_buffer.h

diff --git a/iree/hal/cuda/graph_command_buffer.h b/iree/hal/cuda/graph_command_buffer.h
index 7ad5251..8ef4fda 100644
--- a/iree/hal/cuda/graph_command_buffer.h
+++ b/iree/hal/cuda/graph_command_buffer.h
@@ -17,12 +17,18 @@
 extern "C" {
 #endif  // __cplusplus
 
-// Creates a cuda graph.
+typedef struct iree_arena_block_pool_t iree_arena_block_pool_t;
+
+// Creates a command buffer that records into a CUDA graph.
+//
+// NOTE: the |block_pool| must remain live for the lifetime of the command
+// buffers that use it.
 iree_status_t iree_hal_cuda_graph_command_buffer_create(
     iree_hal_device_t* device, iree_hal_cuda_context_wrapper_t* context,
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
     iree_hal_queue_affinity_t queue_affinity,
+    iree_arena_block_pool_t* block_pool,
     iree_hal_command_buffer_t** out_command_buffer);
 
 // Returns true if |command_buffer| is a CUDA graph-based command buffer.

iree/hal/vulkan/BUILD

diff --git a/iree/hal/vulkan/BUILD b/iree/hal/vulkan/BUILD
index 2184d1a..5e02eab 100644
--- a/iree/hal/vulkan/BUILD
+++ b/iree/hal/vulkan/BUILD
@@ -91,6 +91,7 @@
         "//iree/base:logging",
         "//iree/base:tracing",
         "//iree/base/internal",
+        "//iree/base/internal:arena",
         "//iree/base/internal:synchronization",
         "//iree/base/internal/flatcc:parsing",
         "//iree/hal",

iree/hal/vulkan/CMakeLists.txt

diff --git a/iree/hal/vulkan/CMakeLists.txt b/iree/hal/vulkan/CMakeLists.txt
index 9d27e03..c78a5bd 100644
--- a/iree/hal/vulkan/CMakeLists.txt
+++ b/iree/hal/vulkan/CMakeLists.txt
@@ -79,6 +79,7 @@
     iree::base::cc
     iree::base::core_headers
     iree::base::internal
+    iree::base::internal::arena
     iree::base::internal::flatcc::parsing
     iree::base::internal::synchronization
     iree::base::logging

iree/hal/vulkan/direct_command_buffer.cc

diff --git a/iree/hal/vulkan/direct_command_buffer.cc b/iree/hal/vulkan/direct_command_buffer.cc
index 9e6d799..b8ce770 100644
--- a/iree/hal/vulkan/direct_command_buffer.cc
+++ b/iree/hal/vulkan/direct_command_buffer.cc
@@ -34,6 +34,7 @@
   iree_hal_command_buffer_t base;
   VkDeviceHandle* logical_device;
   iree_hal_vulkan_tracing_context_t* tracing_context;
+  iree_arena_block_pool_t* block_pool;
 
   VkCommandPoolHandle* command_pool;
   VkCommandBuffer handle;
@@ -81,10 +82,12 @@
     iree_hal_vulkan_tracing_context_t* tracing_context,
     iree::hal::vulkan::DescriptorPoolCache* descriptor_pool_cache,
     iree::hal::vulkan::BuiltinExecutables* builtin_executables,
+    iree_arena_block_pool_t* block_pool,
     iree_hal_command_buffer_t** out_command_buffer) {
   IREE_ASSERT_ARGUMENT(logical_device);
   IREE_ASSERT_ARGUMENT(command_pool);
   IREE_ASSERT_ARGUMENT(descriptor_pool_cache);
+  IREE_ASSERT_ARGUMENT(block_pool);
   IREE_ASSERT_ARGUMENT(out_command_buffer);
   IREE_TRACE_ZONE_BEGIN(z0);
 
@@ -109,6 +112,7 @@
         &iree_hal_vulkan_direct_command_buffer_vtable, &command_buffer->base);
     command_buffer->logical_device = logical_device;
     command_buffer->tracing_context = tracing_context;
+    command_buffer->block_pool = block_pool;
     command_buffer->command_pool = command_pool;
     command_buffer->handle = handle;
     command_buffer->syms = logical_device->syms().get();

iree/hal/vulkan/direct_command_buffer.h

diff --git a/iree/hal/vulkan/direct_command_buffer.h b/iree/hal/vulkan/direct_command_buffer.h
index 071d369..57c15ad 100644
--- a/iree/hal/vulkan/direct_command_buffer.h
+++ b/iree/hal/vulkan/direct_command_buffer.h
@@ -18,7 +18,12 @@
 extern "C" {
 #endif  // __cplusplus
 
+typedef struct iree_arena_block_pool_t iree_arena_block_pool_t;
+
 // Creates a command buffer that directly records into a VkCommandBuffer.
+//
+// NOTE: the |block_pool| must remain live for the lifetime of the command
+// buffers that use it.
 iree_status_t iree_hal_vulkan_direct_command_buffer_allocate(
     iree_hal_device_t* device,
     iree::hal::vulkan::VkDeviceHandle* logical_device,
@@ -29,6 +34,7 @@
     iree_hal_vulkan_tracing_context_t* tracing_context,
     iree::hal::vulkan::DescriptorPoolCache* descriptor_pool_cache,
     iree::hal::vulkan::BuiltinExecutables* builtin_executables,
+    iree_arena_block_pool_t* block_pool,
     iree_hal_command_buffer_t** out_command_buffer);
 
 // Returns the native Vulkan VkCommandBuffer handle.

iree/hal/vulkan/vulkan_device.cc

diff --git a/iree/hal/vulkan/vulkan_device.cc b/iree/hal/vulkan/vulkan_device.cc
index 73e4a82..966e4e3 100644
--- a/iree/hal/vulkan/vulkan_device.cc
+++ b/iree/hal/vulkan/vulkan_device.cc
@@ -11,6 +11,7 @@
 #include <cstring>
 #include <vector>
 
+#include "iree/base/internal/arena.h"
 #include "iree/base/internal/math.h"
 #include "iree/base/tracing.h"
 #include "iree/hal/utils/buffer_transfer.h"
@@ -362,6 +363,10 @@
   VkCommandPoolHandle* dispatch_command_pool;
   VkCommandPoolHandle* transfer_command_pool;
 
+  // Block pool used for command buffers with a larger block size (as command
+  // buffers can contain inlined data uploads).
+  iree_arena_block_pool_t block_pool;
+
   // Used only for emulated timeline semaphores.
   TimePointSemaphorePool* semaphore_pool;
   TimePointFencePool* fence_pool;
@@ -563,6 +568,9 @@
   device->logical_device = logical_device;
   device->logical_device->AddReference();
 
+  iree_arena_block_pool_initialize(32 * 1024, host_allocator,
+                                   &device->block_pool);
+
   // Point the queue storage into the new device allocation. The queues
   // themselves are populated
   device->queues = (CommandQueue**)buffer_ptr;
@@ -667,6 +675,9 @@
   // There should be no more buffers live that use the allocator.
   iree_hal_allocator_release(device->device_allocator);
 
+  // All arena blocks should have been returned.
+  iree_arena_block_pool_deinitialize(&device->block_pool);
+
   // Finally, destroy the device.
   device->logical_device->ReleaseReference();
   iree_hal_driver_release(device->driver);
@@ -919,7 +930,7 @@
 static iree_status_t iree_hal_vulkan_device_trim(
     iree_hal_device_t* base_device) {
   iree_hal_vulkan_device_t* device = iree_hal_vulkan_device_cast(base_device);
-  // TODO(benvanik): trim of vulkan resources, whenever we care.
+  iree_arena_block_pool_trim(&device->block_pool);
   return iree_hal_allocator_trim(device->device_allocator);
 }
 
@@ -1007,7 +1018,7 @@
       base_device, device->logical_device, command_pool, mode,
       command_categories, queue_affinity, queue->tracing_context(),
       device->descriptor_pool_cache, device->builtin_executables,
-      out_command_buffer);
+      &device->block_pool, out_command_buffer);
 }
 
 static iree_status_t iree_hal_vulkan_device_create_descriptor_set(