Issue/1386

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,77 @@
+#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;
+    case INFINI_DEVICE_CAMBRICON:
+        return ::infini::rt::Device::Type::kCambricon;
+    case INFINI_DEVICE_ASCEND:
+        return ::infini::rt::Device::Type::kAscend;
+    case INFINI_DEVICE_METAX:
+        return ::infini::rt::Device::Type::kMetax;
+    case INFINI_DEVICE_MOORE:
+        return ::infini::rt::Device::Type::kMoore;
+    case INFINI_DEVICE_ILUVATAR:
+        return ::infini::rt::Device::Type::kIluvatar;
+    case INFINI_DEVICE_HYGON:
+        return ::infini::rt::Device::Type::kHygon;
+    case INFINI_DEVICE_ALI:
+        return ::infini::rt::Device::Type::kThead;
+    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;
+    case ::infini::rt::Device::Type::kCambricon:
+        return INFINI_DEVICE_CAMBRICON;
+    case ::infini::rt::Device::Type::kAscend:
+        return INFINI_DEVICE_ASCEND;
+    case ::infini::rt::Device::Type::kMetax:
+        return INFINI_DEVICE_METAX;
+    case ::infini::rt::Device::Type::kMoore:
+        return INFINI_DEVICE_MOORE;
+    case ::infini::rt::Device::Type::kIluvatar:
+        return INFINI_DEVICE_ILUVATAR;
+    case ::infini::rt::Device::Type::kHygon:
+        return INFINI_DEVICE_HYGON;
+    case ::infini::rt::Device::Type::kThead:
+        return INFINI_DEVICE_ALI;
+    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,29 +62,40 @@ 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);
+
+#define INIT_RUNTIME(RT_DEVICE, CORE_DEVICE)                                                                 \
+    if constexpr (infini::rt::DeviceEnabled<RT_DEVICE>::value) {                                             \
+        auto device_type = static_cast<Device::Type>(CORE_DEVICE);                                           \
+        auto device_index = static_cast<int>(device_type);                                                   \
+        infini::rt::set_runtime_device_type(RT_DEVICE);                                                      \
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(                                                \
+            infini::rt::runtime::GetDeviceCount(&device_counter[device_index])));                            \
+        runtime_table_[device_index].resize(device_counter[device_index]);                                   \
+        if (device_counter[device_index] > 0) {                                                              \
+            runtime_table_[device_index][0] = std::unique_ptr<Runtime>(new Runtime(Device(device_type, 0))); \
+            if (device_type != Device::Type::CPU && current_runtime_ == nullptr) {                           \
+                current_runtime_ = runtime_table_[device_index][0].get();                                    \
+            }                                                                                                \
+        }                                                                                                    \
     }
 
-    if (current_runtime_ == nullptr) {
-        current_runtime_ = runtime_table_[0][0].get();
+    INIT_RUNTIME(infini::rt::Device::Type::kCpu, INFINI_DEVICE_CPU);
+    INIT_RUNTIME(infini::rt::Device::Type::kNvidia, INFINI_DEVICE_NVIDIA);
+    INIT_RUNTIME(infini::rt::Device::Type::kCambricon, INFINI_DEVICE_CAMBRICON);
+    INIT_RUNTIME(infini::rt::Device::Type::kAscend, INFINI_DEVICE_ASCEND);
+    INIT_RUNTIME(infini::rt::Device::Type::kMetax, INFINI_DEVICE_METAX);
+    INIT_RUNTIME(infini::rt::Device::Type::kMoore, INFINI_DEVICE_MOORE);
+    INIT_RUNTIME(infini::rt::Device::Type::kIluvatar, INFINI_DEVICE_ILUVATAR);
+    INIT_RUNTIME(infini::rt::Device::Type::kHygon, INFINI_DEVICE_HYGON);
+    INIT_RUNTIME(infini::rt::Device::Type::kThead, INFINI_DEVICE_ALI);
+
+#undef INIT_RUNTIME
+
+    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();
     }
 }
-
 namespace context {
 
 void setDevice(Device device) {