fix: correct device guards and reject pool growth during capture

include/infini/rt/device_guard.h

@@ -2,27 +2,30 @@
 #define INFINI_RT_DEVICE_GUARD_H_
 
 #include <infini/rt.h>
+
 #include <cstdio>
 #include <stdexcept>
 
 namespace infini::rt {
 
-// Select the owning resource's device and restore this thread's prior selection.
+// Select the owning resource's device and restore this thread's prior
+// selection.
 class DeviceGuard {
  public:
   DeviceGuard(Device::Type target_type, int target_index) noexcept
       : previous_type_(runtime_device_type()) {
+    if (!enabled(target_type, AllDeviceTypes{})) {
+      std::fputs("InfiniRT device scope unsupported target device type\n",
+                 stderr);
+      restored_ = true;
+      return;
+    }
     const auto previous_status = runtime::GetDevice(&previous_index_);
     restore_index_ = previous_status == runtime::kSuccess;
     warn(previous_status, "read previous device");
-    if (target_type == Device::Type::kCount) {
-      std::fputs("InfiniRT device scope invalid target device type\n", stderr);
-      restore();
-      return;
-    }
-    if (restore_index_ && previous_type_ == target_type && previous_index_ == target_index) {
+    if (restore_index_ && previous_type_ == target_type &&
+        previous_index_ == target_index) {
       active_ = true;
-      restored_ = true;
       return;
     }
     set_runtime_device_type(target_type);
@@ -34,19 +37,25 @@ class DeviceGuard {
 
   DeviceGuard(Device::Type target_type, int target_index, bool required)
       : DeviceGuard(target_type, target_index) {
-    if (required && !active_) throw std::runtime_error("cannot select runtime device");
+    if (required && !active_)
+      throw std::runtime_error("cannot select runtime device");
   }
 
-  DeviceGuard(const DeviceGuard &) = delete;
-  DeviceGuard &operator=(const DeviceGuard &) = delete;
-  DeviceGuard(DeviceGuard &&) = delete;
-  DeviceGuard &operator=(DeviceGuard &&) = delete;
+  DeviceGuard(const DeviceGuard&) = delete;
+  DeviceGuard& operator=(const DeviceGuard&) = delete;
+  DeviceGuard(DeviceGuard&&) = delete;
+  DeviceGuard& operator=(DeviceGuard&&) = delete;
   ~DeviceGuard() noexcept { restore(); }
 
   bool active() const noexcept { return active_; }
 
  private:
-  static void warn(runtime::Error status, const char *operation) noexcept {
+  template <Device::Type... types>
+  static bool enabled(Device::Type type, List<types...>) noexcept {
+    return ((type == types && DeviceEnabled<types>::value) || ...);
+  }
+
+  static void warn(runtime::Error status, const char* operation) noexcept {
     if (status != runtime::kSuccess)
       std::fprintf(stderr, "InfiniRT device scope %s error=%lld\n", operation,
                    static_cast<long long>(status));
@@ -55,7 +64,8 @@ class DeviceGuard {
     if (restored_) return;
     restored_ = true;
     set_runtime_device_type(previous_type_);
-    if (restore_index_) warn(runtime::SetDevice(previous_index_), "restore previous device");
+    if (restore_index_)
+      warn(runtime::SetDevice(previous_index_), "restore previous device");
   }
 
   Device::Type previous_type_;

include/infini/rt/reusable_memory_pool.h

@@ -2,6 +2,7 @@
 #define INFINI_RT_REUSABLE_MEMORY_POOL_H_
 
 #include <infini/rt.h>
+
 #include <cstddef>
 #include <memory>
 
@@ -13,43 +14,48 @@ namespace infini::rt {
 // Reuse/deallocation is ordered on this stream by the caller. A foreign-thread
 // release must first prove completion; uses on other streams must be joined.
 class ReusableMemoryPool {
-    struct State;
-    struct PinState;
-public:
-    class PinLease {
-    public:
-        ~PinLease() noexcept;
-        PinLease(const PinLease &) = delete;
-        PinLease &operator=(const PinLease &) = delete;
-    private:
-        friend class ReusableMemoryPool;
-        explicit PinLease(std::shared_ptr<PinState> state) noexcept;
-        std::shared_ptr<PinState> state_;
-    };
-
-    explicit ReusableMemoryPool(Device device);
-    ~ReusableMemoryPool() noexcept;
-    ReusableMemoryPool(const ReusableMemoryPool &) = delete;
-    ReusableMemoryPool &operator=(const ReusableMemoryPool &) = delete;
-
-    Device device() const;
-    runtime::Stream stream() const;
-    std::byte *allocate(std::size_t size);
-    std::byte *allocate_exact(std::size_t size);
-    void deallocate(std::byte *ptr);
-    // Retain an additional logical owner and return the physical block size.
-    std::size_t retain_allocation(void *ptr);
-
-    void begin_pin_mode();
-    std::shared_ptr<PinLease> commit_pin_mode();
-    void cancel_pin_mode();
-    // True if the logical allocation predates this pin transaction.
-    bool retain_for_capture(void *ptr);
-    std::size_t pinned_bytes();
-    void trim();
-private:
-    std::shared_ptr<State> state_;
+  struct State;
+  struct PinState;
+
+ public:
+  class PinLease {
+   public:
+    ~PinLease() noexcept;
+    PinLease(const PinLease&) = delete;
+    PinLease& operator=(const PinLease&) = delete;
+
+   private:
+    friend class ReusableMemoryPool;
+    explicit PinLease(std::shared_ptr<PinState> state) noexcept;
+    std::shared_ptr<PinState> state_;
+  };
+
+  explicit ReusableMemoryPool(Device device);
+  ~ReusableMemoryPool() noexcept;
+  ReusableMemoryPool(const ReusableMemoryPool&) = delete;
+  ReusableMemoryPool& operator=(const ReusableMemoryPool&) = delete;
+
+  Device device() const;
+  runtime::Stream stream() const;
+  // During native CUDA stream capture, only cached blocks may be reused.
+  // Preallocate before capture; a cache miss throws without invalidating it.
+  std::byte* allocate(std::size_t size);
+  std::byte* allocate_exact(std::size_t size);
+  void deallocate(std::byte* ptr);
+  // Retain an additional logical owner and return the physical block size.
+  std::size_t retain_allocation(void* ptr);
+
+  void begin_pin_mode();
+  std::shared_ptr<PinLease> commit_pin_mode();
+  void cancel_pin_mode();
+  // True if the logical allocation predates this pin transaction.
+  bool retain_for_capture(void* ptr);
+  std::size_t pinned_bytes();
+  void trim();
+
+ private:
+  std::shared_ptr<State> state_;
 };
 
-} // namespace infini::rt
+}  // namespace infini::rt
 #endif

scripts/generate_public_headers.py

@@ -222,7 +222,9 @@ def _write_generated_header(include_root, source_root, devices):
         f"{function.signature()};" for function in public_runtime_functions
     )
     if any(function.name == "GetDeviceCount" for function in public_runtime_functions):
-        runtime_declarations += "\n\nError GetDeviceCount(Device::Type type, int* count);"
+        runtime_declarations += (
+            "\n\nError GetDeviceCount(Device::Type type, int* count);"
+        )
     if any(function.name == "EventQuery" for function in public_runtime_functions):
         runtime_declarations += "\n\nbool IsEventNotReady(Error status);"
 
@@ -594,7 +596,10 @@ def _write_runtime_dispatch(source_path, source_root, devices):
             f"    case {_DEVICE_TYPES[device]}:\n"
             f"      return status == static_cast<Error>(Runtime<{_DEVICE_TYPES[device]}>::kEventNotReady);"
             for device in _devices_for_function(event_query, devices, source_root)
-            if re.search(r"\bkEventNotReady\b", _runtime_header_for_device(source_root, device).read_text())
+            if re.search(
+                r"\bkEventNotReady\b",
+                _runtime_header_for_device(source_root, device).read_text(),
+            )
         )
         event_not_ready_dispatch = f"""
 bool IsEventNotReady(Error status) {{

src/native/cuda/hygon/runtime_.h

@@ -32,6 +32,8 @@ struct Runtime<Device::Type::kHygon>
 
   static constexpr Error kSuccess = cudaSuccess;
 
+  static constexpr Error kEventNotReady = cudaErrorNotReady;
+
   static constexpr auto SetDevice = cudaSetDevice;
 
   static constexpr auto GetDevice = cudaGetDevice;

src/native/cuda/iluvatar/runtime_.h

@@ -32,6 +32,8 @@ struct Runtime<Device::Type::kIluvatar>
 
   static constexpr Error kSuccess = cudaSuccess;
 
+  static constexpr Error kEventNotReady = cudaErrorNotReady;
+
   static constexpr auto SetDevice = cudaSetDevice;
 
   static constexpr auto GetDevice = cudaGetDevice;

src/native/cuda/mars/runtime_.h

@@ -30,6 +30,8 @@ struct Runtime<Device::Type::kMars>
 
   static constexpr Error kSuccess = hcSuccess;
 
+  static constexpr Error kEventNotReady = hcErrorNotReady;
+
   static constexpr auto SetDevice = hcSetDevice;
 
   static constexpr auto GetDevice = hcGetDevice;

src/native/cuda/metax/runtime_.h

@@ -30,6 +30,8 @@ struct Runtime<Device::Type::kMetax>
 
   static constexpr Error kSuccess = mcSuccess;
 
+  static constexpr Error kEventNotReady = mcErrorNotReady;
+
   static constexpr auto SetDevice = mcSetDevice;
 
   static constexpr auto GetDevice = mcGetDevice;

src/native/cuda/moore/runtime_.h

@@ -31,6 +31,8 @@ struct Runtime<Device::Type::kMoore>
 
   static constexpr Error kSuccess = musaSuccess;
 
+  static constexpr Error kEventNotReady = musaErrorNotReady;
+
   static constexpr auto SetDevice = musaSetDevice;
 
   static constexpr auto GetDevice = [](auto&&... args) {

src/native/cuda/thead/runtime_.h

@@ -33,6 +33,8 @@ struct Runtime<Device::Type::kThead>
 
   static constexpr Error kSuccess = cudaSuccess;
 
+  static constexpr Error kEventNotReady = cudaErrorNotReady;
+
   static constexpr auto SetDevice = cudaSetDevice;
 
   static constexpr auto GetDevice = cudaGetDevice;