From da1bd045abbe7a01c606b70cdb55c14795d2d5f2 Mon Sep 17 00:00:00 2001
From: xuchen-amd <xuchen@amd.com>
Date: Wed, 29 Jan 2025 13:14:57 -0500
Subject: [PATCH] Enable kokkos trace/features from rocprofv3 in
 rocprof-compute. (#545)

Signed-off-by: Xuan Chen <xuchen@amd.com>
---
 src/argparser.py                                 | 16 ++++++++++++++++
 .../profiler_rocprof_v3.py                       | 10 +++++++++-
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/argparser.py b/src/argparser.py
index d0113b223c..5342331914 100644
--- a/src/argparser.py
+++ b/src/argparser.py
@@ -139,6 +139,22 @@ Examples:
         required=False,
         help="\t\t\tSpecify the type of subpath to save workload: node_name, gpu_model.",
     )
+    profile_group.add_argument(
+        "--hip-trace",
+        dest="hip_trace",
+        required=False,
+        default=False,
+        action="store_true",
+        help="\t\t\tHIP trace, execturion trace for the entire application at the HIP level.",
+    )
+    profile_group.add_argument(
+        "--kokkos-trace",
+        dest="kokkos_trace",
+        required=False,
+        default=False,
+        action="store_true",
+        help="\t\t\tKokkos trace, traces Kokkos API calls.",
+    )
     profile_group.add_argument(
         "-k",
         "--kernel",
diff --git a/src/rocprof_compute_profile/profiler_rocprof_v3.py b/src/rocprof_compute_profile/profiler_rocprof_v3.py
index d27384cff1..b81e424537 100644
--- a/src/rocprof_compute_profile/profiler_rocprof_v3.py
+++ b/src/rocprof_compute_profile/profiler_rocprof_v3.py
@@ -42,9 +42,17 @@ class rocprof_v3_profiler(RocProfCompute_Base):
 
     def get_profiler_options(self, fname):
         app_cmd = shlex.split(self.get_args().remaining)
+        trace_option = "--kernel-trace"
         rocprof_out_format = "json"
+
         if self.get_args().format_rocprof_output == "csv":
             rocprof_out_format = "csv"
+
+        if self.get_args().hip_trace:
+            trace_option += " " + "--hip-trace"
+        if self.get_args().kokkos_trace:
+            trace_option += " " + "--kokkos-trace"
+
         args = [
             "-E",
             os.path.join(
@@ -56,7 +64,7 @@ class rocprof_v3_profiler(RocProfCompute_Base):
             # v3 requires output directory argument
             "-d",
             self.get_args().path + "/" + "out",
-            "--kernel-trace",
+            trace_option,
             "--output-format",
             rocprof_out_format,
             "--",