From 1393082dca0ff1a824a3b7979161fb33a0c8af11 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Augonnet?=
 <158148890+caugonnet@users.noreply.github.com>
Date: Sun, 15 Dec 2024 08:03:37 +0100
Subject: [PATCH] Implement CUDASTF_DOT_TIMING for the host_launch construct
 (#3170)

---
 .../__stf/internal/backend_ctx.cuh            | 42 ++++++++++++++++++-
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh b/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh
index 495c6d435b4..dce51b4ed32 100644
--- a/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh
+++ b/cudax/include/cuda/experimental/__stf/internal/backend_ctx.cuh
@@ -120,6 +120,9 @@ public:
   template <typename Fun>
   void operator->*(Fun&& f)
   {
+    auto& dot        = *ctx.get_dot();
+    auto& statistics = reserved::task_statistics::instance();
+
     auto t = ctx.task(exec_place::host);
     t.add_deps(deps);
     if (!symbol.empty())
@@ -127,13 +130,48 @@ public:
       t.set_symbol(symbol);
     }
 
+    cudaEvent_t start_event, end_event;
+    const bool record_time = t.schedule_task() || statistics.is_calibrating_to_file();
+
     t.start();
+
+    if constexpr (::std::is_same_v<Ctx, stream_ctx>)
+    {
+      if (record_time)
+      {
+        cuda_safe_call(cudaEventCreate(&start_event));
+        cuda_safe_call(cudaEventCreate(&end_event));
+        cuda_safe_call(cudaEventRecord(start_event, t.get_stream()));
+      }
+    }
+
     SCOPE(exit)
     {
-      t.end();
+      t.end_uncleared();
+      if constexpr (::std::is_same_v<Ctx, stream_ctx>)
+      {
+        if (record_time)
+        {
+          cuda_safe_call(cudaEventRecord(end_event, t.get_stream()));
+          cuda_safe_call(cudaEventSynchronize(end_event));
+
+          float milliseconds = 0;
+          cuda_safe_call(cudaEventElapsedTime(&milliseconds, start_event, end_event));
+
+          if (dot.is_tracing())
+          {
+            dot.template add_vertex_timing<typename Ctx::task_type>(t, milliseconds, -1);
+          }
+
+          if (statistics.is_calibrating())
+          {
+            statistics.log_task_time(t, milliseconds);
+          }
+        }
+      }
+      t.clear();
     };
 
-    auto& dot = *ctx.get_dot();
     if (dot.is_tracing())
     {
       dot.template add_vertex<typename Ctx::task_type, logical_data_untyped>(t);