From b39a7aa1c96f2f4b403899aa96d682b3f94f25b5 Mon Sep 17 00:00:00 2001 From: zhushuang Date: Mon, 6 Jul 2026 12:04:40 +0000 Subject: [PATCH] issue/1311 - feat: refactor infinicore cpu runtime with infinirt --- include/infiniop/handle.h | 2 + src/bridge/infini/rt.hpp | 49 +++++++++++++++ .../allocators/device_pinned_allocator.cc | 8 +-- .../context/allocators/host_allocator.cc | 2 - .../allocators/pinnable_block_allocator.cc | 13 ++-- .../allocators/stream_ordered_allocator.cc | 14 ++++- src/infinicore/context/context_impl.cc | 37 ++++++------ src/infinicore/context/runtime/runtime.cc | 54 +++++++++-------- src/infinicore/graph/graph.cc | 19 ++++-- src/infiniop/devices/handle.cc | 23 +++++-- test/infiniop/libinfiniop/liboperators.py | 24 ++++---- test/infiniop/libinfiniop/utils.py | 19 +++--- xmake.lua | 60 +++---------------- xmake/cpu.lua | 1 + xmake/nvidia.lua | 2 +- xmake/test.lua | 10 +++- 16 files changed, 195 insertions(+), 142 deletions(-) create mode 100644 src/bridge/infini/rt.hpp diff --git a/include/infiniop/handle.h b/include/infiniop/handle.h index da8689542..207ca2ce9 100644 --- a/include/infiniop/handle.h +++ b/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); diff --git a/src/bridge/infini/rt.hpp b/src/bridge/infini/rt.hpp new file mode 100644 index 000000000..d53cc49ac --- /dev/null +++ b/src/bridge/infini/rt.hpp @@ -0,0 +1,49 @@ +#pragma once + +#include "infinicore.h" +#include "infinirt.h" + +#include + +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 diff --git a/src/infinicore/context/allocators/device_pinned_allocator.cc b/src/infinicore/context/allocators/device_pinned_allocator.cc index 36ee599ff..a8764d5f6 100644 --- a/src/infinicore/context/allocators/device_pinned_allocator.cc +++ b/src/infinicore/context/allocators/device_pinned_allocator.cc @@ -1,6 +1,6 @@ #include "device_pinned_allocator.hpp" -#include +#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(); } } diff --git a/src/infinicore/context/allocators/host_allocator.cc b/src/infinicore/context/allocators/host_allocator.cc index 41a55b453..0f1cc7072 100644 --- a/src/infinicore/context/allocators/host_allocator.cc +++ b/src/infinicore/context/allocators/host_allocator.cc @@ -1,7 +1,5 @@ #include "host_allocator.hpp" -#include - namespace infinicore { std::byte *HostAllocator::allocate(size_t size) { if (size == 0) { diff --git a/src/infinicore/context/allocators/pinnable_block_allocator.cc b/src/infinicore/context/allocators/pinnable_block_allocator.cc index 32e5c5e9b..3c70376f0 100644 --- a/src/infinicore/context/allocators/pinnable_block_allocator.cc +++ b/src/infinicore/context/allocators/pinnable_block_allocator.cc @@ -4,8 +4,9 @@ #include "../../utils.hpp" +#include "../../../bridge/infini/rt.hpp" + #include -#include #include 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(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 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(); diff --git a/src/infinicore/context/allocators/stream_ordered_allocator.cc b/src/infinicore/context/allocators/stream_ordered_allocator.cc index 3bd0aaab1..00df7fb47 100644 --- a/src/infinicore/context/allocators/stream_ordered_allocator.cc +++ b/src/infinicore/context/allocators/stream_ordered_allocator.cc @@ -1,6 +1,6 @@ #include "stream_ordered_allocator.hpp" -#include +#include "../../../bridge/infini/rt.hpp" #include "../../utils.hpp" @@ -12,7 +12,11 @@ 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; } @@ -20,6 +24,10 @@ 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 diff --git a/src/infinicore/context/context_impl.cc b/src/infinicore/context/context_impl.cc index f126beb1d..06c85955c 100644 --- a/src/infinicore/context/context_impl.cc +++ b/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 device_counter(static_cast(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(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(new Runtime(Device(Device::Type(i), 0))); - current_runtime_ = runtime_table_[i][0].get(); - } + std::vector device_counter(static_cast(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(Device::Type::CPU)]))); + + runtime_table_[static_cast(Device::Type::CPU)].resize(device_counter[static_cast(Device::Type::CPU)]); + if (device_counter[static_cast(Device::Type::CPU)] > 0) { + runtime_table_[static_cast(Device::Type::CPU)][0] = std::unique_ptr(new Runtime(Device(Device::Type::CPU, 0))); + } + + if constexpr (infini::rt::DeviceEnabled::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(Device::Type::NVIDIA)]))); + runtime_table_[static_cast(Device::Type::NVIDIA)].resize(device_counter[static_cast(Device::Type::NVIDIA)]); + if (device_counter[static_cast(Device::Type::NVIDIA)] > 0) { + runtime_table_[static_cast(Device::Type::NVIDIA)][0] = std::unique_ptr(new Runtime(Device(Device::Type::NVIDIA, 0))); + current_runtime_ = runtime_table_[static_cast(Device::Type::NVIDIA)][0].get(); } } - if (current_runtime_ == nullptr) { - current_runtime_ = runtime_table_[0][0].get(); + if (current_runtime_ == nullptr && !runtime_table_[static_cast(Device::Type::CPU)].empty()) { + current_runtime_ = runtime_table_[static_cast(Device::Type::CPU)][0].get(); } } diff --git a/src/infinicore/context/runtime/runtime.cc b/src/infinicore/context/runtime/runtime.cc index 92ba0f737..783b018c2 100644 --- a/src/infinicore/context/runtime/runtime.cc +++ b/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()) { 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(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(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(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,43 +107,47 @@ std::shared_ptr Runtime::reinstantiateBlob(std::shared_ptr 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; } @@ -147,26 +155,24 @@ 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 { diff --git a/src/infinicore/graph/graph.cc b/src/infinicore/graph/graph.cc index 3b8fc57e5..3511b10a1 100644 --- a/src/infinicore/graph/graph.cc +++ b/src/infinicore/graph/graph.cc @@ -2,7 +2,10 @@ #include "../utils.hpp" #include "infinicore/context/context.hpp" + +#ifdef USE_INFINIRT_GRAPH #include +#endif namespace infinicore::graph { @@ -31,6 +34,7 @@ DispatchableGraphOperator::~DispatchableGraphOperator() { * Graph * ========================= */ +#ifdef USE_INFINIRT_GRAPH struct Graph::DeviceGraph { infinirtGraph_t graph; infinirtGraphExec_t exec; @@ -54,17 +58,22 @@ struct Graph::DeviceGraph { INFINICORE_CHECK_ERROR(infinirtGraphLuanch(exec, context::getStream())); } }; +#else +struct Graph::DeviceGraph {}; +#endif Graph::Graph() { } void Graph::run() const { +#ifdef USE_INFINIRT_GRAPH if (device_graph_ != nullptr && device_graph_.get()->exec != nullptr) { device_graph_.get()->launch(); - } else { - for (auto &op : op_list_) { - op->run(); - } + return; + } +#endif + for (auto &op : op_list_) { + op->run(); } } @@ -73,6 +82,7 @@ void Graph::add_operator(std::shared_ptr op) { } void Graph::instantiate() { +#ifdef USE_INFINIRT_GRAPH // Reset device graph device_graph_ = std::make_unique(); @@ -112,6 +122,7 @@ void Graph::instantiate() { spdlog::warn("Fail to instantiate device graph: {}", std::string(device_graph_.get()->log_buffer.data())); } } +#endif } Graph::~Graph() = default; diff --git a/src/infiniop/devices/handle.cc b/src/infiniop/devices/handle.cc index b36d3b6bf..9ded0587b 100644 --- a/src/infiniop/devices/handle.cc +++ b/src/infiniop/devices/handle.cc @@ -1,6 +1,6 @@ #include "infiniop/handle.h" +#include "../../bridge/infini/rt.hpp" #include "../../utils.h" -#include "infinirt.h" #ifdef ENABLE_CPU_API #include "cpu/cpu_handle.h" @@ -24,14 +24,29 @@ #include "metax/metax_handle.h" #endif +__INFINI_C infiniStatus_t infiniopSetRuntimeDevice(infiniDevice_t device, int device_id) { + auto rt_device = infinicore::bridge::infini::rt::translate(device); + if (rt_device == infini::rt::Device::Type::kCount) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + + infini::rt::set_runtime_device_type(rt_device); + CHECK_STATUS(infinicore::bridge::infini::rt::translate(infini::rt::runtime::SetDevice(device_id))); + return INFINI_STATUS_SUCCESS; +} + __INFINI_C infiniStatus_t infiniopCreateHandle(infiniopHandle_t *handle_ptr) { if (handle_ptr == nullptr) { return INFINI_STATUS_NULL_POINTER; } - infiniDevice_t device; - int device_id; - CHECK_STATUS(infinirtGetDevice(&device, &device_id)); + infiniDevice_t device = infinicore::bridge::infini::rt::translate(infini::rt::runtime_device_type()); + if (device == INFINI_DEVICE_TYPE_COUNT) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + + int device_id = 0; + CHECK_STATUS(infinicore::bridge::infini::rt::translate(infini::rt::runtime::GetDevice(&device_id))); #define CREATE(CASE, NAMESPACE) \ case CASE: \ diff --git a/test/infiniop/libinfiniop/liboperators.py b/test/infiniop/libinfiniop/liboperators.py index ceffd5141..48ce04fe4 100644 --- a/test/infiniop/libinfiniop/liboperators.py +++ b/test/infiniop/libinfiniop/liboperators.py @@ -1,11 +1,12 @@ +import ctypes import os import platform -import ctypes -from ctypes import c_int, c_int64, c_uint64, POINTER +from ctypes import POINTER, c_int, c_int64, c_uint64 +from pathlib import Path + from .datatypes import * from .devices import * from .op_register import OpRegister -from pathlib import Path from .structs import * INFINI_ROOT = os.getenv("INFINI_ROOT") or str(Path.home() / ".infini") @@ -45,12 +46,12 @@ def find_library_in_ld_path(subdir, library_name): libop_path = find_library_in_ld_path("lib", "libinfiniop.so") librt_path = find_library_in_ld_path("lib", "libinfinirt.so") - assert ( - libop_path is not None - ), f"Cannot find infiniop.dll or libinfiniop.so. Check if INFINI_ROOT is set correctly." - assert ( - librt_path is not None - ), f"Cannot find infinirt.dll or libinfinirt.so. Check if INFINI_ROOT is set correctly." + assert libop_path is not None, ( + "Cannot find infiniop.dll or libinfiniop.so. Check if INFINI_ROOT is set correctly." + ) + assert librt_path is not None, ( + "Cannot find infinirt.dll or libinfinirt.so. Check if INFINI_ROOT is set correctly." + ) librt = ctypes.CDLL(librt_path) libop = ctypes.CDLL(libop_path) @@ -65,13 +66,12 @@ def find_library_in_ld_path(subdir, library_name): lib.infiniopCreateTensorDescriptor.restype = c_int lib.infiniopDestroyTensorDescriptor.argtypes = [infiniopTensorDescriptor_t] lib.infiniopDestroyTensorDescriptor.restype = c_int + lib.infiniopSetRuntimeDevice.argtypes = [c_int, c_int] + lib.infiniopSetRuntimeDevice.restype = c_int lib.infiniopCreateHandle.argtypes = [POINTER(infiniopHandle_t)] lib.infiniopCreateHandle.restype = c_int lib.infiniopDestroyHandle.argtypes = [infiniopHandle_t] lib.infiniopDestroyHandle.restype = c_int - lib.infinirtSetDevice.argtypes = [c_int, c_int] - lib.infinirtSetDevice.restype = c_int - OpRegister.register_lib(lib) return lib diff --git a/test/infiniop/libinfiniop/utils.py b/test/infiniop/libinfiniop/utils.py index c0560dbde..8b0b1c173 100644 --- a/test/infiniop/libinfiniop/utils.py +++ b/test/infiniop/libinfiniop/utils.py @@ -1,10 +1,12 @@ -from typing import Sequence -import torch import ctypes +from typing import Sequence + import numpy as np +import torch + from .datatypes import * from .devices import * -from .liboperators import infiniopTensorDescriptor_t, LIBINFINIOP, infiniopHandle_t +from .liboperators import LIBINFINIOP, infiniopHandle_t, infiniopTensorDescriptor_t def check_error(status): @@ -602,9 +604,10 @@ def print_discrepancy( if actual.shape != expected.shape: raise ValueError("Tensors must have the same shape to compare.") - import torch import sys + import torch + is_terminal = sys.stdout.isatty() actual = actual.to("cpu") @@ -730,7 +733,7 @@ def test_operator(device, test_func, test_cases, tensor_dtypes): to be passed to `test_func`. - tensor_dtypes (list): A list of tensor data types (e.g., `torch.float32`) to test. """ - LIBINFINIOP.infinirtSetDevice(device, ctypes.c_int(0)) + check_error(LIBINFINIOP.infiniopSetRuntimeDevice(device, ctypes.c_int(0))) handle = create_handle() tensor_dtypes = filter_tensor_dtypes_by_device(device, tensor_dtypes) try: @@ -768,12 +771,9 @@ def get_test_devices(args): if args.qy: devices_to_test.append(InfiniDeviceEnum.QY) if args.cambricon: - import torch_mlu - devices_to_test.append(InfiniDeviceEnum.CAMBRICON) if args.ascend: import torch - import torch_npu torch.npu.set_device(0) # Ascend NPU needs explicit device initialization devices_to_test.append(InfiniDeviceEnum.ASCEND) @@ -783,12 +783,9 @@ def get_test_devices(args): devices_to_test.append(InfiniDeviceEnum.METAX) if args.moore: import torch - import torch_musa devices_to_test.append(InfiniDeviceEnum.MOORE) if args.kunlun: - import torch_xmlir - devices_to_test.append(InfiniDeviceEnum.KUNLUN) if args.hygon: import torch diff --git a/xmake.lua b/xmake.lua index fcbcf423d..d98fe6e01 100644 --- a/xmake.lua +++ b/xmake.lua @@ -461,57 +461,13 @@ target("infini-utils") add_files("src/utils/*.cc") target_end() -target("infinirt") - set_kind("shared") - - if has_config("cpu") then - add_deps("infinirt-cpu") - end - if has_config("nv-gpu") then - add_deps("infinirt-nvidia") - end - if has_config("cambricon-mlu") then - add_deps("infinirt-cambricon") - end - if has_config("ascend-npu") then - add_deps("infinirt-ascend") - end - if has_config("metax-gpu") then - add_deps("infinirt-metax") - end - if has_config("moore-gpu") then - add_deps("infinirt-moore") - end - if has_config("iluvatar-gpu") then - add_deps("infinirt-iluvatar") - end - if has_config("ali-ppu") then - add_deps("infinirt-ali") - end - if has_config("qy-gpu") then - add_deps("infinirt-qy") - add_files("build/.objs/infinirt-qy/rules/qy.cuda/src/infinirt/cuda/*.cu.o", {public = true}) - end - if has_config("kunlun-xpu") then - add_deps("infinirt-kunlun") - end - if has_config("hygon-dcu") then - add_deps("infinirt-hygon") - end - set_languages("cxx17") - if not is_plat("windows") then - add_cxflags("-fPIC") - add_cxxflags("-fPIC") - add_ldflags("-fPIC", {force = true}) - end - set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini")) - add_files("src/infinirt/*.cc") - add_installfiles("include/infinirt.h", {prefixdir = "include"}) -target_end() - target("infiniop") set_kind("shared") - add_deps("infinirt") + local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") + add_includedirs(INFINI_ROOT.."/include", { public = true }) + add_linkdirs(INFINI_ROOT.."/lib") + add_links("infinirt") + add_rpathdirs(INFINI_ROOT.."/lib") local public_cuda_root = get_config("cuda") or os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") if public_cuda_root and public_cuda_root ~= "" then @@ -589,7 +545,6 @@ target_end() target("infiniccl") set_kind("shared") - add_deps("infinirt") if has_config("nv-gpu") then add_deps("infiniccl-nvidia") @@ -635,7 +590,7 @@ target_end() target("infinicore_c_api") set_kind("phony") - add_deps("infiniop", "infinirt", "infiniccl") + add_deps("infiniop", "infiniccl") after_build(function (target) print(YELLOW .. "[Congratulations!] Now you can install the libraries with \"xmake install\"" .. NC) end) target_end() @@ -650,7 +605,7 @@ target_end() target("infinicore_cpp_api") set_kind("shared") - add_deps("infiniop", "infinirt", "infiniccl") + add_deps("infiniop", "infiniccl") set_languages("cxx17") set_symbols("visibility") @@ -715,6 +670,7 @@ target("infinicore_cpp_api") add_linkdirs(INFINI_ROOT.."/lib") add_links("infiniop", "infinirt", "infiniccl") + add_rpathdirs(INFINI_ROOT.."/lib") if get_config("flash-attn") and get_config("flash-attn") ~= "" then add_installfiles("(builddir)/$(plat)/$(arch)/$(mode)/flash-attn*.so", {prefixdir = "lib"}) diff --git a/xmake/cpu.lua b/xmake/cpu.lua index e192fbbbd..a1974b823 100644 --- a/xmake/cpu.lua +++ b/xmake/cpu.lua @@ -26,6 +26,7 @@ target_end() target("infinirt-cpu") set_kind("static") + set_default(false) add_deps("infini-utils") on_install(function (target) end) diff --git a/xmake/nvidia.lua b/xmake/nvidia.lua index 282f5876f..6fbe90dfa 100644 --- a/xmake/nvidia.lua +++ b/xmake/nvidia.lua @@ -223,6 +223,7 @@ target_end() target("infinirt-nvidia") set_kind("static") + set_default(false) add_deps("infini-utils") on_install(function (target) end) @@ -246,7 +247,6 @@ target_end() target("infiniccl-nvidia") set_kind("static") - add_deps("infinirt") on_install(function (target) end) if has_config("ccl") then set_policy("build.cuda.devlink", true) diff --git a/xmake/test.lua b/xmake/test.lua index 64533163b..d4ef4abe7 100644 --- a/xmake/test.lua +++ b/xmake/test.lua @@ -54,14 +54,20 @@ target_end() target("infinirt-test") set_kind("binary") - add_deps("infinirt") + set_default(false) on_install(function (target) end) set_languages("cxx17") set_warnings("all", "error") + local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") + add_includedirs(INFINI_ROOT.."/include") + add_linkdirs(INFINI_ROOT.."/lib") + add_links("infinirt") + add_rpathdirs(INFINI_ROOT.."/lib") + add_files(os.projectdir().."/src/infinirt-test/*.cc") - set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini")) + set_installdir(INFINI_ROOT) target_end() target("infinicore-test")