Refactor/modern infini stack review

include/infini/rt/device_guard.h

@@ -0,0 +1,70 @@
+#ifndef INFINI_RT_DEVICE_GUARD_H_
+#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.
+class DeviceGuard {
+ public:
+  DeviceGuard(Device::Type target_type, int target_index) noexcept
+      : previous_type_(runtime_device_type()) {
+    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) {
+      active_ = true;
+      restored_ = true;
+      return;
+    }
+    set_runtime_device_type(target_type);
+    const auto status = runtime::SetDevice(target_index);
+    active_ = status == runtime::kSuccess;
+    warn(status, "select target device");
+    if (!active_) restore();
+  }
+
+  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");
+  }
+
+  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 {
+    if (status != runtime::kSuccess)
+      std::fprintf(stderr, "InfiniRT device scope %s error=%lld\n", operation,
+                   static_cast<long long>(status));
+  }
+  void restore() noexcept {
+    if (restored_) return;
+    restored_ = true;
+    set_runtime_device_type(previous_type_);
+    if (restore_index_) warn(runtime::SetDevice(previous_index_), "restore previous device");
+  }
+
+  Device::Type previous_type_;
+  int previous_index_ = 0;
+  bool restore_index_ = false;
+  bool active_ = false;
+  bool restored_ = false;
+};
+
+}  // namespace infini::rt
+
+#endif

include/infini/rt/reusable_memory_pool.h

@@ -0,0 +1,55 @@
+#ifndef INFINI_RT_REUSABLE_MEMORY_POOL_H_
+#define INFINI_RT_REUSABLE_MEMORY_POOL_H_
+
+#include <infini/rt.h>
+#include <cstddef>
+#include <memory>
+
+namespace infini::rt {
+
+// Block reuse and capture retention, independent of models/requests. The pool
+// owns one ordinary execution stream; consumers borrow stream(), never destroy
+// it. Raw allocation pointers require a live pool or a committed pin lease.
+// 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_;
+};
+
+} // namespace infini::rt
+#endif

scripts/generate_public_headers.py

@@ -221,6 +221,10 @@ def _write_generated_header(include_root, source_root, devices):
     runtime_declarations = "\n\n".join(
         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);"
+    if any(function.name == "EventQuery" for function in public_runtime_functions):
+        runtime_declarations += "\n\nbool IsEventNotReady(Error status);"
 
     path = include_root / "infini" / "rt" / "generated.h"
     path.parent.mkdir(parents=True, exist_ok=True)
@@ -521,6 +525,38 @@ def _write_runtime_dispatch_function(function, devices):
 """
 
 
+def _write_explicit_device_count_query(devices):
+    cases = "\n".join(
+        f"    case {_DEVICE_TYPES[device]}:\n"
+        f"      return QueryDeviceCount<{_DEVICE_TYPES[device]}>(count);"
+        for device in devices
+    )
+    return f"""namespace {{
+
+template <Device::Type type>
+Error QueryDeviceCount(int* count) {{
+  using DeviceRuntime = Runtime<type>;
+  int queried_count = -1;
+  const auto status = DeviceRuntime::GetDeviceCount(&queried_count);
+  if (status != DeviceRuntime::kSuccess) return static_cast<Error>(status);
+  if (queried_count < 0) return InvalidValueError();
+  *count = queried_count;
+  return kSuccess;
+}}
+
+}}  // namespace
+
+Error GetDeviceCount(Device::Type type, int* count) {{
+  if (count == nullptr) return InvalidValueError();
+  switch (type) {{
+{cases}
+    default:
+      return InvalidValueError();
+  }}
+}}
+"""
+
+
 def _runtime_header_for_device(source_root, device):
     for _, header_name, target in _DEVICE_HEADERS[device]:
         if header_name == "runtime_.h":
@@ -549,6 +585,27 @@ def _public_runtime_functions_for_devices(devices, source_root):
 
 def _write_runtime_dispatch(source_path, source_root, devices):
     functions = _public_runtime_functions_for_devices(devices, source_root)
+    event_not_ready_dispatch = ""
+    event_query = next(
+        (function for function in functions if function.name == "EventQuery"), None
+    )
+    if event_query is not None:
+        event_not_ready_cases = "\n".join(
+            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())
+        )
+        event_not_ready_dispatch = f"""
+bool IsEventNotReady(Error status) {{
+  if (status == kSuccess) return false;
+  switch (infini::rt::runtime_device_type()) {{
+{event_not_ready_cases}
+    default:
+      return false;
+  }}
+}}
+"""
     stream_capture_mode_helper = (
         """
 template <Device::Type device_type>
