Fix xrt driver

nod-ai · Dec 11, 2024 · bff07d5 · bff07d5
1 parent a9905f6
commit bff07d5
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 6 deletions.
diff --git a/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.cc b/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.cc
@@ -197,7 +197,8 @@ static iree_status_t iree_hal_xrt_allocator_allocate_buffer(
         xrt_buffer.release(), placement, compat_params.type,
         compat_params.access, compat_params.usage, allocation_size,
         /*byte_offset=*/0, /*byte_length=*/allocation_size,
-        iree_hal_buffer_release_callback_null(), &buffer);
+        iree_hal_buffer_release_callback_null(), allocator->host_allocator,
+        &buffer);
   }
 
   if (iree_status_is_ok(status)) {

diff --git a/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.cc b/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.cc
@@ -39,13 +39,10 @@ iree_status_t iree_hal_xrt_buffer_wrap(
     iree_hal_buffer_usage_t allowed_usage, iree_device_size_t allocation_size,
     iree_device_size_t byte_offset, iree_device_size_t byte_length,
     iree_hal_buffer_release_callback_t release_callback,
-    iree_hal_buffer_t** out_buffer) {
-  IREE_ASSERT_ARGUMENT(allocator);
+    iree_allocator_t host_allocator, iree_hal_buffer_t** out_buffer) {
   IREE_ASSERT_ARGUMENT(out_buffer);
   IREE_TRACE_ZONE_BEGIN(z0);
 
-  iree_allocator_t host_allocator =
-      iree_hal_allocator_host_allocator(allocator);
   iree_hal_xrt_buffer_t* buffer = nullptr;
   IREE_RETURN_AND_END_ZONE_IF_ERROR(
       z0,

diff --git a/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.h b/runtime/src/iree-amd-aie/driver/xrt/xrt_buffer.h
@@ -23,7 +23,7 @@ iree_status_t iree_hal_xrt_buffer_wrap(
     iree_hal_buffer_usage_t allowed_usage, iree_device_size_t allocation_size,
     iree_device_size_t byte_offset, iree_device_size_t byte_length,
     iree_hal_buffer_release_callback_t release_callback,
-    iree_hal_buffer_t** out_buffer);
+    iree_allocator_t host_allocator, iree_hal_buffer_t** out_buffer);
 
 // Returns the underlying XRT buffer handle for the given |buffer|.
 xrt::bo* iree_hal_xrt_buffer_handle(const iree_hal_buffer_t* buffer);