Cleaning up iree_hal_module_debug_sink_t destroy. (#20841)

It was not properly added to the inline HAL, did not handle the status,
and was misnamed for what it was doing (destroy = delete). It's still
not great that this grew this method but it was confusing as it was. If
debug sinks need lifetime management they should be promoted to a full
ref counted type instead of an on-stack by-value struct as was the
original intent.

Author: benvanik

runtime/bindings/python/hal.cc

@@ -1137,10 +1137,10 @@ iree_hal_module_debug_sink_t HalModuleDebugSink::AsIreeHalModuleDebugSink()
     const {
   iree_hal_module_debug_sink_t res;
   memset(&res, 0, sizeof(res));
+  res.release.fn = HalModuleDebugSink::ReleaseCallback;
+  res.release.user_data = const_cast<HalModuleDebugSink*>(this);
   res.buffer_view_trace.fn = HalModuleDebugSink::IreeHalModuleBufferViewTrace;
   res.buffer_view_trace.user_data = const_cast<HalModuleDebugSink*>(this);
-  res.destroy.fn = HalModuleDebugSink::DestroyCallback;
-  res.destroy.user_data = const_cast<HalModuleDebugSink*>(this);
   return res;
 }
 
@@ -1161,11 +1161,10 @@ static std::vector<HalBufferView> CreateHalBufferViewVector(
   return res;
 }
 
-iree_status_t HalModuleDebugSink::DestroyCallback(void* user_data) {
+void HalModuleDebugSink::ReleaseCallback(void* user_data) {
   HalModuleDebugSink* debug_sink =
       reinterpret_cast<HalModuleDebugSink*>(user_data);
   debug_sink->dec_ref();
-  return iree_ok_status();
 }
 
 iree_status_t HalModuleDebugSink::IreeHalModuleBufferViewTrace(

runtime/bindings/python/hal.h

@@ -334,7 +334,7 @@ class HalModuleDebugSink : public py::intrusive_base {
  private:
   HalModuleBufferViewTraceCallback buffer_view_trace_callback_;
 
-  static iree_status_t DestroyCallback(void* user_data);
+  static void ReleaseCallback(void* user_data);
 
   static iree_status_t IreeHalModuleBufferViewTrace(
       void* user_data, iree_string_view_t key,

runtime/src/iree/modules/check/check_test.cc

@@ -233,27 +233,6 @@ iree_vm_instance_t* CheckTest::instance_ = nullptr;
 iree_vm_module_t* CheckTest::check_module_ = nullptr;
 iree_vm_module_t* CheckTest::hal_module_ = nullptr;
 
-TEST_F(CheckTest, HalModuleDebugSinkDestroyCallbackIsCalled) {
-  struct UserData {
-    bool is_callback_called = false;
-  };
-
-  iree_hal_module_debug_sink_t sink = {};
-  sink.destroy.fn = [](void* user_data) {
-    reinterpret_cast<UserData*>(user_data)->is_callback_called = true;
-    return iree_ok_status();
-  };
-  UserData user_data;
-  sink.destroy.user_data = &user_data;
-  iree_vm_module_t* hal_module;
-  IREE_ASSERT_OK(iree_hal_module_create(
-      instance(), /*device_count=*/1, &device(), IREE_HAL_MODULE_FLAG_NONE,
-      sink, iree_allocator_system(), &hal_module));
-  IREE_ASSERT_FALSE(user_data.is_callback_called);
-  iree_vm_module_release(hal_module);
-  IREE_ASSERT_TRUE(user_data.is_callback_called);
-}
-
 TEST_F(CheckTest, ExpectTrueSuccess) {
   IREE_ASSERT_OK(InvokeValue("expect_true", {iree_vm_value_make_i32(1)}));
 }

runtime/src/iree/modules/hal/debugging.h

@@ -19,7 +19,14 @@ extern "C" {
 // Debug Sink
 //===----------------------------------------------------------------------===//
 
-// Receives a
+typedef void(IREE_API_PTR* iree_hal_module_debug_sink_release_callback_fn_t)(
+    void* user_data);
+
+typedef struct iree_hal_module_debug_sink_release_callback_t {
+  iree_hal_module_debug_sink_release_callback_fn_t fn;
+  void* user_data;
+} iree_hal_module_debug_sink_release_callback_t;
+
 typedef iree_status_t(
     IREE_API_PTR* iree_hal_module_buffer_view_trace_callback_fn_t)(
     void* user_data, iree_string_view_t key, iree_host_size_t buffer_view_count,
@@ -30,23 +37,16 @@ typedef struct iree_hal_module_buffer_view_trace_callback_t {
   void* user_data;
 } iree_hal_module_buffer_view_trace_callback_t;
 
-typedef iree_status_t(
-    IREE_API_PTR* iree_hal_module_debug_sink_destroy_callback_fn_t)(
-    void* user_data);
-
-// Called by the runtime when the HAL module no longer needs the debug sink.
-typedef struct iree_hal_module_debug_sink_destroy_callback_t {
-  iree_hal_module_debug_sink_destroy_callback_fn_t fn;
-  void* user_data;
-} iree_hal_module_debug_sink_destroy_callback_t;
-
 // Interface for a HAL module debug event sink.
-// Any referenced user data must remain live for the lifetime of the HAL module
-// the sink is provided to.
+// Any referenced user data must remain live until all users of the sink have
+// released it via the release callback.
 typedef struct iree_hal_module_debug_sink_t {
+  // Called once per HAL module that has been provided the debug sink when it
+  // will no longer be used. Optional if there's no lifetime management of user
+  // data required.
+  iree_hal_module_debug_sink_release_callback_t release;
   // Called on each hal.buffer_view.trace.
   iree_hal_module_buffer_view_trace_callback_t buffer_view_trace;
-  iree_hal_module_debug_sink_destroy_callback_t destroy;
 } iree_hal_module_debug_sink_t;
 
 // Returns a default debug sink that outputs nothing.

runtime/src/iree/modules/hal/inline/module.c

@@ -152,9 +152,18 @@ typedef struct iree_hal_inline_module_state_t {
 } iree_hal_inline_module_state_t;
 
 static void IREE_API_PTR iree_hal_inline_module_destroy(void* base_module) {
+  IREE_TRACE_ZONE_BEGIN(z0);
+
   iree_hal_inline_module_t* module = IREE_HAL_INLINE_MODULE_CAST(base_module);
+
+  if (module->debug_sink.release.fn) {
+    module->debug_sink.release.fn(module->debug_sink.release.user_data);
+  }
+
   iree_hal_allocator_release(module->device_allocator);
   module->device_allocator = NULL;
+
+  IREE_TRACE_ZONE_END(z0);
 }
 
 static iree_status_t IREE_API_PTR

runtime/src/iree/modules/hal/module.c

@@ -47,15 +47,24 @@ typedef struct iree_hal_module_t {
   (iree_hal_module_t*)((uint8_t*)(module) + iree_vm_native_module_size());
 
 static void IREE_API_PTR iree_hal_module_destroy(void* base_module) {
+  IREE_TRACE_ZONE_BEGIN(z0);
+
   iree_hal_module_t* module = IREE_HAL_MODULE_CAST(base_module);
 
-  if (module->debug_sink.destroy.fn) {
-    module->debug_sink.destroy.fn(module->debug_sink.destroy.user_data);
+  // Release the debug sink prior to releasing devices as it may be caching
+  // device-specific information that will be unavailable once the devices are
+  // released.
+  if (module->debug_sink.release.fn) {
+    module->debug_sink.release.fn(module->debug_sink.release.user_data);
   }
 
+  // Release all devices. The module may be the last retainer and the devices
+  // (and their corresponding drivers) may be immediately unloaded.
   for (iree_host_size_t i = 0; i < module->device_count; ++i) {
     iree_hal_device_release(module->devices[i]);
   }
+
+  IREE_TRACE_ZONE_END(z0);
 }
 
 typedef struct iree_hal_module_state_t {