add interface to device_mem for tests

This commit is contained in:
Kevin Abraham
2025-12-11 07:15:27 +00:00
parent 60a476420d
commit 43efeb67e6

View File

@@ -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