diff --git a/src/native/cuda/iluvatar/runtime_.h b/src/native/cuda/iluvatar/runtime_.h index 9425559..f864482 100644 --- a/src/native/cuda/iluvatar/runtime_.h +++ b/src/native/cuda/iluvatar/runtime_.h @@ -14,13 +14,20 @@ namespace infini::rt::runtime { template <> struct Runtime - : CudaRuntime> { + : GraphRuntime, + CudaRuntime>> { using Error = cudaError_t; using Stream = cudaStream_t; + using Graph = cudaGraph_t; + + using GraphExec = cudaGraphExec_t; + using Event = cudaEvent_t; + using StreamCaptureMode = cudaStreamCaptureMode; + static constexpr Device::Type kDeviceType = Device::Type::kIluvatar; static constexpr Error kSuccess = cudaSuccess; @@ -120,6 +127,38 @@ struct Runtime static constexpr auto EventElapsedTime = [](auto&&... args) { return cudaEventElapsedTime(std::forward(args)...); }; + + static constexpr auto kStreamCaptureModeGlobal = cudaStreamCaptureModeGlobal; + + static constexpr auto kStreamCaptureModeThreadLocal = + cudaStreamCaptureModeThreadLocal; + + static constexpr auto kStreamCaptureModeRelaxed = + cudaStreamCaptureModeRelaxed; + + static constexpr auto StreamBeginCapture = [](auto&&... args) { + return cudaStreamBeginCapture(std::forward(args)...); + }; + + static constexpr auto StreamEndCapture = [](auto&&... args) { + return cudaStreamEndCapture(std::forward(args)...); + }; + + static constexpr auto GraphDestroy = [](auto&&... args) { + return cudaGraphDestroy(std::forward(args)...); + }; + + static Error GraphInstantiate(GraphExec* graph_exec, Graph graph) { + return cudaGraphInstantiate(graph_exec, graph, nullptr, nullptr, 0); + } + + static constexpr auto GraphExecDestroy = [](auto&&... args) { + return cudaGraphExecDestroy(std::forward(args)...); + }; + + static constexpr auto GraphLaunch = [](auto&&... args) { + return cudaGraphLaunch(std::forward(args)...); + }; }; static_assert(Runtime::Validate()); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 66a10c6..ecc9893 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -67,6 +67,8 @@ if(WITH_ILUVATAR) ILUVATAR infini::rt::Device::Type::kIluvatar infini/rt/iluvatar/runtime_.h 1 1 1 1 1 1 1 1) + add_infini_rt_backend_graph_test( + ILUVATAR infini::rt::Device::Type::kIluvatar 1) endif() if(WITH_HYGON)