issue/1311 - feat: refactor InfiniCore cpu and cuda-like runtime to InfiniRT

include/infiniop/handle.h

@@ -7,6 +7,8 @@ struct InfiniopHandle;
 
 typedef struct InfiniopHandle *infiniopHandle_t;
 
+__INFINI_C __export infiniStatus_t infiniopSetRuntimeDevice(infiniDevice_t device, int device_id);
+
 __INFINI_C __export infiniStatus_t infiniopCreateHandle(infiniopHandle_t *handle_ptr);
 
 __INFINI_C __export infiniStatus_t infiniopDestroyHandle(infiniopHandle_t handle);

src/bridge/infini/rt.hpp

@@ -0,0 +1,49 @@
+#pragma once
+
+#include "infinicore.h"
+#include "infinirt.h"
+
+#include <infini/rt.h>
+
+namespace infinicore::bridge::infini::rt {
+
+inline infiniStatus_t translate(::infini::rt::runtime::Error error) {
+    switch (error) {
+    case ::infini::rt::runtime::kSuccess:
+        return INFINI_STATUS_SUCCESS;
+    default:
+        return INFINI_STATUS_INTERNAL_ERROR;
+    }
+}
+
+inline ::infini::rt::Device::Type translate(infiniDevice_t device) {
+    switch (device) {
+    case INFINI_DEVICE_CPU:
+        return ::infini::rt::Device::Type::kCpu;
+    case INFINI_DEVICE_NVIDIA:
+        return ::infini::rt::Device::Type::kNvidia;
+    default:
+        return ::infini::rt::Device::Type::kCount;
+    }
+}
+
+inline infiniDevice_t translate(::infini::rt::Device::Type device) {
+    switch (device) {
+    case ::infini::rt::Device::Type::kCpu:
+        return INFINI_DEVICE_CPU;
+    case ::infini::rt::Device::Type::kNvidia:
+        return INFINI_DEVICE_NVIDIA;
+    default:
+        return INFINI_DEVICE_TYPE_COUNT;
+    }
+}
+
+inline ::infini::rt::runtime::Stream translate(infinirtStream_t stream) {
+    return reinterpret_cast<::infini::rt::runtime::Stream>(stream);
+}
+
+inline ::infini::rt::runtime::Stream *translate(infinirtStream_t *stream) {
+    return reinterpret_cast<::infini::rt::runtime::Stream *>(stream);
+}
+
+} // namespace infinicore::bridge::infini::rt

src/infinicore/context/allocators/device_pinned_allocator.cc

@@ -1,6 +1,6 @@
 #include "device_pinned_allocator.hpp"
 
-#include <infinirt.h>
+#include "../../../bridge/infini/rt.hpp"
 
 #include "../../utils.hpp"
 
@@ -16,7 +16,7 @@ std::byte *DevicePinnedHostAllocator::allocate(size_t size) {
         return nullptr;
     }
     void *ptr;
-    INFINICORE_CHECK_ERROR(infinirtMallocHost(&ptr, size));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::MallocHost(&ptr, size)));
     return (std::byte *)ptr;
 }
 
@@ -25,7 +25,7 @@ void DevicePinnedHostAllocator::deallocate(std::byte *ptr) {
         return;
     }
     if (owner_ == context::getDevice()) {
-        INFINICORE_CHECK_ERROR(infinirtFreeHost(ptr));
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::FreeHost(ptr)));
         gc();
     } else {
         gc_queue_.push(ptr);
@@ -35,7 +35,7 @@ void DevicePinnedHostAllocator::deallocate(std::byte *ptr) {
 void DevicePinnedHostAllocator::gc() {
     while (gc_queue_.empty() == false) {
         std::byte *p = gc_queue_.front();
-        INFINICORE_CHECK_ERROR(infinirtFreeHost(p));
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::FreeHost(p)));
         gc_queue_.pop();
     }
 }

src/infinicore/context/allocators/host_allocator.cc

@@ -1,7 +1,5 @@
 #include "host_allocator.hpp"
 
-#include <infinirt.h>
-
 namespace infinicore {
 std::byte *HostAllocator::allocate(size_t size) {
     if (size == 0) {

src/infinicore/context/allocators/pinnable_block_allocator.cc

@@ -4,8 +4,9 @@
 
 #include "../../utils.hpp"
 
+#include "../../../bridge/infini/rt.hpp"
+
 #include <algorithm>
-#include <infinirt.h>
 #include <stdexcept>
 
 namespace infinicore {
@@ -74,7 +75,7 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) {
             block->in_use = true;
             block->use_count = 1;
 
-            INFINICORE_CHECK_ERROR(infinirtMalloc(&block->ptr, block->size));
+            INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Malloc(&block->ptr, block->size)));
 
             all_blocks_[block->ptr] = block;
             return reinterpret_cast<std::byte *>(block->ptr);
@@ -101,7 +102,7 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) {
     block->in_use = true;
     block->use_count = 1;
 
-    INFINICORE_CHECK_ERROR(infinirtMalloc(&block->ptr, block->size));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Malloc(&block->ptr, block->size)));
 
     large_blocks_.push_back(block);
     all_blocks_[block->ptr] = block;
