22
33#include " ../utils.hpp"
44#include " infinicore/context/context.hpp"
5+ #include " standalone_infinirt_graph_bridge.hpp"
56#include < infinirt.h>
67
78namespace infinicore ::graph {
@@ -32,9 +33,11 @@ DispatchableGraphOperator::~DispatchableGraphOperator() {
3233 * ========================= */
3334
3435struct 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) {
7593void 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