diff --git a/src/trace_link/kineto_operator.py b/src/trace_link/kineto_operator.py index 5078aeba..76858a66 100644 --- 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: """