From 36cf5aeb7edc53ffeace330ad2634ab8e5baffac Mon Sep 17 00:00:00 2001
From: Yue Zhu <zhuyue_134@qq.com>
Date: Tue, 28 Jul 2026 07:59:25 +0000
Subject: [PATCH] feat(cambricon): add TaskTopo graph runtime support

Map InfiniRT graph operations to CNRT queue capture and TaskTopo APIs so Cambricon can record, instantiate, and replay real device graphs.

Enable async memset required by captured workloads and register Cambricon graph/runtime dispatch coverage.
---
 src/native/cambricon/runtime_.h | 46 ++++++++++++++++++++++++++++++---
 tests/CMakeLists.txt            |  4 ++-
 tests/test_runtime_dispatch.cc  |  2 +-
 3 files changed, 46 insertions(+), 6 deletions(-)

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<Device::Type::kCambricon>
-    : DeviceRuntime<Runtime<Device::Type::kCambricon>> {
+    : GraphRuntime<Runtime<Device::Type::kCambricon>> {
   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<Device::Type::kCambricon>
 
   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<Device::Type::kCambricon>
 
   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<Error>(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)