Hot reload events

src/sgl/device/device.h

@@ -174,6 +174,10 @@ struct ShaderCacheStats {
     size_t miss_count;
 };
 
+/// Event data for hot reload hook.
+struct ShaderHotReloadEvent { };
+using ShaderHotReloadCallback = std::function<void(const ShaderHotReloadEvent&)>;
+
 class SGL_API Device : public Object {
     SGL_OBJECT(Device)
 public:
@@ -541,6 +545,7 @@ class SGL_API Device : public Object {
     /// - Vulkan: VkQueue (Vulkan)
     NativeHandle get_native_command_queue_handle(CommandQueueType queue = CommandQueueType::graphics) const;
 
+
     /// Enumerates all available adapters of a given device type.
     static std::vector<AdapterInfo> enumerate_adapters(DeviceType type = DeviceType::automatic);
 
@@ -565,13 +570,25 @@ class SGL_API Device : public Object {
      */
     static bool enable_agility_sdk();
 
+    /// Register a hot reload hook, called immediately after any module is reloaded.
+    void register_shader_hot_reload_callback(ShaderHotReloadCallback call_back)
+    {
+        m_shader_hot_reload_callbacks.push_back(call_back);
+    }
+
     cuda::Device* cuda_device() const { return m_cuda_device.get(); }
 
     std::string to_string() const override;
 
     Blitter* _blitter();
     HotReload* _hot_reload() { return m_hot_reload; }
 
+    /// Called by hot reload system after reload occurs, to trigger the hooks.
+    void _on_hot_reload()
+    {
+        for (auto& hook : m_shader_hot_reload_callbacks)
+            hook({});
+    }
 
 private:
     DeviceDesc m_desc;
@@ -613,6 +630,9 @@ class SGL_API Device : public Object {
     /// Transient resource heaps that are currently in flight.
     std::queue<std::pair<Slang::ComPtr<gfx::ITransientResourceHeap>, uint64_t>> m_in_flight_transient_resource_heaps;
 
+    /// List of callbacks for hot reload event
+    std::vector<ShaderHotReloadCallback> m_shader_hot_reload_callbacks;
+
     struct DeferredRelease {
         uint64_t fence_value;
         Slang::ComPtr<ISlangUnknown> object;

src/sgl/device/hot_reload.cpp

@@ -3,6 +3,7 @@
 #include "hot_reload.h"
 
 #include "sgl/core/file_system_watcher.h"
+#include "sgl/device/device.h"
 #include "sgl/device/shader.h"
 
 namespace sgl {
@@ -86,6 +87,9 @@ void HotReload::recreate_all_sessions()
 
     // Set has reloaded flag so testing system can detect changes
     m_has_reloaded = true;
+
+    // Notify device so it can notify hooks.
+    m_device->_on_hot_reload();
 }
 
 void HotReload::update_watched_paths_for_session(SlangSession* session)

src/sgl/device/python/device.cpp

@@ -145,6 +145,8 @@ SGL_PY_EXPORT(device_device)
         .def_ro("hit_count", &ShaderCacheStats::hit_count, D(ShaderCacheStats, hit_count))
         .def_ro("miss_count", &ShaderCacheStats::miss_count, D(ShaderCacheStats, miss_count));
 
+    nb::class_<ShaderHotReloadEvent>(m, "ShaderHotReloadEvent", D_NA(ShaderHotReloadEvent));
+
     nb::class_<Device, Object> device(m, "Device", D(Device));
     device.def(
         "__init__",
@@ -751,6 +753,12 @@ SGL_PY_EXPORT(device_device)
     device.def("flush_print_to_string", &Device::flush_print_to_string, D(Device, flush_print_to_string));
     device.def("run_garbage_collection", &Device::run_garbage_collection, D(Device, run_garbage_collection));
     device.def("wait", &Device::wait, D(Device, wait));
+    device.def(
+        "register_shader_hot_reload_callback",
+        &Device::register_shader_hot_reload_callback,
+        "callback"_a,
+        D_NA(Device, register_shader_hot_reload_callback)
+    );
 
     device.def_static(
         "enumerate_adapters",

src/sgl/device/tests/test_device.py

@@ -63,5 +63,32 @@ def test_global_buffer_alignment(device_type: sgl.DeviceType):
     assert np.allclose(val, texture_data, atol=1e-6)
 
 
+# Tests the hot reload event callback.
+@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
+def test_hot_reload_event(device_type: sgl.DeviceType):
+    device = helpers.get_device(type=device_type, use_cache=False)
+
+    # Load a shader
+    program = device.load_program(
+        module_name="test_shader_foo.slang",
+        entry_point_names=["main_a", "main_b", "main_vs", "main_fs"],
+    )
+
+    # Setup a hook that increments a counter on hot reload.
+    count = 0
+
+    def inc_count(x: sgl.ShaderHotReloadEvent):
+        nonlocal count
+        count += 1
+
+    device.register_shader_hot_reload_callback(inc_count)
+
+    # Force hot reload.
+    device.reload_all_programs()
+
+    # Check count.
+    assert count == 1
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])