MadSpace stream handling for external users

madspace/include/madspace/driver/backend.hpp

@@ -19,6 +19,7 @@ class Runtime {
         const std::vector<bool>& eval_grad,
         bool return_contiguous_grads = false
     ) = 0;
+    virtual void release_inputs() {}
     friend std::unique_ptr<Runtime>
     build_runtime(const Function& function, ContextPtr context, bool concurrent);
 

madspace/include/madspace/driver/context.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
-#include <stdint.h>
+#include <cstdint>
+#include <optional>
 #include <unordered_map>
 
 #include "madspace/compgraphs.hpp"
@@ -200,6 +201,10 @@ ContextPtr default_cuda_context(std::size_t index = 0);
 ContextPtr default_hip_context(std::size_t index = 0);
 ContextPtr default_device_context(DevicePtr device);
 
+// thread-local stream to run on, empty to use our own and synchronize. 0 keeps ours
+std::optional<std::uintptr_t> caller_stream();
+void set_caller_stream(std::optional<std::uintptr_t> stream);
+
 inline std::string prefixed_name(const std::string& prefix, const std::string& name) {
     return prefix == "" ? name : std::format("{}.{}", prefix, name);
 }

madspace/include/madspace/driver/tensor.hpp

@@ -9,6 +9,7 @@
 #include <cstdint>
 #include <functional>
 #include <initializer_list>
