feat: add standalone infinirt graph bridge

include/infinicore/ops/mul_scalar.hpp

@@ -0,0 +1,14 @@
+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(MulScalar, Tensor, const Tensor &, double);
+
+Tensor mul_scalar(const Tensor &a, double alpha);
+void mul_scalar_(Tensor c, const Tensor &a, double alpha);
+
+} // namespace infinicore::op

include/infiniop.h

@@ -97,6 +97,7 @@
 #include "infiniop/ops/moe_topk_softmax.h"
 #include "infiniop/ops/mrope.h"
 #include "infiniop/ops/mul.h"
+#include "infiniop/ops/mul_scalar.h"
 #include "infiniop/ops/multi_margin_loss.h"
 #include "infiniop/ops/nrm2.h"
 #include "infiniop/ops/nsa_compress_paged_cache.h"

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);

include/infiniop/ops/mul_scalar.h

@@ -0,0 +1,25 @@
+#ifndef __INFINIOP_MUL_SCALAR_API_H__
+#define __INFINIOP_MUL_SCALAR_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopMulScalarDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateMulScalarDescriptor(infiniopHandle_t handle,
+                                                                     infiniopMulScalarDescriptor_t *desc_ptr,
+                                                                     infiniopTensorDescriptor_t output,
+                                                                     infiniopTensorDescriptor_t input);
+
+__INFINI_C __export infiniStatus_t infiniopGetMulScalarWorkspaceSize(infiniopMulScalarDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopMulScalar(infiniopMulScalarDescriptor_t desc,
+                                                     void *workspace,
+                                                     size_t workspace_size,
+                                                     void *output,
+                                                     const void *input,
+                                                     double alpha,
+                                                     void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyMulScalarDescriptor(infiniopMulScalarDescriptor_t desc);
+
+#endif

python/infinicore/__init__.py

@@ -114,6 +114,7 @@
 )
 from infinicore.ops.mrope import mrope
 from infinicore.ops.mul import mul
+from infinicore.ops.mul_scalar import mul_scalar
 from infinicore.ops.narrow import narrow
 from infinicore.ops.nrm2 import nrm2
 from infinicore.ops.paged_attention import paged_attention
@@ -234,6 +235,7 @@
     "matmul",
     "equal",
     "mul",
+    "mul_scalar",
     "diff",
     "digamma",
     "dist",

python/infinicore/ops/mul.py

@@ -1,8 +1,14 @@
+import numbers
+
 from infinicore.lib import _infinicore
+from infinicore.ops.mul_scalar import mul_scalar
 from infinicore.tensor import Tensor
 
 
 def mul(input, other, *, out=None):
+    if isinstance(other, numbers.Real):
+        return mul_scalar(input, other, out=out)
+
     if out is None:
         return Tensor(_infinicore.mul(input._underlying, other._underlying))
 

python/infinicore/ops/mul_scalar.py

@@ -0,0 +1,12 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def mul_scalar(input, alpha, *, out=None):
+    alpha = float(alpha)
+    if out is None:
+        return Tensor(_infinicore.mul_scalar(input._underlying, alpha))
+
+    _infinicore.mul_scalar_(out._underlying, input._underlying, alpha)
+
+    return out

python/infinicore/tensor.py

@@ -128,6 +128,9 @@ def __matmul__(self, other):
     def __mul__(self, other):
         return infinicore.mul(self, other)
 
+    def __rmul__(self, other):
+        return infinicore.mul(self, other)
+
     def narrow(self, dim, start, length):
         return infinicore.narrow(self, dim, start, length)
 

src/bridge/infini/rt.hpp

@@ -0,0 +1,80 @@
+#pragma once
+
+#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_KUNLUN:
+        return ::infini::rt::Device::Type::kKunlun;
+    case INFINI_DEVICE_HYGON:
+        return ::infini::rt::Device::Type::kHygon;
+    case INFINI_DEVICE_QY:
+        return ::infini::rt::Device::Type::kQy;
+    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::kKunlun:
+        return INFINI_DEVICE_KUNLUN;
+    case ::infini::rt::Device::Type::kHygon:
+        return INFINI_DEVICE_HYGON;
+    case ::infini::rt::Device::Type::kQy:
+        return INFINI_DEVICE_QY;
+    default:
+        return INFINI_DEVICE_TYPE_COUNT;
+    }
+}
+
+inline ::infini::rt::runtime::Stream to_rt_stream(infinirtStream_t stream) {
+    return reinterpret_cast<::infini::rt::runtime::Stream>(stream);
+}
+
+inline infinirtStream_t to_core_stream(::infini::rt::runtime::Stream stream) {
+    return reinterpret_cast<infinirtStream_t>(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,27 @@ 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::to_rt_stream(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::to_rt_stream(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,38 @@ 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 constexpr (infini::rt::DeviceEnabled<infini::rt::Device::Type::kAscend>::value) {
+        infini::rt::set_runtime_device_type(bridge::infini::rt::translate(INFINI_DEVICE_ASCEND));
+        INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::GetDeviceCount(&device_counter[static_cast<int>(Device::Type::ASCEND)])));
+        runtime_table_[static_cast<int>(Device::Type::ASCEND)].resize(device_counter[static_cast<int>(Device::Type::ASCEND)]);
+        if (device_counter[static_cast<int>(Device::Type::ASCEND)] > 0) {
+            runtime_table_[static_cast<int>(Device::Type::ASCEND)][0] = std::unique_ptr<Runtime>(new Runtime(Device(Device::Type::ASCEND, 0)));
+            current_runtime_ = runtime_table_[static_cast<int>(Device::Type::ASCEND)][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();
     }
 }