diff --git a/src/infiniop/devices/metax/metax_ht2mc.h b/src/infiniop/devices/metax/metax_ht2mc.h index 2db1087d4..a1c8c1ffe 100644 --- a/src/infiniop/devices/metax/metax_ht2mc.h +++ b/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 diff --git a/src/infinirt/metax/infinirt_metax.cc b/src/infinirt/metax/infinirt_metax.cc index aca187366..a9d69a9d8 100644 --- a/src/infinirt/metax/infinirt_metax.cc +++ b/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