@@ -167,7 +168,7 @@ void PinnableBlockAllocator::trim() {
     for (auto &cls : size_classes_) {
         for (auto it = cls.free_blocks.begin(); it != cls.free_blocks.end();) {
             if (!(*it)->frozen) {
-                INFINICORE_CHECK_ERROR(infinirtFree((*it)->ptr));
+                INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Free((*it)->ptr)));
                 all_blocks_.erase((*it)->ptr);
                 it = cls.free_blocks.erase(it);
             } else {
@@ -178,7 +179,7 @@ void PinnableBlockAllocator::trim() {
     // Free non-frozen large blocks
     for (auto it = large_blocks_.begin(); it != large_blocks_.end();) {
         if (!(*it)->frozen && !(*it)->in_use) {
-            INFINICORE_CHECK_ERROR(infinirtFree((*it)->ptr));
+            INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Free((*it)->ptr)));
             all_blocks_.erase((*it)->ptr);
             it = large_blocks_.erase(it);
         } else {
@@ -192,7 +193,7 @@ PinnableBlockAllocator::~PinnableBlockAllocator() {
     std::lock_guard<std::mutex> lock(mutex_);
     for (auto &p : all_blocks_) {
         if (p.second->ptr) {
-            infinirtFree(p.second->ptr);
+            (void)infini::rt::runtime::Free(p.second->ptr);
         }
     }
     all_blocks_.clear();

src/infinicore/context/allocators/stream_ordered_allocator.cc

@@ -1,6 +1,6 @@
 #include "stream_ordered_allocator.hpp"
 
-#include <infinirt.h>
+#include "../../../bridge/infini/rt.hpp"
 
 #include "../../utils.hpp"
 
@@ -12,14 +12,22 @@ std::byte *StreamOrderedAllocator::allocate(size_t size) {
         return nullptr;
     }
     void *ptr = nullptr;
-    INFINICORE_CHECK_ERROR(infinirtMallocAsync(&ptr, size, context::getStream()));
+    if (device_.getType() != Device::Type::CPU) {
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::MallocAsync(&ptr, size, bridge::infini::rt::translate(context::getStream()))));
+    } else {
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Malloc(&ptr, size)));
+    }
     return (std::byte *)ptr;
 }
 
 void StreamOrderedAllocator::deallocate(std::byte *ptr) {
     if (ptr == nullptr) {
         return;
     }
-    INFINICORE_CHECK_ERROR(infinirtFreeAsync(ptr, context::getStream()));
+    if (device_.getType() != Device::Type::CPU) {
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::FreeAsync(ptr, bridge::infini::rt::translate(context::getStream()))));
+    } else {
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Free(ptr)));
+    }
 }
 } // namespace infinicore

src/infinicore/context/context_impl.cc

@@ -1,4 +1,5 @@
 #include "context_impl.hpp"
+#include "../../bridge/infini/rt.hpp"
 #include "internal.hpp"
 
 #include "../utils.hpp"
