From d43aa215ad1d91f7cfff3ea8de2161e29539f118 Mon Sep 17 00:00:00 2001
From: Andrew Woloszyn <andrew.woloszyn@gmail.com>
Date: Tue, 10 Dec 2024 15:10:36 -0500
Subject: [PATCH] Fix iree_hal_buffer_subspan after change in iree.

---
 shortfin/src/shortfin/array/storage.cc | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/shortfin/src/shortfin/array/storage.cc b/shortfin/src/shortfin/array/storage.cc
index ffbbd9ba2..1b59cc510 100644
--- a/shortfin/src/shortfin/array/storage.cc
+++ b/shortfin/src/shortfin/array/storage.cc
@@ -92,7 +92,9 @@ storage storage::subspan(iree_device_size_t byte_offset,
                          iree_device_size_t byte_length) {
   storage new_storage(device_, {}, timeline_resource_);
   SHORTFIN_THROW_IF_ERROR(iree_hal_buffer_subspan(
-      buffer_, byte_offset, byte_length, new_storage.buffer_.for_output()));
+      buffer_, byte_offset, byte_length,
+      iree_hal_device_host_allocator(device_.raw_device()->hal_device()),
+        new_storage.buffer_.for_output()));
   return new_storage;
 }