[Store] Refactor accelerator device registry and staging copies

mooncake-store/include/device/accelerator_device.h

@@ -0,0 +1,68 @@
+#pragma once
+
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+
+#include "pinned_host_buffer.h"
+
+namespace mooncake {
+namespace device {
+
+enum class AcceleratorVendor {
+    kNvidia,
+    kMusa,
+    kMaca,
+    kHygon,
+    kCorex,
+    kHip,
+    kAscend,
+    kSunrise,
+};
+
+enum class MemoryKind {
+    kHost,
+    kDevice,
+    kUnknown,
+};
+
+enum class CopyDirection {
+    kHostToHost,
+    kHostToDevice,
+    kDeviceToHost,
+    kDeviceToDevice,
+    kAuto,
+};
+
+struct PointerInfo {
+    MemoryKind kind = MemoryKind::kUnknown;
+    int32_t device_id = -1;
+};
+
+class AcceleratorDevice {
+   public:
+    virtual ~AcceleratorDevice() = default;
+
+    virtual AcceleratorVendor Vendor() const = 0;
+    virtual bool Available(bool ensure = false) const = 0;
+    virtual PointerInfo QueryPointer(const void* ptr) const = 0;
+    virtual int32_t CurrentDeviceId() const = 0;
+    virtual void SetContext(int32_t device_id) const = 0;
+    virtual bool Copy(void* dst, const void* src, size_t size,
+                      CopyDirection direction) const = 0;
+    virtual PinnedHostBuffer AllocatePinnedHost(size_t size) const = 0;
+};
+
+class ProbeCachedAcceleratorDevice : public AcceleratorDevice {
+   public:
+    bool Available(bool ensure = false) const override;
+
+   protected:
+    virtual bool ProbeAvailable() const = 0;
+
+   private:
+    mutable std::atomic<uint8_t> available_state_{0};
+};
+
+}  // namespace device
+}  // namespace mooncake

mooncake-store/include/device/accelerator_registry.h

@@ -0,0 +1,32 @@
+#pragma once
+
+#include <span>
+#include <vector>
+
+#include "device/accelerator_device.h"
+#include "device/runtime_accelerator.h"
+
+namespace mooncake {
+namespace device {
+
+class AcceleratorRegistry {
+   public:
+    virtual ~AcceleratorRegistry() = default;
+
+    virtual std::span<const AcceleratorDevice* const> RegisteredDevices()
+        const = 0;
+    virtual RuntimeAccelerator RuntimeAccelerators(
+        bool ensure = false) const = 0;
+    virtual const AcceleratorDevice* GetDevice(
+        AcceleratorVendor vendor) const = 0;
+};
+
+const AcceleratorRegistry& GetAcceleratorRegistry();
+
+class AcceleratorDeviceRegistrar {
+   public:
+    explicit AcceleratorDeviceRegistrar(const AcceleratorDevice& device);
+};
+
+}  // namespace device
+}  // namespace mooncake

mooncake-store/include/device/runtime_accelerator.h

@@ -0,0 +1,33 @@
+#pragma once
+
+#include <cstddef>
+#include <span>
+#include <vector>
+
+#include "device/accelerator_device.h"
+
+namespace mooncake {
+namespace device {
+
+class RuntimeAccelerator {
+   public:
+    RuntimeAccelerator() = default;
+    explicit RuntimeAccelerator(std::vector<const AcceleratorDevice*> devices);
+
+    std::span<const AcceleratorDevice* const> Devices() const;
+
+    const AcceleratorDevice* FindDeviceForPointer(
+        const void* ptr, PointerInfo* out_info = nullptr) const;
+
+    bool CopyToHost(void* dst, const void* src, size_t size) const;
+
+    bool CopyFromHost(void* dst, const void* src, size_t size) const;
+
+    bool CopyMaybeAccelerator(void* dst, const void* src, size_t size) const;
+
+   private:
+    std::vector<const AcceleratorDevice*> devices_;
+};
+
+}  // namespace device
+}  // namespace mooncake