@@ -61,26 +62,28 @@ ContextImpl &ContextImpl::singleton() {
 }
 
 ContextImpl::ContextImpl() {
-    std::vector<int> device_counter(static_cast<size_t>(Device::Type::COUNT));
-    INFINICORE_CHECK_ERROR(infinirtGetAllDeviceCount(device_counter.data()));
-
-    // Reserve runtime slot for all devices.
-    runtime_table_[0].resize(device_counter[0]);
-    runtime_table_[0][0] = std::unique_ptr<Runtime>(new Runtime(Device(Device::Type::CPU, 0)));
-
-    // Context will try to use the first non-cpu available device as the default runtime.
-    for (int i = int(Device::Type::COUNT) - 1; i > 0; i--) {
-        if (device_counter[i] > 0) {
-            runtime_table_[i].resize(device_counter[i]);
-            if (current_runtime_ == nullptr) {
-                runtime_table_[i][0] = std::unique_ptr<Runtime>(new Runtime(Device(Device::Type(i), 0)));
-                current_runtime_ = runtime_table_[i][0].get();
-            }
+    std::vector<int> device_counter(static_cast<size_t>(Device::Type::COUNT), 0);
+
+    infini::rt::set_runtime_device_type(bridge::infini::rt::translate(INFINI_DEVICE_CPU));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::GetDeviceCount(&device_counter[static_cast<int>(Device::Type::CPU)])));
+
+    runtime_table_[static_cast<int>(Device::Type::CPU)].resize(device_counter[static_cast<int>(Device::Type::CPU)]);
+    if (device_counter[static_cast<int>(Device::Type::CPU)] > 0) {
+        runtime_table_[static_cast<int>(Device::Type::CPU)][0] = std::unique_ptr<Runtime>(new Runtime(Device(Device::Type::CPU, 0)));
+    }
+
+    if constexpr (infini::rt::DeviceEnabled<infini::rt::Device::Type::kNvidia>::value) {
+        infini::rt::set_runtime_device_type(bridge::infini::rt::translate(INFINI_DEVICE_NVIDIA));
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::GetDeviceCount(&device_counter[static_cast<int>(Device::Type::NVIDIA)])));
+        runtime_table_[static_cast<int>(Device::Type::NVIDIA)].resize(device_counter[static_cast<int>(Device::Type::NVIDIA)]);
+        if (device_counter[static_cast<int>(Device::Type::NVIDIA)] > 0) {
+            runtime_table_[static_cast<int>(Device::Type::NVIDIA)][0] = std::unique_ptr<Runtime>(new Runtime(Device(Device::Type::NVIDIA, 0)));
+            current_runtime_ = runtime_table_[static_cast<int>(Device::Type::NVIDIA)][0].get();
         }
     }
 
-    if (current_runtime_ == nullptr) {
-        current_runtime_ = runtime_table_[0][0].get();
+    if (current_runtime_ == nullptr && !runtime_table_[static_cast<int>(Device::Type::CPU)].empty()) {
+        current_runtime_ = runtime_table_[static_cast<int>(Device::Type::CPU)][0].get();
     }
 }
 

src/infinicore/context/runtime/runtime.cc

@@ -1,4 +1,5 @@
 #include "runtime.hpp"
+#include "../../../bridge/infini/rt.hpp"
 
 #include "../../utils.hpp"
 
@@ -10,7 +11,7 @@
 namespace infinicore {
 Runtime::Runtime(Device device) : device_(device), graph_manager_(std::make_unique<graph::GraphManager>()) {
     activate();
-    INFINICORE_CHECK_ERROR(infinirtStreamCreate(&stream_));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::StreamCreate(bridge::infini::rt::translate(&stream_))));
     INFINICORE_CHECK_ERROR(infiniopCreateHandle(&infiniop_handle_));
     if (device_.getType() == Device::Type::CPU) {
         device_memory_allocator_ = std::make_unique<PinnableBlockAllocator>(device);
@@ -26,11 +27,14 @@ Runtime::~Runtime() {
     }
     device_memory_allocator_.reset();
     infiniopDestroyHandle(infiniop_handle_);
-    infinirtStreamDestroy(stream_);
+    (void)infini::rt::runtime::StreamDestroy(bridge::infini::rt::translate(stream_));
 }
 
 Runtime *Runtime::activate() {
-    INFINICORE_CHECK_ERROR(infinirtSetDevice((infiniDevice_t)device_.getType(), (int)device_.getIndex()));
+    auto rt_device = bridge::infini::rt::translate(static_cast<infiniDevice_t>(device_.getType()));
+    INFINICORE_ASSERT(rt_device != infini::rt::Device::Type::kCount);
+    infini::rt::set_runtime_device_type(rt_device);
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::SetDevice(static_cast<int>(device_.getIndex()))));
     return this;
 }
 