@@ -582,6 +639,16 @@ def _write_runtime_dispatch(source_path, source_root, devices):
         )
         for function in functions
     )
+    device_count_query = next(
+        (function for function in functions if function.name == "GetDeviceCount"), None
+    )
+    explicit_device_count_query = (
+        _write_explicit_device_count_query(
+            _devices_for_function(device_count_query, devices, source_root)
+        )
+        if device_count_query is not None
+        else ""
+    )
     set_device_type_cases = "\n".join(
         f"""    case {_DEVICE_TYPES[device]}:
       runtime_device_type_ = device_type;
@@ -660,6 +727,8 @@ def _write_runtime_dispatch(source_path, source_root, devices):
 }}  // namespace
 
 {dispatch_functions}
+{explicit_device_count_query}
+{event_not_ready_dispatch}
 }}  // namespace infini::rt::runtime
 """
     )

src/data_type.h

@@ -93,12 +93,14 @@ struct Float16 {
       // Inf
       return {static_cast<std::uint16_t>(sign | 0x7C00)};
     } else if (exponent >= -14) {
-      return {static_cast<std::uint16_t>(sign | ((exponent + 15) << 10) |
-                                         (mantissa >> 13))};
-    } else if (exponent >= -24) {
+      mantissa += 0x0FFF + ((mantissa >> 13) & 1);
+      return {static_cast<std::uint16_t>(
+          sign | (((exponent + 15) << 10) + (mantissa >> 13)))};
+    } else if (exponent >= -25) {
       mantissa |= 0x800000;
-      mantissa >>= (-14 - exponent);
-      return {static_cast<std::uint16_t>(sign | (mantissa >> 13))};
+      const unsigned shift = static_cast<unsigned>(-exponent - 1);
+      const auto bias = ((1u << (shift - 1)) - 1) + ((mantissa >> shift) & 1);
+      return {static_cast<std::uint16_t>(sign | ((mantissa + bias) >> shift))};
     }
     // Too small for subnormal: return signed zero.
     return {sign};

src/native/ascend/runtime_.h

@@ -35,12 +35,17 @@ struct Runtime<Device::Type::kAscend>
 
   static constexpr Error kSuccess = ACL_SUCCESS;
 
-  static constexpr auto SetDevice = aclrtSetDevice;
+  static Error SetDevice(int device) {
+    const auto status = EnsureInitialized();
+    return status == kSuccess ? aclrtSetDevice(device) : status;
+  }
 
   static constexpr auto GetDevice = aclrtGetDevice;
 
   static auto GetDeviceCount(int* count) {
     assert(count != nullptr);
+    const auto initialized = EnsureInitialized();
+    if (initialized != kSuccess) return initialized;
     std::uint32_t device_count = 0;
     auto status = aclrtGetDeviceCount(&device_count);
     *count = static_cast<int>(device_count);
@@ -182,6 +187,16 @@ struct Runtime<Device::Type::kAscend>
   }
 
  private:
+  static Error EnsureInitialized() {
+    // Cache one live epoch; external owners must not finalize while RT is live.
+    // No static destructor finalization or post-finalize reinitialization.
+    static const Error status = [] {
+      const auto result = aclInit(nullptr);
+      return result == ACL_ERROR_REPEAT_INITIALIZE ? kSuccess : result;
+    }();
+    return status;
+  }
+
   struct RiApi {
     using CaptureBeginFn = Error (*)(Stream, int);
     using CaptureGetInfoFn = Error (*)(Stream, int*, void**);

src/native/cambricon/data_type_.h

@@ -1,21 +1,19 @@
 #ifndef INFINI_RT_CAMBRICON_DATA_TYPE__H_
 #define INFINI_RT_CAMBRICON_DATA_TYPE__H_
 
-#include "bang_bf16.h"
-#include "bang_fp16.h"
 #include "data_type.h"
 #include "native/cambricon/device_.h"
 
 namespace infini::rt {
 
 template <>
 struct TypeMap<Device::Type::kCambricon, DataType::kFloat16> {
-  using type = __half;
+  using type = Float16;
 };
 
 template <>
 struct TypeMap<Device::Type::kCambricon, DataType::kBFloat16> {
-  using type = __bang_bfloat16;
+  using type = BFloat16;
 };
 
 }  // namespace infini::rt

src/native/cuda/nvidia/runtime_.h

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