From 66f2e05f450e24a86e8dcb3f0bdd8e92ecdc112f Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident@gmail.com>
Date: Wed, 19 Feb 2025 06:37:55 -0800
Subject: [PATCH] Fix dealloc barrier.

---
 shortfin/src/shortfin/array/storage.cc  | 24 +++++++++++++-----------
 shortfin/src/shortfin/local/scheduler.h |  4 ++--
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/shortfin/src/shortfin/array/storage.cc b/shortfin/src/shortfin/array/storage.cc
index 8da5d986f..3339d69f1 100644
--- a/shortfin/src/shortfin/array/storage.cc
+++ b/shortfin/src/shortfin/array/storage.cc
@@ -70,27 +70,29 @@ storage storage::allocate_device(ScopedDevice &device,
       .semaphores = &timeline_sem,
       .payload_values = &signal_timepoint,
   };
-  SHORTFIN_SCHED_LOG(
-      "storage::allocate_device(device={}, affinity={:x}):[{}, Wait@{} ->"
-      "Signal:@{}]",
-      static_cast<void *>(device.raw_device()->hal_device()),
-      device.affinity().queue_affinity(), static_cast<void *>(timeline_sem),
-      current_timepoint, signal_timepoint);
   // Async allocate.
   SHORTFIN_THROW_IF_ERROR(iree_hal_device_queue_alloca(
       device.raw_device()->hal_device(), device.affinity().queue_affinity(),
       wait_semaphore_list, signal_semaphore_list,
       IREE_HAL_ALLOCATOR_POOL_DEFAULT, params, allocation_size,
       buffer.for_output()));
+  SHORTFIN_SCHED_LOG(
+      "storage::allocate_device(device={}, affinity={:x}):[{}, Wait@{}:"
+      "Signal:@{}] -> buffer={}",
+      static_cast<void *>(device.raw_device()->hal_device()),
+      device.affinity().queue_affinity(), static_cast<void *>(timeline_sem),
+      current_timepoint, signal_timepoint, static_cast<void *>(buffer.get()));
 
   // Device allocations are always async.
   // TODO: DO NOT SUBMIT: Enable async destruction.
-  // TimelineResourceDestructor dtor =
-  //     TimelineResource::CreateAsyncBufferDestructor(device, buffer);
+  TimelineResourceDestructor dtor =
+      TimelineResource::CreateAsyncBufferDestructor(device, buffer);
+  auto resource = device.fiber().NewTimelineResource(std::move(dtor));
+  resource->set_mutation_barrier(timeline_sem, signal_timepoint);
+  resource->use_barrier_insert(timeline_sem, signal_timepoint);
+  return storage(device, std::move(buffer), std::move(resource));
   // return storage(device, std::move(buffer),
-  //                device.fiber().NewTimelineResource(std::move(dtor)));
-  return storage(device, std::move(buffer),
-                 device.fiber().NewTimelineResource());
+  //                device.fiber().NewTimelineResource());
 }
 
 storage storage::allocate_host(ScopedDevice &device,
diff --git a/shortfin/src/shortfin/local/scheduler.h b/shortfin/src/shortfin/local/scheduler.h
index 0ed172ad7..eb7e7a3c9 100644
--- a/shortfin/src/shortfin/local/scheduler.h
+++ b/shortfin/src/shortfin/local/scheduler.h
@@ -288,8 +288,8 @@ class SHORTFIN_API Scheduler {
   TimelineResource::Ref NewTimelineResource(
       std::shared_ptr<Fiber> fiber,
       TimelineResourceDestructor destructor = nullptr) {
-    return TimelineResource::Ref(
-        new TimelineResource(std::move(fiber), semaphore_count_, destructor));
+    return TimelineResource::Ref(new TimelineResource(
+        std::move(fiber), semaphore_count_, std::move(destructor)));
   }
 
   // Creates a new fence with capacity for all semaphores that are extant at