From 780456f1ce4da8aa2c00defa85d5ac0d5c0c48ce Mon Sep 17 00:00:00 2001
From: Hosang <156028780+hyoon1@users.noreply.github.com>
Date: Mon, 29 Sep 2025 18:34:04 -0400
Subject: [PATCH] Fix timing issue in CK_TILE GEMM example (#2940)

[ROCm/composable_kernel commit: 28ad8ae5d8558e147f29aba29db569fe25210947]
---
 example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp | 4 +++-
 example/ck_tile/03_gemm/universal_gemm_invoker.hpp        | 4 +++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp b/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp
index f0a71670ad..8c7589dabb 100644
--- a/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp
+++ b/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp
@@ -231,7 +231,7 @@ struct SplitKTwoStageInvoker
                 preprocess = clear_gemm_output;
             }
 
-            return ck_tile::launch_kernel_time_mask(
+            ave_time = ck_tile::launch_kernel_time_mask(
                 s,
                 preprocess,
                 ck_tile::make_kernel<GemmConfig::kBlockPerCu>(
@@ -245,6 +245,8 @@ struct SplitKTwoStageInvoker
                                                   ck_tile::make_tuple(args.N, 1), // Output Stride
                                                   input_tensors,
                                                   static_cast<CDataType*>(c_ptr)));
+
+            return ave_time;
         };
 
         const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
diff --git a/example/ck_tile/03_gemm/universal_gemm_invoker.hpp b/example/ck_tile/03_gemm/universal_gemm_invoker.hpp
index 19855c7f72..e0d97a50db 100644
--- a/example/ck_tile/03_gemm/universal_gemm_invoker.hpp
+++ b/example/ck_tile/03_gemm/universal_gemm_invoker.hpp
@@ -174,10 +174,12 @@ struct UniversalInvoker
                 preprocess = clear_gemm_output;
             }
 
-            return ck_tile::launch_kernel_time_mask(
+            ave_time = ck_tile::launch_kernel_time_mask(
                 s,
                 preprocess,
                 ck_tile::make_kernel<GemmConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
+
+            return ave_time;
         };
 
         const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {