diff --git a/src/infinicore/graph/graph.cc b/src/infinicore/graph/graph.cc index 3b8fc57e5..ab38f46e4 100644 --- a/src/infinicore/graph/graph.cc +++ b/src/infinicore/graph/graph.cc @@ -2,6 +2,7 @@ #include "../utils.hpp" #include "infinicore/context/context.hpp" +#include "standalone_infinirt_graph_bridge.hpp" #include namespace infinicore::graph { @@ -32,9 +33,11 @@ DispatchableGraphOperator::~DispatchableGraphOperator() { * ========================= */ struct Graph::DeviceGraph { - infinirtGraph_t graph; - infinirtGraphExec_t exec; - infinirtGraphNode_t node; + infinirtGraph_t graph = nullptr; + infinirtGraphExec_t exec = nullptr; + infinirtGraphNode_t node = nullptr; + infinirtStream_t stream = nullptr; + bool standalone = false; std::vector log_buffer; DeviceGraph() { @@ -43,15 +46,27 @@ struct Graph::DeviceGraph { ~DeviceGraph() { if (exec) { - infinirtGraphExecDestroy(exec); + if (standalone) { + standalone_infinirt::graph_exec_destroy(exec); + } else { + infinirtGraphExecDestroy(exec); + } } if (graph) { - infinirtGraphDestroy(graph); + if (standalone) { + standalone_infinirt::graph_destroy(graph); + } else { + infinirtGraphDestroy(graph); + } } } void launch() { - INFINICORE_CHECK_ERROR(infinirtGraphLuanch(exec, context::getStream())); + if (standalone) { + INFINICORE_CHECK_ERROR(standalone_infinirt::graph_launch(exec, stream)); + } else { + INFINICORE_CHECK_ERROR(infinirtGraphLuanch(exec, context::getStream())); + } } }; @@ -75,6 +90,22 @@ void Graph::add_operator(std::shared_ptr op) { void Graph::instantiate() { // Reset device graph device_graph_ = std::make_unique(); + device_graph_->standalone = standalone_infinirt::available(context::getDevice()); + device_graph_->stream = context::getStream(); + if (device_graph_->standalone) { + auto set_device_status = standalone_infinirt::set_device(context::getDevice()); + if (set_device_status != INFINI_STATUS_SUCCESS) { + spdlog::warn("Standalone InfiniRT graph bridge failed to select the current device. Falling back to eager execution."); + device_graph_.reset(); + return; + } + + static bool logged_once = false; + if (!logged_once) { + logged_once = true; + spdlog::info("Using standalone InfiniRT graph bridge for graph capture and replay."); + } + } // warmup for (size_t iter = 0; iter < 5; ++iter) { @@ -82,35 +113,42 @@ void Graph::instantiate() { } infinicore::context::syncStream(); - if (infinirtStreamBeginCapture( - context::getStream(), - INFINIRT_STREAM_CAPTURE_MODE_RELAXED) - != INFINI_STATUS_SUCCESS) { + auto begin_status = device_graph_->standalone + ? standalone_infinirt::stream_begin_capture(device_graph_->stream, INFINIRT_STREAM_CAPTURE_MODE_RELAXED) + : infinirtStreamBeginCapture(context::getStream(), INFINIRT_STREAM_CAPTURE_MODE_RELAXED); + if (begin_status != INFINI_STATUS_SUCCESS) { + spdlog::warn("Fail to begin device graph capture."); + device_graph_.reset(); return; } // Run and record this->run(); - if (infinirtStreamEndCapture( - context::getStream(), - &device_graph_.get()->graph) - != INFINI_STATUS_SUCCESS) { + auto end_status = device_graph_->standalone + ? standalone_infinirt::stream_end_capture(device_graph_->stream, &device_graph_.get()->graph) + : infinirtStreamEndCapture(context::getStream(), &device_graph_.get()->graph); + if (end_status != INFINI_STATUS_SUCCESS) { + spdlog::warn("Fail to end device graph capture."); + device_graph_.reset(); return; } - if (infinirtGraphInstantiate( - &device_graph_.get()->exec, - device_graph_.get()->graph, - &device_graph_.get()->node, - device_graph_.get()->log_buffer.data(), - device_graph_.get()->log_buffer.size()) - != INFINI_STATUS_SUCCESS) { + auto instantiate_status = device_graph_->standalone + ? standalone_infinirt::graph_instantiate(&device_graph_.get()->exec, device_graph_.get()->graph) + : infinirtGraphInstantiate( + &device_graph_.get()->exec, + device_graph_.get()->graph, + &device_graph_.get()->node, + device_graph_.get()->log_buffer.data(), + device_graph_.get()->log_buffer.size()); + if (instantiate_status != INFINI_STATUS_SUCCESS) { static bool warned_once = false; if (!warned_once) { warned_once = true; spdlog::warn("Fail to instantiate device graph: {}", std::string(device_graph_.get()->log_buffer.data())); } + device_graph_.reset(); } } diff --git a/src/infinicore/graph/standalone_infinirt_graph_bridge.cc b/src/infinicore/graph/standalone_infinirt_graph_bridge.cc new file mode 100644 index 000000000..be4c90940 --- /dev/null +++ b/src/infinicore/graph/standalone_infinirt_graph_bridge.cc @@ -0,0 +1,173 @@ +#include "standalone_infinirt_graph_bridge.hpp" + +#ifdef USE_STANDALONE_INFINIRT_GRAPH + +#include +#include + +#include + +namespace infinicore::graph::standalone_infinirt { +namespace { + +using StandaloneDevice = infini::rt::Device; +using StandaloneRuntime = infini::rt::runtime::Runtime; + +bool truthy_env(const char *name) { + auto value = std::getenv(name); + if (value == nullptr) { + return false; + } + std::string text{value}; + return text == "1" || text == "ON" || text == "on" || text == "true" || text == "TRUE"; +} + +bool supports_device(Device::Type type) { + return type == Device::Type::NVIDIA; +} + +template +infiniStatus_t to_core_status(Status status) { + return status == StandaloneRuntime::kSuccess + ? INFINI_STATUS_SUCCESS + : INFINI_STATUS_INTERNAL_ERROR; +} + +StandaloneRuntime::Stream to_standalone_stream(infinirtStream_t stream) { + return reinterpret_cast(stream); +} + +StandaloneRuntime::Graph to_standalone_graph(infinirtGraph_t graph) { + return reinterpret_cast(graph); +} + +StandaloneRuntime::GraphExec to_standalone_graph_exec(infinirtGraphExec_t graph_exec) { + return reinterpret_cast(graph_exec); +} + +decltype(StandaloneRuntime::kStreamCaptureModeRelaxed) +to_standalone_capture_mode(infinirtStreamCaptureMode_t mode) { + switch (mode) { + case INFINIRT_STREAM_CAPTURE_MODE_GLOBAL: + return StandaloneRuntime::kStreamCaptureModeGlobal; + case INFINIRT_STREAM_CAPTURE_MODE_THREAD_LOCAL: + return StandaloneRuntime::kStreamCaptureModeThreadLocal; + case INFINIRT_STREAM_CAPTURE_MODE_RELAXED: + return StandaloneRuntime::kStreamCaptureModeRelaxed; + } + return StandaloneRuntime::kStreamCaptureModeRelaxed; +} + +} // namespace + +bool enabled() { + return truthy_env("INFINICORE_USE_STANDALONE_INFINIRT_GRAPH"); +} + +bool available(const Device &device) { + return enabled() && supports_device(device.getType()); +} + +infiniStatus_t set_device(const Device &device) { + if (!supports_device(device.getType())) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + return to_core_status(StandaloneRuntime::SetDevice(static_cast(device.getIndex()))); +} + +infiniStatus_t stream_begin_capture(infinirtStream_t stream, infinirtStreamCaptureMode_t mode) { + if (stream == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + return to_core_status(StandaloneRuntime::StreamBeginCapture( + to_standalone_stream(stream), + to_standalone_capture_mode(mode))); +} + +infiniStatus_t stream_end_capture(infinirtStream_t stream, infinirtGraph_t *graph) { + if (stream == nullptr || graph == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + return to_core_status(StandaloneRuntime::StreamEndCapture( + to_standalone_stream(stream), + reinterpret_cast(graph))); +} + +infiniStatus_t graph_destroy(infinirtGraph_t graph) { + if (graph == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + return to_core_status(StandaloneRuntime::GraphDestroy(to_standalone_graph(graph))); +} + +infiniStatus_t graph_instantiate(infinirtGraphExec_t *graph_exec, infinirtGraph_t graph) { + if (graph_exec == nullptr || graph == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + return to_core_status(StandaloneRuntime::GraphInstantiate( + reinterpret_cast(graph_exec), + to_standalone_graph(graph))); +} + +infiniStatus_t graph_exec_destroy(infinirtGraphExec_t graph_exec) { + if (graph_exec == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + return to_core_status(StandaloneRuntime::GraphExecDestroy( + to_standalone_graph_exec(graph_exec))); +} + +infiniStatus_t graph_launch(infinirtGraphExec_t graph_exec, infinirtStream_t stream) { + if (graph_exec == nullptr || stream == nullptr) { + return INFINI_STATUS_NULL_POINTER; + } + return to_core_status(StandaloneRuntime::GraphLaunch( + to_standalone_graph_exec(graph_exec), + to_standalone_stream(stream))); +} + +} // namespace infinicore::graph::standalone_infinirt + +#else + +namespace infinicore::graph::standalone_infinirt { + +bool enabled() { + return false; +} + +bool available(const Device &) { + return false; +} + +infiniStatus_t set_device(const Device &) { + return INFINI_STATUS_NOT_IMPLEMENTED; +} + +infiniStatus_t stream_begin_capture(infinirtStream_t, infinirtStreamCaptureMode_t) { + return INFINI_STATUS_NOT_IMPLEMENTED; +} + +infiniStatus_t stream_end_capture(infinirtStream_t, infinirtGraph_t *) { + return INFINI_STATUS_NOT_IMPLEMENTED; +} + +infiniStatus_t graph_destroy(infinirtGraph_t) { + return INFINI_STATUS_NOT_IMPLEMENTED; +} + +infiniStatus_t graph_instantiate(infinirtGraphExec_t *, infinirtGraph_t) { + return INFINI_STATUS_NOT_IMPLEMENTED; +} + +infiniStatus_t graph_exec_destroy(infinirtGraphExec_t) { + return INFINI_STATUS_NOT_IMPLEMENTED; +} + +infiniStatus_t graph_launch(infinirtGraphExec_t, infinirtStream_t) { + return INFINI_STATUS_NOT_IMPLEMENTED; +} + +} // namespace infinicore::graph::standalone_infinirt + +#endif diff --git a/src/infinicore/graph/standalone_infinirt_graph_bridge.hpp b/src/infinicore/graph/standalone_infinirt_graph_bridge.hpp new file mode 100644 index 000000000..0ebbb0a6b --- /dev/null +++ b/src/infinicore/graph/standalone_infinirt_graph_bridge.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "infinicore/device.hpp" +#include + +namespace infinicore::graph::standalone_infinirt { + +bool enabled(); + +bool available(const Device &device); + +infiniStatus_t set_device(const Device &device); + +infiniStatus_t stream_begin_capture(infinirtStream_t stream, infinirtStreamCaptureMode_t mode); + +infiniStatus_t stream_end_capture(infinirtStream_t stream, infinirtGraph_t *graph); + +infiniStatus_t graph_destroy(infinirtGraph_t graph); + +infiniStatus_t graph_instantiate(infinirtGraphExec_t *graph_exec, infinirtGraph_t graph); + +infiniStatus_t graph_exec_destroy(infinirtGraphExec_t graph_exec); + +infiniStatus_t graph_launch(infinirtGraphExec_t graph_exec, infinirtStream_t stream); + +} // namespace infinicore::graph::standalone_infinirt diff --git a/xmake.lua b/xmake.lua index 820ccd0f9..9b9a63d38 100644 --- a/xmake.lua +++ b/xmake.lua @@ -253,10 +253,20 @@ option("graph") set_description("Whether to use device graph instantiating feature, such as cuda graph for nvidia") option_end() -if has_config("graph") then +option("standalone-infinirt-graph") + set_default(false) + set_showmenu(true) + set_description("Whether to route graph lifecycle calls through an installed standalone InfiniRT") +option_end() + +if has_config("graph") or has_config("standalone-infinirt-graph") then add_defines("USE_INFINIRT_GRAPH") end +if has_config("standalone-infinirt-graph") then + add_defines("USE_STANDALONE_INFINIRT_GRAPH") +end + -- InfiniCCL option("ccl") @@ -320,6 +330,30 @@ end local infiniops_external_built = false +local function find_standalone_infinirt(infinirt_root, xmake_os) + local standalone_infinirt = path.join(infinirt_root, "lib", "libinfinirt.so") + if not xmake_os.isfile(standalone_infinirt) then + standalone_infinirt = path.join(infinirt_root, "lib64", "libinfinirt.so") + end + if not xmake_os.isfile(standalone_infinirt) then + raise("Standalone InfiniRT library not found under: " .. infinirt_root) + end + return standalone_infinirt +end + +local function patch_infiniops_private_infinirt(xmake_os, infiniops_lib, standalone_infinirt, private_infinirt) + local private_soname = path.filename(private_infinirt) + xmake_os.cp(standalone_infinirt, private_infinirt) + xmake_os.execv("patchelf", {"--set-soname", private_soname, private_infinirt}) + xmake_os.execv("patchelf", {"--replace-needed", standalone_infinirt, private_soname, infiniops_lib}) + xmake_os.execv("patchelf", {"--replace-needed", "libinfinirt.so", private_soname, infiniops_lib}) + + local needed = xmake_os.iorunv("patchelf", {"--print-needed", infiniops_lib}) + if not needed:find(private_soname, 1, true) then + xmake_os.execv("patchelf", {"--add-needed", private_soname, infiniops_lib}) + end +end + local function filter_infiniops_ops_for_backend(infiniops_ops) if not infiniops_ops or #infiniops_ops == 0 then return infiniops_ops @@ -401,20 +435,10 @@ local function build_infiniops_external(xmake_os) xmake_os.execv("cmake", {"--build", infiniops_builddir, "--target", "infiniops"}) xmake_os.execv("cmake", {"--install", infiniops_builddir, "--prefix", INFINI_ROOT}) if infinirt_root and infinirt_root ~= "" then - local standalone_infinirt = path.join(infinirt_root, "lib", "libinfinirt.so") - if not xmake_os.isfile(standalone_infinirt) then - standalone_infinirt = path.join(infinirt_root, "lib64", "libinfinirt.so") - end - if not xmake_os.isfile(standalone_infinirt) then - raise("Standalone InfiniRT library not found under: " .. infinirt_root) - end + local standalone_infinirt = find_standalone_infinirt(infinirt_root, xmake_os) local infiniops_lib = path.join(INFINI_ROOT, "lib", "libinfiniops.so") - local private_soname = "libinfiniops_infinirt.so" - local private_infinirt = path.join(INFINI_ROOT, "lib", private_soname) - xmake_os.cp(standalone_infinirt, private_infinirt) - xmake_os.execv("patchelf", {"--set-soname", private_soname, private_infinirt}) - xmake_os.execv("patchelf", {"--replace-needed", standalone_infinirt, private_soname, infiniops_lib}) - xmake_os.execv("patchelf", {"--replace-needed", "libinfinirt.so", private_soname, infiniops_lib}) + local private_infinirt = path.join(INFINI_ROOT, "lib", "libinfiniops_infinirt.so") + patch_infiniops_private_infinirt(xmake_os, infiniops_lib, standalone_infinirt, private_infinirt) end infiniops_external_built = true end @@ -645,6 +669,13 @@ target("infinicore_cpp_api") add_defines("CHAR_BIT=8", "INT_MIN=(-2147483647 - 1)", "INT_MAX=2147483647", "UINT_MAX=4294967295U") end add_includedirs(INFINI_ROOT.."/include", { public = true }) + local graph_infinirt_root = get_standalone_infinirt_root() + if has_config("standalone-infinirt-graph") then + if not graph_infinirt_root or graph_infinirt_root == "" then + raise("--standalone-infinirt-graph requires --infinirt-root or INFINI_RT_ROOT") + end + add_includedirs(graph_infinirt_root .. "/include") + end if has_config("nv-gpu") then local cuda_root = os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or get_config("cuda") or "/usr/local/cuda" add_includedirs(cuda_root .. "/include") @@ -672,19 +703,10 @@ target("infinicore_cpp_api") local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") local infinirt_root = get_standalone_infinirt_root() if infinirt_root and infinirt_root ~= "" then - local standalone_infinirt = path.join(infinirt_root, "lib", "libinfinirt.so") - if not os.isfile(standalone_infinirt) then - standalone_infinirt = path.join(infinirt_root, "lib64", "libinfinirt.so") - end - if not os.isfile(standalone_infinirt) then - raise("Standalone InfiniRT library not found under: " .. infinirt_root) - end + local standalone_infinirt = find_standalone_infinirt(infinirt_root, os) local infiniops_lib = path.join(INFINI_ROOT, "lib", "libinfiniops.so") - local private_soname = "libinfiniops_infinirt.so" - local private_infinirt = path.join(INFINI_ROOT, "lib", private_soname) - os.cp(standalone_infinirt, private_infinirt) - os.execv("patchelf", {"--set-soname", private_soname, private_infinirt}) - os.execv("patchelf", {"--replace-needed", "libinfinirt.so", private_soname, infiniops_lib}) + local private_infinirt = path.join(INFINI_ROOT, "lib", "libinfiniops_infinirt.so") + patch_infiniops_private_infinirt(os, infiniops_lib, standalone_infinirt, private_infinirt) end end) end @@ -918,19 +940,9 @@ target("_infinicore") end local infinirt_root = get_standalone_infinirt_root() if infinirt_root and infinirt_root ~= "" then - local standalone_infinirt = path.join(infinirt_root, "lib", "libinfinirt.so") - if not os.isfile(standalone_infinirt) then - standalone_infinirt = path.join(infinirt_root, "lib64", "libinfinirt.so") - end - if not os.isfile(standalone_infinirt) then - raise("Standalone InfiniRT library not found under: " .. infinirt_root) - end - local private_soname = "libinfiniops_infinirt.so" - local private_infinirt = path.join(INFINI_ROOT, "lib", private_soname) - os.cp(standalone_infinirt, private_infinirt) - os.execv("patchelf", {"--set-soname", private_soname, private_infinirt}) - os.execv("patchelf", {"--replace-needed", standalone_infinirt, private_soname, infiniops_lib}) - os.execv("patchelf", {"--replace-needed", "libinfinirt.so", private_soname, infiniops_lib}) + local standalone_infinirt = find_standalone_infinirt(infinirt_root, os) + local private_infinirt = path.join(INFINI_ROOT, "lib", "libinfiniops_infinirt.so") + patch_infiniops_private_infinirt(os, infiniops_lib, standalone_infinirt, private_infinirt) end os.mkdir(path.join(INFINI_ROOT, "lib")) if not infiniops_lib_installed then