issue/978 - metax cuda graph impl and wrappings

src/infiniop/devices/metax/metax_ht2mc.h

@@ -85,4 +85,17 @@
 #define hcclSuccess mcclSuccess
 #define hcclCommDestroy mcclCommDestroy
 #define hcclAllReduce mcclAllReduce
+#define hcStreamCaptureMode mcStreamCaptureMode
+#define hcStreamCaptureModeGlobal mcStreamCaptureModeGlobal
+#define hcStreamCaptureModeThreadLocal mcStreamCaptureModeThreadLocal
+#define hcStreamCaptureModeRelaxed mcStreamCaptureModeRelaxed
+#define hcStreamBeginCapture mcStreamBeginCapture
+#define hcStreamEndCapture mcStreamEndCapture
+#define hcGraph_t mcGraph_t
+#define hcGraphExec_t mcGraphExec_t
+#define hcGraphNode_t mcGraphNode_t
+#define hcGraphInstantiate mcGraphInstantiate
+#define hcGraphDestroy mcGraphDestroy
+#define hcGraphExecDestroy mcGraphExecDestroy
+#define hcGraphLaunch mcGraphLaunch
 #endif

src/infinirt/metax/infinirt_metax.cc

@@ -154,15 +154,32 @@ infiniStatus_t freeAsync(void *ptr, infinirtStream_t stream) {
 }
 
 infiniStatus_t streamBeginCapture(infinirtStream_t stream, infinirtStreamCaptureMode_t mode) {
-    return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    hcStreamCaptureMode graph_mode;
+    if (mode == INFINIRT_STREAM_CAPTURE_MODE_GLOBAL) {
+        graph_mode = hcStreamCaptureModeGlobal;
+    } else if (mode == INFINIRT_STREAM_CAPTURE_MODE_THREAD_LOCAL) {
+        graph_mode = hcStreamCaptureModeThreadLocal;
+    } else if (mode == INFINIRT_STREAM_CAPTURE_MODE_RELAXED) {
+        graph_mode = hcStreamCaptureModeRelaxed;
+    } else {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+
+    CHECK_MACART(hcStreamBeginCapture((hcStream_t)stream, graph_mode));
+
+    return INFINI_STATUS_SUCCESS;
 }
 
 infiniStatus_t streamEndCapture(infinirtStream_t stream, infinirtGraph_t *graph_ptr) {
-    return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    hcGraph_t graph;
+    CHECK_MACART(hcStreamEndCapture((hcStream_t)stream, &graph));
+    *graph_ptr = graph;
+    return INFINI_STATUS_SUCCESS;
 }
 
 infiniStatus_t graphDestroy(infinirtGraph_t graph) {
-    return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    CHECK_MACART(hcGraphDestroy((hcGraph_t)graph));
+    return INFINI_STATUS_SUCCESS;
 }
 
 infiniStatus_t graphInstantiate(
@@ -171,15 +188,23 @@ infiniStatus_t graphInstantiate(
     infinirtGraphNode_t *node_ptr,
     char *log_buffer,
     size_t buffer_size) {
-    return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    CHECK_MACART(hcGraphInstantiate(
+        (hcGraphExec_t *)graph_exec_ptr,
+        (hcGraph_t)graph,
+        (hcGraphNode_t *)node_ptr,
+        log_buffer,
+        buffer_size));
+    return INFINI_STATUS_SUCCESS;
 }
 
 infiniStatus_t graphExecDestroy(infinirtGraphExec_t graph_exec) {
-    return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    CHECK_MACART(hcGraphExecDestroy((hcGraphExec_t)graph_exec));
+    return INFINI_STATUS_SUCCESS;
 }
 
 infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream) {
-    return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    CHECK_MACART(hcGraphLaunch((hcGraphExec_t)graph_exec, (hcStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
 }
 
 } // namespace infinirt::metax