From 09e10ad38fdcdcba065f2e5871a219fa9a9f0067 Mon Sep 17 00:00:00 2001
From: zejun <zejun.chen@intel.com>
Date: Wed, 22 Jan 2025 09:37:18 -0800
Subject: [PATCH] Add following eligible unified runtime ops for drawing flow
 in traced JSON file (#1024)

Summary:
Intel GPU toolchain is using the unified runtime for now. As the profiling tools XPTI, it should be compatible with the unified runtime and previous SYCL runtime, so this PR adds the following eligible unified runtime ops for drawing flow in traced JSON file when using unified runtime.
```
urEnqueueUSMFill
urEnqueueUSMFill2D
urEnqueueUSMMemcpy
urEnqueueUSMMemcpy2D
urEnqueueKernelLaunch
urEnqueueKernelLaunchCustomExp
urEnqueueCooperativeKernelLaunchExp
```

Pull Request resolved: https://github.com/pytorch/kineto/pull/1024

Reviewed By: aaronenyeshi

Differential Revision: D68457492

Pulled By: sraikund16

fbshipit-source-id: fae9357c9c8c8b3793dbeb0104c72787127db3b0
---
 libkineto/src/plugin/xpupti/XpuptiActivityProfiler.cpp | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/libkineto/src/plugin/xpupti/XpuptiActivityProfiler.cpp b/libkineto/src/plugin/xpupti/XpuptiActivityProfiler.cpp
index e722e49ae..b2a56c6fa 100644
--- a/libkineto/src/plugin/xpupti/XpuptiActivityProfiler.cpp
+++ b/libkineto/src/plugin/xpupti/XpuptiActivityProfiler.cpp
@@ -10,14 +10,21 @@ std::vector<std::array<unsigned char, 16>>
     XpuptiActivityProfilerSession::deviceUUIDs_ = {};
 std::vector<std::string> XpuptiActivityProfilerSession::correlateRuntimeOps_ = {
     "piextUSMEnqueueFill",
+    "urEnqueueUSMFill",
     "piextUSMEnqueueFill2D",
+    "urEnqueueUSMFill2D",
     "piextUSMEnqueueMemcpy",
+    "urEnqueueUSMMemcpy",
     "piextUSMEnqueueMemset",
     "piextUSMEnqueueMemcpy2D",
+    "urEnqueueUSMMemcpy2D",
     "piextUSMEnqueueMemset2D",
     "piEnqueueKernelLaunch",
+    "urEnqueueKernelLaunch",
     "piextEnqueueKernelLaunchCustom",
-    "piextEnqueueCooperativeKernelLaunch"};
+    "urEnqueueKernelLaunchCustomExp",
+    "piextEnqueueCooperativeKernelLaunch",
+    "urEnqueueCooperativeKernelLaunchExp"};
 
 // =========== Session Constructor ============= //
 XpuptiActivityProfilerSession::XpuptiActivityProfilerSession(