Skip to content

Commit 4e5b2e6

Browse files
committed
feat: add standalone infinirt graph bridge
Route graph capture and replay through a dlopen-based standalone InfiniRT bridge when enabled. Package the standalone InfiniRT runtime as libinfiniops_infinirt.so and ensure libinfiniops.so records that private dependency.
1 parent 84ee85f commit 4e5b2e6

4 files changed

Lines changed: 431 additions & 60 deletions

File tree

src/infinicore/graph/graph.cc

Lines changed: 62 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
#include "../utils.hpp"
44
#include "infinicore/context/context.hpp"
5+
#include "standalone_infinirt_graph_bridge.hpp"
56
#include <infinirt.h>
67

78
namespace infinicore::graph {
@@ -32,9 +33,11 @@ DispatchableGraphOperator::~DispatchableGraphOperator() {
3233
* ========================= */
3334

3435
struct Graph::DeviceGraph {
35-
infinirtGraph_t graph;
36-
infinirtGraphExec_t exec;
37-
infinirtGraphNode_t node;
36+
infinirtGraph_t graph = nullptr;
37+
infinirtGraphExec_t exec = nullptr;
38+
infinirtGraphNode_t node = nullptr;
39+
infinirtStream_t stream = nullptr;
40+
bool standalone = false;
3841
std::vector<char> log_buffer;
3942

4043
DeviceGraph() {
@@ -43,15 +46,30 @@ struct Graph::DeviceGraph {
4346

4447
~DeviceGraph() {
4548
if (exec) {
46-
infinirtGraphExecDestroy(exec);
49+
if (standalone) {
50+
standalone_infinirt::graph_exec_destroy(exec);
51+
} else {
52+
infinirtGraphExecDestroy(exec);
53+
}
4754
}
4855
if (graph) {
49-
infinirtGraphDestroy(graph);
56+
if (standalone) {
57+
standalone_infinirt::graph_destroy(graph);
58+
} else {
59+
infinirtGraphDestroy(graph);
60+
}
61+
}
62+
if (standalone && stream) {
63+
standalone_infinirt::destroy_wrapped_stream(stream);
5064
}
5165
}
5266

5367
void launch() {
54-
INFINICORE_CHECK_ERROR(infinirtGraphLuanch(exec, context::getStream()));
68+
if (standalone) {
69+
INFINICORE_CHECK_ERROR(standalone_infinirt::graph_launch(exec, stream));
70+
} else {
71+
INFINICORE_CHECK_ERROR(infinirtGraphLuanch(exec, context::getStream()));
72+
}
5573
}
5674
};
5775

@@ -75,42 +93,65 @@ void Graph::add_operator(std::shared_ptr<GraphOperator> op) {
7593
void Graph::instantiate() {
7694
// Reset device graph
7795
device_graph_ = std::make_unique<DeviceGraph>();
96+
device_graph_->standalone = standalone_infinirt::available();
97+
device_graph_->stream = device_graph_->standalone
98+
? standalone_infinirt::wrap_stream(context::getDevice(), context::getStream())
99+
: context::getStream();
100+
if (device_graph_->standalone && device_graph_->stream == nullptr) {
101+
spdlog::warn("Standalone InfiniRT graph bridge is enabled but failed to wrap the current stream. Falling back to eager execution.");
102+
device_graph_.reset();
103+
return;
104+
}
105+
if (device_graph_->standalone) {
106+
static bool logged_once = false;
107+
if (!logged_once) {
108+
logged_once = true;
109+
spdlog::info("Using standalone InfiniRT graph bridge for graph capture and replay.");
110+
}
111+
}
78112

79113
// warmup
80114
for (size_t iter = 0; iter < 5; ++iter) {
81115
this->run();
82116
}
83117
infinicore::context::syncStream();
84118

85-
if (infinirtStreamBeginCapture(
86-
context::getStream(),
87-
INFINIRT_STREAM_CAPTURE_MODE_RELAXED)
88-
!= INFINI_STATUS_SUCCESS) {
119+
auto begin_status = device_graph_->standalone
120+
? standalone_infinirt::stream_begin_capture(device_graph_->stream, INFINIRT_STREAM_CAPTURE_MODE_RELAXED)
121+
: infinirtStreamBeginCapture(context::getStream(), INFINIRT_STREAM_CAPTURE_MODE_RELAXED);
122+
if (begin_status != INFINI_STATUS_SUCCESS) {
123+
spdlog::warn("Fail to begin device graph capture.");
124+
device_graph_.reset();
89125
return;
90126
}
91127

92128
// Run and record
93129
this->run();
94130

95-
if (infinirtStreamEndCapture(
96-
context::getStream(),
97-
&device_graph_.get()->graph)
98-
!= INFINI_STATUS_SUCCESS) {
131+
auto end_status = device_graph_->standalone
132+
? standalone_infinirt::stream_end_capture(device_graph_->stream, &device_graph_.get()->graph)
133+
: infinirtStreamEndCapture(context::getStream(), &device_graph_.get()->graph);
134+
if (end_status != INFINI_STATUS_SUCCESS) {
135+
spdlog::warn("Fail to end device graph capture.");
136+
device_graph_.reset();
99137
return;
100138
}
101139

102-
if (infinirtGraphInstantiate(
103-
&device_graph_.get()->exec,
104-
device_graph_.get()->graph,
105-
&device_graph_.get()->node,
106-
device_graph_.get()->log_buffer.data(),
107-
device_graph_.get()->log_buffer.size())
108-
!= INFINI_STATUS_SUCCESS) {
140+
auto instantiate_status = device_graph_->standalone
141+
? standalone_infinirt::graph_instantiate(&device_graph_.get()->exec, device_graph_.get()->graph)
142+
: infinirtGraphInstantiate(
143+
&device_graph_.get()->exec,
144+
device_graph_.get()->graph,
145+
&device_graph_.get()->node,
146+
device_graph_.get()->log_buffer.data(),
147+
device_graph_.get()->log_buffer.size());
148+
if (instantiate_status != INFINI_STATUS_SUCCESS) {
109149
static bool warned_once = false;
110150
if (!warned_once) {
111151
warned_once = true;
112152
spdlog::warn("Fail to instantiate device graph: {}", std::string(device_graph_.get()->log_buffer.data()));
113153
}
154+
device_graph_.reset();
114155
}
115156
}
116157

0 commit comments

Comments
 (0)