[CK] s_prefetch unit test fixes.

Signed-off-by: Michal Kulikowski <Michal.Kulikowski@amd.com>


[ROCm/composable_kernel commit: cd8af997e6]

test/s_prefetch_op/s_prefetch_op.cpp

@@ -1,54 +1,44 @@
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
 
-#include <algorithm>
-#include <cstdlib>
-#include <iostream>
-#include <numeric>
-#include <tuple>
-#include <vector>
-#include <chrono>
-
 #include "ck/ck.hpp"
-#include "ck/library/utility/device_memory.hpp"
-#include "ck/library/utility/host_tensor.hpp"
-#include "ck/library/utility/check_err.hpp"
-#include "ck/host_utility/hip_check_error.hpp"
 #include "ck/host_utility/device_prop.hpp"
 
-#include <hip/hip_runtime.h>
-
-#if __clang_major__ >= 20
-#include "ck/utility/amd_buffer_addressing_builtins.hpp"
-#else
-#include "ck/utility/amd_buffer_addressing.hpp"
-#endif
-
 #include "s_prefetch_op_util.hpp"
 
-template <typename T>
-bool run_test()
+template <typename T, uint32_t NUM_THREADS, uint32_t NUM_SCALARS>
+bool run_test(bool time_kernels)
 {
     bool pass = true;
 
-    const auto s_prefetch_kernel = ck::s_prefetch_op_util::kernel_with_scalar_prefetch<T>;
-    const auto s_buffer_prefetch_kernel =
-        ck::s_prefetch_op_util::kernel_with_scalar_buffer_prefetch<T>;
+    const auto s_prefetch_kernel =
+        ck::s_prefetch_op_util::kernel_with_prefetch<T,
+                                                     NUM_THREADS,
+                                                     NUM_SCALARS,
+                                                     ck::s_prefetch_op_util::SPrefetchDataOp<T>>;
+    const auto s_buffer_prefetch_kernel = ck::s_prefetch_op_util::kernel_with_prefetch<
+        T,
+        NUM_THREADS,
+        NUM_SCALARS,
+        ck::s_prefetch_op_util::SBufferPrefetchDataOp<T, NUM_SCALARS>>;
 
     const auto prefetch_kernel_container =
         std::make_tuple(s_prefetch_kernel, s_buffer_prefetch_kernel);
 
     ck::static_for<0, 2, 1>{}([&](auto i) {
         std::string kernel_name = (i == 1 ? "s_buffer_prefetch" : "s_prefetch");
-        pass &= ck::s_prefetch_op_util::test_constant_prefetch_impl<
-            decltype(std::get<ck::Number<i>{}>(prefetch_kernel_container)),
-            T>(std::get<ck::Number<i>{}>(prefetch_kernel_container), kernel_name);
+
+        auto kernel = std::get<ck::Number<i>{}>(prefetch_kernel_container);
+
+        pass &= ck::s_prefetch_op_util::
+            test_prefetch_impl<decltype(kernel), T, NUM_THREADS, NUM_SCALARS>(
+                time_kernels, kernel, kernel_name);
     });
 
     return pass;
 }
 
-int main(int, char*[])
+int main(int argc, char* argv[])
 {
     if(!ck::is_gfx12_supported())
     {
@@ -56,13 +46,20 @@ int main(int, char*[])
         return 0;
     }
 
+    bool time_kernels = false;
+
+    if(argc == 2)
+    {
+        time_kernels = std::stoi(argv[1]);
+    }
+
     bool pass = true;
 
     std::cout << "=== Testing Constant Cache Prefetch ===" << std::endl;
 
     // Test different data types
-    pass &= run_test<float>();
-    pass &= run_test<double>();
+    pass &= run_test<float, 4096, 1024>(time_kernels);
+    pass &= run_test<double, 4096, 512>(time_kernels);
 
     std::cout << "TestConstantPrefetch ..... " << (pass ? "SUCCESS" : "FAILURE") << std::endl;
     return pass ? 0 : 1;