feat(metax): add graph capture and replay support

docs/backends.md

@@ -33,7 +33,7 @@ Current test expectations are:
 | CPU | No | Yes | No | Yes | No | Yes | No |
 | NVIDIA | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
 | Iluvatar | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
-| MetaX | Yes | Yes | Yes | Yes | Yes | Yes | No |
+| MetaX | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
 | Moore | Yes | Yes | No | Yes | Yes | Yes | No |
 | Hygon | Yes | Yes | Yes | Yes | Yes | Yes | No |
 | Cambricon | Yes | No | No | No | No | No | No |

src/native/cuda/metax/runtime_.h

@@ -12,13 +12,20 @@ namespace infini::rt::runtime {
 
 template <>
 struct Runtime<Device::Type::kMetax>
-    : CudaRuntime<Runtime<Device::Type::kMetax>> {
+    : GraphRuntime<Runtime<Device::Type::kMetax>,
+                   CudaRuntime<Runtime<Device::Type::kMetax>>> {
   using Error = mcError_t;
 
   using Stream = mcStream_t;
 
+  using Graph = mcGraph_t;
+
+  using GraphExec = mcGraphExec_t;
+
   using Event = mcEvent_t;
 
+  using StreamCaptureMode = mcStreamCaptureMode;
+
   static constexpr Device::Type kDeviceType = Device::Type::kMetax;
 
   static constexpr Error kSuccess = mcSuccess;
@@ -124,6 +131,39 @@ struct Runtime<Device::Type::kMetax>
   static constexpr auto EventElapsedTime = [](auto&&... args) {
     return mcEventElapsedTime(std::forward<decltype(args)>(args)...);
   };
+
+  static constexpr auto kStreamCaptureModeGlobal =
+      mcStreamCaptureModeGlobal;
+
+  static constexpr auto kStreamCaptureModeThreadLocal =
+      mcStreamCaptureModeThreadLocal;
+
+  static constexpr auto kStreamCaptureModeRelaxed =
+      mcStreamCaptureModeRelaxed;
+
+  static constexpr auto StreamBeginCapture = [](auto&&... args) {
+    return mcStreamBeginCapture(std::forward<decltype(args)>(args)...);
+  };
+
+  static constexpr auto StreamEndCapture = [](auto&&... args) {
+    return mcStreamEndCapture(std::forward<decltype(args)>(args)...);
+  };
+
+  static constexpr auto GraphDestroy = [](auto&&... args) {
+    return mcGraphDestroy(std::forward<decltype(args)>(args)...);
+  };
+
+  static Error GraphInstantiate(GraphExec* graph_exec, Graph graph) {
+    return mcGraphInstantiate(graph_exec, graph, nullptr, nullptr, 0);
+  }
+
+  static constexpr auto GraphExecDestroy = [](auto&&... args) {
+    return mcGraphExecDestroy(std::forward<decltype(args)>(args)...);
+  };
+
+  static constexpr auto GraphLaunch = [](auto&&... args) {
+    return mcGraphLaunch(std::forward<decltype(args)>(args)...);
+  };
 };
 
 static_assert(Runtime<Device::Type::kMetax>::Validate());

tests/CMakeLists.txt

@@ -85,6 +85,8 @@ if(WITH_METAX)
         METAX infini::rt::Device::Type::kMetax
         infini/rt/metax/runtime_.h
         1 1 1 1 1 1 1 1)
+    add_infini_rt_backend_graph_test(
+        METAX infini::rt::Device::Type::kMetax 1)
 endif()
 
 if(WITH_MOORE)