Added Kernel Launch Op names for AMD GPUs

mlcommons · Dec 3, 2024 · df5204c · df5204c
1 parent b915ab8
commit df5204c
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
@@ -147,7 +147,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:
         """