[Level Zero] Adding iree_hal_device_replace_allocator API to level ze…

…ro backend (iree-org#36)
nithinsubbiah · Feb 27, 2023 · 6798ba8 · 6798ba8
1 parent 8e5f7cf
commit 6798ba8
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/experimental/level_zero/level_zero_device.c b/experimental/level_zero/level_zero_device.c
@@ -209,6 +209,14 @@ static iree_hal_allocator_t* iree_hal_level_zero_device_allocator(
   return device->device_allocator;
 }
 
+static void iree_hal_level_zero_replace_device_allocator(
+    iree_hal_device_t* base_device, iree_hal_allocator_t* new_allocator) {
+  iree_hal_level_zero_device_t* device = iree_hal_level_zero_device_cast(base_device);
+  iree_hal_allocator_retain(new_allocator);
+  iree_hal_allocator_release(device->device_allocator);
+  device->device_allocator = new_allocator;
+}
+
 static iree_status_t iree_hal_level_zero_device_query_i64(
     iree_hal_device_t* base_device, iree_string_view_t category,
     iree_string_view_t key, int64_t* out_value) {
@@ -389,6 +397,7 @@ static const iree_hal_device_vtable_t iree_hal_level_zero_device_vtable = {
     .id = iree_hal_level_zero_device_id,
     .host_allocator = iree_hal_level_zero_device_host_allocator,
     .device_allocator = iree_hal_level_zero_device_allocator,
+    .replace_device_allocator = iree_hal_level_zero_replace_device_allocator,
     .trim = iree_hal_level_zero_device_trim,
     .query_i64 = iree_hal_level_zero_device_query_i64,
     .create_command_buffer = iree_hal_level_zero_device_create_command_buffer,