diff --git a/src/native/cambricon/runtime_.h b/src/native/cambricon/runtime_.h index 4e89928..da56365 100644 --- a/src/native/cambricon/runtime_.h +++ b/src/native/cambricon/runtime_.h @@ -13,13 +13,20 @@ namespace infini::rt::runtime { template <> struct Runtime - : DeviceRuntime> { + : GraphRuntime> { using Error = cnrtRet_t; using Stream = cnrtQueue_t; + // CNRT exposes graph capture and replay as TaskTopo objects. + using Graph = cnrtTaskTopo_t; + + using GraphExec = cnrtTaskTopoEntity_t; + using Event = void*; + using StreamCaptureMode = cnrtQueueCaptureMode_t; + static constexpr Device::Type kDeviceType = Device::Type::kCambricon; #ifdef CNRT_RET_SUCCESS @@ -79,9 +86,8 @@ struct Runtime static constexpr auto Memset = cnrtMemset; - static Error MemsetAsync(void*, int, std::size_t, Stream) { - return Unsupported(); - } + // InfiniCore emits zero-fill work on the captured queue. + static constexpr auto MemsetAsync = cnrtMemsetAsync; static constexpr auto StreamCreate = cnrtQueueCreate; @@ -109,6 +115,38 @@ struct Runtime static Error EventElapsedTime(float*, Event, Event) { return Unsupported(); } + static constexpr auto kStreamCaptureModeGlobal = + cnrtQueueCaptureModeGlobal; + + static constexpr auto kStreamCaptureModeThreadLocal = + cnrtQueueCaptureModeThreadLocal; + + static constexpr auto kStreamCaptureModeRelaxed = + cnrtQueueCaptureModeRelaxed; + + static constexpr auto StreamBeginCapture = cnrtQueueBeginCapture; + + static Error StreamEndCapture(Stream stream, Graph* graph) { + assert(graph != nullptr); + return cnrtQueueEndCapture(stream, graph); + } + + static Error GraphDestroy(Graph graph) { + return graph == nullptr ? kSuccess : cnrtTaskTopoDestroy(graph); + } + + static Error GraphInstantiate(GraphExec* graph_exec, Graph graph) { + assert(graph_exec != nullptr); + return cnrtTaskTopoInstantiate(graph_exec, graph, nullptr, nullptr, 0); + } + + static Error GraphExecDestroy(GraphExec graph_exec) { + return graph_exec == nullptr ? kSuccess + : cnrtTaskTopoEntityDestroy(graph_exec); + } + + static constexpr auto GraphLaunch = cnrtTaskTopoEntityInvoke; + private: static Error Unsupported() { return static_cast(1); } }; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 66577c0..9343f3a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -128,7 +128,9 @@ if(WITH_CAMBRICON) add_infini_rt_backend_runtime_test( CAMBRICON infini::rt::Device::Type::kCambricon infini/rt/cambricon/runtime_.h - 1 0 0 0 0 0 0 0) + 1 0 0 0 1 0 0 0) + add_infini_rt_backend_graph_test( + CAMBRICON infini::rt::Device::Type::kCambricon 1) endif() if(WITH_ASCEND) diff --git a/tests/test_runtime_dispatch.cc b/tests/test_runtime_dispatch.cc index 35e0d03..637efe2 100644 --- a/tests/test_runtime_dispatch.cc +++ b/tests/test_runtime_dispatch.cc @@ -442,7 +442,7 @@ int main() { #if defined(INFINI_RT_TEST_WITH_CAMBRICON) TestDispatch(&context, infini::rt::Device::Type::kCambricon, "CAMBRICON", - {true, false, false, false, false, false, false, false}); + {true, false, false, false, true, false, false, false}); #endif #if defined(INFINI_RT_TEST_WITH_ASCEND)