From 43efeb67e633cbacc8b2dbdcc84259ff56544d80 Mon Sep 17 00:00:00 2001
From: Kevin Abraham <kevin.abraham@streamhpc.com>
Date: Thu, 11 Dec 2025 07:15:27 +0000
Subject: [PATCH] add interface to device_mem for tests

---
 include/ck/library/utility/device_memory.hpp | 63 ++++++++++++++++++++
 1 file changed, 63 insertions(+)

diff --git a/include/ck/library/utility/device_memory.hpp b/include/ck/library/utility/device_memory.hpp
index b0ee766ff5..889bdb2829 100644
--- a/include/ck/library/utility/device_memory.hpp
+++ b/include/ck/library/utility/device_memory.hpp
@@ -4,6 +4,8 @@
 #pragma once
 
 #include <hip/hip_runtime.h>
+#include <math.h>
+#include "ck/library/utility/device_tensor_generator.hpp"
 
 namespace ck {
 
@@ -34,6 +36,12 @@ struct DeviceMem
     void SetZero() const;
     template <typename T>
     void SetValue(T x) const;
+    template <typename T>
+    void FillUniformRandInteger(int min_value, int max_value);
+    template <typename T>
+    void FillUniformRandFp(float min_value, float max_value);
+    template <typename T>
+    void FillNormalRandFp(float sigma, float mean);
     ~DeviceMem();
 
     void* mpDeviceBuf;
@@ -51,4 +59,59 @@ void DeviceMem::SetValue(T x) const
     set_buffer_value<T><<<1, 1024>>>(static_cast<T*>(mpDeviceBuf), x, mMemSize / sizeof(T));
 }
 
+template <typename T>
+void DeviceMem::FillUniformRandInteger(int min_value, int max_value)
+{
+    if(mMemSize % sizeof(T) != 0)
+    {
+        throw std::runtime_error("wrong! not entire DeviceMem will be filled");
+    }
+    if(max_value - min_value <= 1)
+    {
+        throw std::runtime_error("Error while filling device tensor with random integer data: max "
+                                 "value must be at least 2 greater than min value, otherwise "
+                                 "tensor will be filled by a constant value (end is exclusive)");
+    }
+    if(max_value - 1 == min_value || max_value - 1 == max_value)
+    {
+        throw std::runtime_error("Error while filling device tensor with random integer data: "
+                                 "insufficient precision in specified range");
+    }
+
+    size_t packed_size = packed_size_v<T>;
+    fill_tensor_uniform_rand_int_values<<<256, 256>>>(
+        static_cast<T*>(mpDeviceBuf), min_value, max_value, (mMemSize * packed_size) / sizeof(T));
+}
+
+template <typename T>
+void DeviceMem::FillUniformRandFp(float min_value, float max_value)
+{
+    if(mMemSize % sizeof(T) != 0)
+    {
+        throw std::runtime_error("wrong! not entire DeviceMem will be filled");
+    }
+
+    if(std::nextafter(min_value, max_value) >= max_value)
+    {
+        // nextafter is not defined for non-standard data types, skip test for now
+        // throw std::runtime_error("Error while filling device tensor with random integer data:
+        // insufficient precision in specified range");
+    }
+
+    size_t packed_size = packed_size_v<T>;
+    fill_tensor_uniform_rand_fp_values<<<256, 256>>>(
+        static_cast<T*>(mpDeviceBuf), min_value, max_value, (mMemSize * packed_size) / sizeof(T));
+}
+
+template <typename T>
+void DeviceMem::FillNormalRandFp(float sigma, float mean)
+{
+    if(mMemSize % sizeof(T) != 0)
+    {
+        throw std::runtime_error("wrong! not entire DeviceMem will be filled");
+    }
+
+    fill_tensor_norm_rand_fp_values<<<256, 256>>>(
+        static_cast<T*>(mpDeviceBuf), sigma, mean, mMemSize / sizeof(T));
+}
 } // namespace ck