@@ -47,11 +51,11 @@ infiniopHandle_t Runtime::infiniopHandle() const {
 }
 
 void Runtime::syncStream() {
-    INFINICORE_CHECK_ERROR(infinirtStreamSynchronize(stream_));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::StreamSynchronize(bridge::infini::rt::translate(stream_))));
 }
 
 void Runtime::syncDevice() {
-    INFINICORE_CHECK_ERROR(infinirtDeviceSynchronize());
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::DeviceSynchronize()));
 }
 
 void Runtime::trimMemory() {
@@ -103,70 +107,72 @@ std::shared_ptr<Memory> Runtime::reinstantiateBlob(std::shared_ptr<Memory> blob)
 }
 
 void Runtime::memcpyH2D(void *dst, const void *src, size_t size, bool async) {
-    if (async) {
-        INFINICORE_CHECK_ERROR(infinirtMemcpyAsync(dst, src, size, INFINIRT_MEMCPY_H2D, stream_));
+    if (async && device_.getType() != Device::Type::CPU) {
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::MemcpyAsync(dst, src, size, infini::rt::runtime::kMemcpyHostToDevice, bridge::infini::rt::translate(stream_))));
     } else {
-        INFINICORE_CHECK_ERROR(infinirtMemcpy(dst, src, size, INFINIRT_MEMCPY_H2D));
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Memcpy(dst, src, size, infini::rt::runtime::kMemcpyHostToDevice)));
     }
 }
 
 void Runtime::memcpyD2H(void *dst, const void *src, size_t size) {
-    INFINICORE_CHECK_ERROR(infinirtMemcpy(dst, src, size, INFINIRT_MEMCPY_D2H));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Memcpy(dst, src, size, infini::rt::runtime::kMemcpyDeviceToHost)));
 }
 
 void Runtime::memcpyD2D(void *dst, const void *src, size_t size, bool async) {
-    if (async) {
-        INFINICORE_CHECK_ERROR(infinirtMemcpyAsync(dst, src, size, INFINIRT_MEMCPY_D2D, stream_));
+    if (async && device_.getType() != Device::Type::CPU) {
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::MemcpyAsync(dst, src, size, infini::rt::runtime::kMemcpyDeviceToDevice, bridge::infini::rt::translate(stream_))));
     } else {
-        INFINICORE_CHECK_ERROR(infinirtMemcpy(dst, src, size, INFINIRT_MEMCPY_D2D));
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Memcpy(dst, src, size, infini::rt::runtime::kMemcpyDeviceToDevice)));
     }
 }
 
 void Runtime::setDeviceMemory(void *ptr, int value, size_t count) {
-    INFINICORE_CHECK_ERROR(infinirtMemset(ptr, value, count));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Memset(ptr, value, count)));
 }
 
 void Runtime::setDeviceMemoryAsync(void *ptr, int value, size_t count, infinirtStream_t stream) {
-    INFINICORE_CHECK_ERROR(infinirtMemsetAsync(ptr, value, count, stream));
+    if (device_.getType() != Device::Type::CPU) {
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::MemsetAsync(ptr, value, count, bridge::infini::rt::translate(stream))));
+    } else {
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Memset(ptr, value, count)));
+    }
 }
 
 // Timing method implementations
 infinirtEvent_t Runtime::createEvent() {
     infinirtEvent_t event;
-    INFINICORE_CHECK_ERROR(infinirtEventCreate(&event));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventCreate(&event)));
     return event;
 }
 
 infinirtEvent_t Runtime::createEventWithFlags(uint32_t flags) {
     infinirtEvent_t event;
-    INFINICORE_CHECK_ERROR(infinirtEventCreateWithFlags(&event, flags));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventCreateWithFlags(&event, flags)));
     return event;
 }
 
 void Runtime::recordEvent(infinirtEvent_t event, infinirtStream_t stream) {
     if (stream == nullptr) {
         stream = stream_;
     }
-    INFINICORE_CHECK_ERROR(infinirtEventRecord(event, stream));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventRecord(event, bridge::infini::rt::translate(stream))));
 }
 
 bool Runtime::queryEvent(infinirtEvent_t event) {
-    infinirtEventStatus_t status;
-    INFINICORE_CHECK_ERROR(infinirtEventQuery(event, &status));
-    return status == INFINIRT_EVENT_COMPLETE;
+    return bridge::infini::rt::translate(infini::rt::runtime::EventQuery(event)) == INFINI_STATUS_SUCCESS;
 }
 
 void Runtime::synchronizeEvent(infinirtEvent_t event) {
-    INFINICORE_CHECK_ERROR(infinirtEventSynchronize(event));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventSynchronize(event)));
 }
 
 void Runtime::destroyEvent(infinirtEvent_t event) {
-    INFINICORE_CHECK_ERROR(infinirtEventDestroy(event));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventDestroy(event)));
 }
 
 float Runtime::elapsedTime(infinirtEvent_t start, infinirtEvent_t end) {
     float ms;
-    INFINICORE_CHECK_ERROR(infinirtEventElapsedTime(&ms, start, end));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventElapsedTime(&ms, start, end)));
     return ms;
 }
 
@@ -175,7 +181,7 @@ void Runtime::streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event) {
     if (stream == nullptr) {
         stream = stream_;
     }
-    INFINICORE_CHECK_ERROR(infinirtStreamWaitEvent(stream, event));
+    INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::StreamWaitEvent(bridge::infini::rt::translate(stream), event, 0)));
 }
 
 bool Runtime::isGraphRecording() const {