adds: hip launch operations

mlcommons · Jul 26, 2024 · 3b867fd · 3b867fd
1 parent adfa2df
commit 3b867fd
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/src/trace_link/kineto_operator.py b/src/trace_link/kineto_operator.py
@@ -142,7 +142,16 @@ def is_kernel_launch_op(self) -> bool:
             "cudaMemcpyToSymbol",
             "cudaLaunchCooperativeKernel",
         }
-        return cuda_launch_categories and self.name in cuda_launch_operations
+
+        hip_launch_operations = {
+            "hipLaunchKernel",
+            "hipExtLaunchKernel",
+            "hipExtModuleLaunchKernel",
+            "hipModuleLaunchKernel",
+            "hipMemcpyWithStream",
+            "hipMemcpyAsync",
+        }
+        return cuda_launch_categories and (self.name in cuda_launch_operations or self.name in hip_launch_operations)
 
     def is_gpu_op(self) -> bool:
         """