+#include <optional>
 #include <vector>
 
 namespace madspace {
@@ -199,10 +200,14 @@ inline bool needs_zero_init(AllocHint hint) {
 
 class Device {
 public:
+    static constexpr bool stream_ordered_alloc = false;
+
     virtual ~Device() = default;
     virtual std::pair<void*, Tensor>
     allocate(std::size_t size, AllocHint hint) const = 0;
     virtual void free(void* ptr) const = 0;
+    virtual void free_on_stream(void* ptr, void* stream) const { free(ptr); }
+    virtual void order_streams(void* from, void* to) const {}
     virtual void memcpy(void* to, void* from, std::size_t size) const = 0;
     virtual void tensor_copy(const Tensor& source, Tensor& target) const = 0;
     virtual void tensor_zero(Tensor& tensor) const = 0;
@@ -386,7 +391,12 @@ class Tensor {
         );
     }
 
-    ~Tensor() { reset(); }
+    ~Tensor() {
+        try {
+            reset();
+        } catch (...) {
+        }
+    }
 
     Tensor& operator=(const Tensor& other) {
         reset();
@@ -398,7 +408,10 @@ class Tensor {
     }
 
     Tensor& operator=(Tensor&& other) noexcept {
-        reset();
+        try {
+            reset();
+        } catch (...) {
+        }
         impl = other.impl;
         other.impl = nullptr;
         return *this;
@@ -474,6 +487,18 @@ class Tensor {
         check_impl();
         return impl->device;
     }
+    std::optional<std::uintptr_t> stream() const {
+        return impl == nullptr ? std::nullopt : storage()->stream;
+    }
+    void set_stream(std::optional<std::uintptr_t> stream) {
+        if (impl == nullptr) {
+            return;
+        }
+        TensorImpl* owner = storage();
+        if (owner->stream_ordered) {
+            owner->stream = stream;
+        }
+    }
     std::size_t index_value() const {
         check_impl();
         if (impl->batch_sizes.size() > 0) {
@@ -499,21 +524,26 @@ class Tensor {
 
     std::size_t byte_size() const { return dtype_size() * shape().product(); }
 
-    void reset() {
+    void reset() { reset_on_stream(stream().value_or(0)); }
+
+    template <typename D>
+    void reset(const D& device) {
         if (impl == nullptr) {
             return;
         }
-        impl->reset(*impl->device);
+        // released before the free, which can throw
+        TensorImpl* owner = impl;
         impl = nullptr;
+        owner->reset(device);
     }
 
-    template <typename D>
-    void reset(const D& device) {
+    void reset_on_stream(std::uintptr_t stream) {
         if (impl == nullptr) {
             return;
         }
-        impl->reset(device);
+        TensorImpl* owner = impl;
         impl = nullptr;
+        owner->reset_on_stream(stream);
     }
 
     Tensor select(std::size_t axis, std::size_t index) const;
@@ -634,6 +664,8 @@ class Tensor {
         Sizes stride;
         std::size_t contiguous_dims;
         SizeVec batch_sizes;
+        bool stream_ordered = false;
+        std::optional<std::uintptr_t> stream;
 
         template <typename D>
         void reset(const D& device) {
@@ -651,9 +683,36 @@ class Tensor {
             delete this;
         }
 
+        void reset_on_stream(std::uintptr_t stream) {
+            if (ref_count.fetch_sub(1, std::memory_order_acq_rel) != 1) {
+                return;
+            }
+            if (owns_data && data != nullptr) {
+                if (stream_ordered) {
+                    device->free_on_stream(data, reinterpret_cast<void*>(stream));
+                } else {
+                    device->free(data);
+                }
+                --Tensor::tensor_count;
+            } else if (data_owner != nullptr) {
+                data_owner->reset_on_stream(stream);
+            } else if (external_reset) {
+                (*external_reset)();
+            }
+            delete this;
+        }
+
         void incref() { ref_count.fetch_add(1, std::memory_order_relaxed); }
     };
 
+    TensorImpl* storage() const {
+        TensorImpl* item = impl;
+        while (item->data_owner != nullptr) {
+            item = item->data_owner;
+        }
+        return item;
+    }
+
     Tensor(TensorImpl* _impl) : impl(_impl) {
         if (impl->data_owner != nullptr) {
             impl->data_owner->incref();
@@ -677,6 +736,10 @@ class Tensor {
             impl->data_owner = parent.impl;
         } else if (data != nullptr) {
             ++tensor_count;
+            if constexpr (D::stream_ordered_alloc) {
+                impl->stream_ordered = true;
+                impl->stream = reinterpret_cast<std::uintptr_t>(device.stream());
+            }
         }
     }
 

madspace/include/madspace/driver/thread_pool.hpp

@@ -72,22 +72,25 @@ class ThreadResource {
         std::optional<std::function<void(T&)>> destructor = std::nullopt
     ) :
         _pool(&pool),
+        _constructor(std::move(constructor)),
         _destructor(destructor),
-        _listener_id(pool.add_listener([this, constructor](std::size_t thread_count) {
+        _listener_id(pool.add_listener([this](std::size_t thread_count) {
             while (_resources.size() < thread_count) {
-                _resources.push_back(constructor());
+                _resources.emplace_back();
             }
         })) {
         for (std::size_t i = 0; i == 0 || i < pool.thread_count(); ++i) {
-            _resources.push_back(constructor());
+            _resources.emplace_back();
         }
+        get();
     }
     ~ThreadResource() {
         reset();
     }
     ThreadResource(ThreadResource&& other) noexcept :
         _pool(std::move(other._pool)),
         _resources(std::move(other._resources)),
+        _constructor(std::move(other._constructor)),
         _listener_id(std::move(other._listener_id)),
         _destructor(std::move(other._destructor)) {
         other._pool = nullptr;
@@ -97,29 +100,44 @@ class ThreadResource {
         reset();
         _pool = std::move(other._pool);
         _resources = std::move(other._resources);
+        _constructor = std::move(other._constructor);
         _listener_id = std::move(other._listener_id);
         _destructor = std::move(other._destructor);
         other._pool = nullptr;
         return *this;
     }
     ThreadResource(const ThreadResource&) = delete;
     ThreadResource& operator=(const ThreadResource&) = delete;
-    T& get() { return _resources.at(ThreadPool::thread_index()); }
-    const T& get() const { return _resources.at(ThreadPool::thread_index()); }
+    T& get() const {
+        auto& [flag, item] = _resources.at(ThreadPool::thread_index());
+        std::call_once(flag, [&] {
+            std::unique_lock<std::mutex> lock(construction_mutex());
+            item.emplace(_constructor());
+        });
+        return *item;
+    }
     void reset() {
         if (_pool) {
             if (_destructor) {
-                for (auto& item : _resources) {
-                    _destructor.value()(item);
+                for (auto& [flag, item] : _resources) {
+                    if (item) {
+                        _destructor.value()(*item);
+                    }
                 }
             }
             _pool->remove_listener(_listener_id);
         }
     }
 
 private:
+    static std::mutex& construction_mutex() {
+        static std::mutex mutex;
+        return mutex;
+    }
+
     ThreadPool* _pool = nullptr;
-    std::vector<T> _resources;
+    mutable std::deque<std::pair<std::once_flag, std::optional<T>>> _resources;
+    std::function<T()> _constructor;
     std::size_t _listener_id;
     std::optional<std::function<void(T&)>> _destructor;
 };

madspace/madspace/_madspace_py_loader.py

@@ -1,3 +1,4 @@
+import contextlib
 import ctypes
 import logging
 import os
@@ -17,6 +18,20 @@
 from ._madspace_py import *
 
 
+@contextlib.contextmanager
+def stream(handle):
+    """
+    Run the calls on the given cuda stream instead of madspace's own, without
+    synchronizing. Keep it alive until everything madspace handed out has been dropped.
+    """
+    previous = get_stream()
+    set_stream(handle)
+    try:
+        yield
+    finally:
+        set_stream(previous)
+
+
 def _init():
     """
     Monkey-patch classes for a more pythonic experience.
@@ -27,7 +42,7 @@ def call_and_convert(runtime, args):
         if len(args) == 0:
             tensorlib = "numpy"
         else:
-            tensorlib = type(args[0]).__module__
+            tensorlib = type(args[0]).__module__.partition(".")[0]
         outputs = runtime.call(args)
         # Convert outputs, lazy-loading torch or numpy
         if tensorlib == "torch":
@@ -122,11 +137,19 @@ def log_handler(level, message):
             case Logger.level_error:
                 py_logger.error(message)
 
+    def release_inputs(self):
+        for name in ("runtime", "forward_runtime", "inverse_runtime"):
+            if hasattr(self, name):
+                getattr(self, name).release_inputs()
+
     FunctionRuntime.__call__ = runtime_call
     Function.__call__ = function_call
     FunctionGenerator.__call__ = function_generator_call
     Mapping.map_forward = map_forward
     Mapping.map_inverse = map_inverse
+    Function.release_inputs = release_inputs
+    FunctionGenerator.release_inputs = release_inputs
+    Mapping.release_inputs = release_inputs
     Tensor.numpy = tensor_numpy
     Tensor.torch = tensor_torch
     # Logger.set_log_handler(log_handler)

madspace/madspace/torch.py

@@ -29,11 +29,14 @@ def __init__(
             ),
         )
 
+    def release_inputs(self) -> None:
+        self.runtime.release_inputs()
+
     def forward(self, *args: torch.Tensor) -> list[torch.Tensor]:
         if torch.is_grad_enabled():
             return AutogradWrapper.apply(self, self.dummy, *args)
         else:
-            outputs = self.runtime.call(args)
+            outputs = self.runtime.call([arg.detach() for arg in args])
             if len(outputs) == 1:
                 return torch.from_dlpack(outputs[0])
             else:
@@ -53,9 +56,8 @@ def forward(
         )
         ctx.module = module
         ctx.eval_grad = eval_grad
-        ctx.save_for_backward(
-            *(None if grad is None else torch.from_dlpack(grad) for grad in local_grads)
-        )
+        ctx.stream = me.get_stream()
+        ctx.stored_locals = local_grads
         if len(outputs) == 1:
             return torch.from_dlpack(outputs[0])
         else:
@@ -64,9 +66,10 @@ def forward(
     @staticmethod
     @once_differentiable
     def backward(ctx: FunctionCtx, *output_grads: torch.Tensor):
-        input_grads, global_grads = ctx.module.runtime.call_backward(
-            output_grads, ctx.saved_tensors, ctx.eval_grad
-        )
+        with me.stream(ctx.stream):
+            input_grads, global_grads = ctx.module.runtime.call_backward(
+                output_grads, ctx.stored_locals, ctx.eval_grad
+            )
         for name, grad in global_grads:
             if grad is None:
                 continue