From 6798ba8c69a885c44f3b42d56b29bf6b894ceb25 Mon Sep 17 00:00:00 2001
From: Nithin Meganathan <18070964+nithinsubbiah@users.noreply.github.com>
Date: Thu, 12 Jan 2023 17:28:21 -0800
Subject: [PATCH] [Level Zero] Adding iree_hal_device_replace_allocator API to
 level zero backend (#36)

---
 experimental/level_zero/level_zero_device.c | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/experimental/level_zero/level_zero_device.c b/experimental/level_zero/level_zero_device.c
index 958c1fca78c0..e28c9e6885bf 100644
--- 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,