From df5204ca629ac872e72a2e4f7367aaa1631fc75a Mon Sep 17 00:00:00 2001
From: spandoesai <spandan.more@amd.com>
Date: Tue, 3 Dec 2024 01:52:12 +0000
Subject: [PATCH] Added Kernel Launch Op names for AMD GPUs

---
 src/trace_link/kineto_operator.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/trace_link/kineto_operator.py b/src/trace_link/kineto_operator.py
index f074f063..f286fa6f 100644
--- 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:
         """