Skip to content

issue/1311 - feat: refactor InfiniCore cpu runtime to InfiniRT #1342

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
spike-zhu wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

issue/1311 - feat: refactor InfiniCore cpu runtime to InfiniRT #1342

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions include/infiniop/handle.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
49 changes: 49 additions & 0 deletions src/bridge/infini/rt.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
8 changes: 4 additions & 4 deletions src/infinicore/context/allocators/device_pinned_allocator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
#include "device_pinned_allocator.hpp"

#include <infinirt.h>
#include "../../../bridge/infini/rt.hpp"

#include "../../utils.hpp"

Expand All @@ -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;
}

Expand All @@ -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);
Expand All @@ -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();
}
}
Expand Down
2 changes: 0 additions & 2 deletions src/infinicore/context/allocators/host_allocator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
#include "host_allocator.hpp"

#include <infinirt.h>

namespace infinicore {
std::byte *HostAllocator::allocate(size_t size) {
if (size == 0) {
Expand Down
13 changes: 7 additions & 6 deletions src/infinicore/context/allocators/pinnable_block_allocator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,9 @@

#include "../../utils.hpp"

#include "../../../bridge/infini/rt.hpp"

#include <algorithm>
#include <infinirt.h>
#include <stdexcept>

namespace infinicore {
Expand Down Expand Up @@ -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);
Expand All @@ -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;
Expand Down Expand Up @@ -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 {
Expand All @@ -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 {
Expand All @@ -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();
Expand Down
14 changes: 11 additions & 3 deletions src/infinicore/context/allocators/stream_ordered_allocator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
#include "stream_ordered_allocator.hpp"

#include <infinirt.h>
#include "../../../bridge/infini/rt.hpp"

#include "../../utils.hpp"

Expand All @@ -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)));
}
Comment on lines +15 to +19

@voltjia voltjia Jul 7, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这块后面可以讨论一下，看看要不要把 CPU 的 Async 改成 noop，这样就不用 if-else 了。

@spike-zhu spike-zhu Jul 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

嘉成，这个是指把原来 cpu runtime 中的：
static Error MallocAsync(void** ptr, std::size_t size, Stream) {
return kErrorInvalidValue;
}

修改为

static Error MallocAsync(void **ptr, std::size_t size, Stream) {
return Malloc(ptr, size);
}

来规避 if-else 吗？

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)));
}
Comment on lines +27 to +31

@voltjia voltjia Jul 7, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同上。

}
} // namespace infinicore
37 changes: 20 additions & 17 deletions src/infinicore/context/context_impl.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
#include "context_impl.hpp"
#include "../../bridge/infini/rt.hpp"
#include "internal.hpp"

#include "../utils.hpp"
Expand Down Expand Up @@ -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();
}
}

Expand Down
54 changes: 30 additions & 24 deletions src/infinicore/context/runtime/runtime.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
#include "runtime.hpp"
#include "../../../bridge/infini/rt.hpp"

#include "../../utils.hpp"

Expand All @@ -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);
Expand All @@ -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;
}

Expand All @@ -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() {
Expand Down Expand Up @@ -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;
}

Expand All @@ -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 {
Expand Down
Loading
Loading