mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-03-21 15:47:38 +00:00
229 lines
9.3 KiB
C++
229 lines
9.3 KiB
C++
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
#include <gtest/gtest.h>
|
|
#include <vector>
|
|
#include <cmath> // For std::abs
|
|
#include <tuple>
|
|
#include <type_traits> // For std::is_same_v, std::is_floating_point_v
|
|
#include <utility> // For std::index_sequence, std::forward
|
|
|
|
#include "ck_tile/core.hpp"
|
|
#include "ck_tile/host.hpp"
|
|
#include "ck_tile/host/kernel_launch.hpp"
|
|
#include "ck_tile/ops/elementwise/kernel/elementwise_kernel.hpp"
|
|
#include "ck_tile/ops/elementwise/pipeline/elementwise_pipeline_problem.hpp"
|
|
#include "ck_tile/ops/elementwise/pipeline/elementwise_pipeline_default_policy.hpp"
|
|
#include "ck_tile/ops/elementwise/pipeline/elementwise_shape.hpp"
|
|
#include "ck_tile/ops/elementwise/binary_elementwise_operation.hpp"
|
|
#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp"
|
|
|
|
// Traits to get number of inputs for an elementwise operation
|
|
template <typename Op>
|
|
struct elementwise_op_traits;
|
|
|
|
template <>
|
|
struct elementwise_op_traits<ck_tile::element_wise::Add>
|
|
{
|
|
static constexpr int num_inputs = 2;
|
|
};
|
|
template <>
|
|
struct elementwise_op_traits<ck_tile::element_wise::Relu>
|
|
{
|
|
static constexpr int num_inputs = 1;
|
|
};
|
|
|
|
using NegRelu =
|
|
ck_tile::element_wise::Compose<ck_tile::element_wise::Relu, ck_tile::element_wise::Neg>;
|
|
template <>
|
|
struct elementwise_op_traits<NegRelu>
|
|
{
|
|
static constexpr int num_inputs = 1;
|
|
};
|
|
|
|
template <std::size_t D, typename F>
|
|
auto make_uniform_array_with_factory(F&& factory)
|
|
{
|
|
return [&]<std::size_t... Is>(std::index_sequence<Is...>) {
|
|
return std::array<std::invoke_result_t<F, std::size_t>, D>{factory(Is)...};
|
|
}(std::make_index_sequence<D>{});
|
|
}
|
|
|
|
template <typename Tuple>
|
|
class TestCkTileElementwise : public ::testing::Test
|
|
{
|
|
protected:
|
|
using XDataType = std::tuple_element_t<0, Tuple>;
|
|
using YDataType = std::tuple_element_t<1, Tuple>;
|
|
using ComputeDataType = std::tuple_element_t<2, Tuple>;
|
|
using ElementwiseOpType = std::tuple_element_t<3, Tuple>;
|
|
using BlockWarps_ = std::tuple_element_t<4, Tuple>;
|
|
using BlockTile_ = std::tuple_element_t<5, Tuple>;
|
|
using WarpTile_ = std::tuple_element_t<6, Tuple>;
|
|
using TestElementWiseShape =
|
|
ck_tile::ElementWiseShape<BlockWarps_, BlockTile_, WarpTile_, XDataType>;
|
|
static constexpr int NumInputs = elementwise_op_traits<ElementwiseOpType>::num_inputs;
|
|
|
|
void RunTest(ck_tile::index_t total_m_elements)
|
|
{
|
|
// Dims and Strides (1D example)
|
|
auto lens = ck_tile::make_tuple(total_m_elements);
|
|
auto strides = ck_tile::make_tuple(
|
|
static_cast<ck_tile::index_t>(1)); // Strides for the single dimension
|
|
|
|
// Host Tensors
|
|
auto h_xs = make_uniform_array_with_factory<NumInputs>([&](std::size_t) {
|
|
auto ret = ck_tile::HostTensor<XDataType>({total_m_elements});
|
|
ck_tile::FillUniformDistribution<XDataType>{0.f, 5.f}(ret);
|
|
return ret;
|
|
});
|
|
ck_tile::HostTensor<YDataType> h_y({total_m_elements});
|
|
h_y.SetZero();
|
|
ck_tile::HostTensor<YDataType> h_y_ref({total_m_elements});
|
|
h_y_ref.SetZero();
|
|
|
|
// Device Buffers
|
|
auto d_xs_mems_owner = make_uniform_array_with_factory<NumInputs>(
|
|
[&](std::size_t i) { return ck_tile::DeviceMem(h_xs[i]); });
|
|
for(int i = 0; i < NumInputs; ++i)
|
|
{
|
|
d_xs_mems_owner[i].ToDevice(h_xs[i].data());
|
|
}
|
|
|
|
ck_tile::DeviceMem d_y_mem(h_y);
|
|
d_y_mem.SetZero();
|
|
|
|
auto d_x_ptrs_tuple = [&]<std::size_t... Is>(std::index_sequence<Is...>) {
|
|
return ck_tile::make_tuple(
|
|
static_cast<const XDataType*>(d_xs_mems_owner[Is].GetDeviceBuffer())...);
|
|
}(std::make_index_sequence<NumInputs>{});
|
|
|
|
YDataType* p_y_device = static_cast<YDataType*>(d_y_mem.GetDeviceBuffer());
|
|
|
|
auto run_elementwise_kernel = [&](auto has_remainder) {
|
|
constexpr bool kPad = decltype(has_remainder)::value;
|
|
using Problem = ck_tile::ElementWisePipelineProblem<XDataType,
|
|
ComputeDataType,
|
|
YDataType,
|
|
TestElementWiseShape,
|
|
ElementwiseOpType,
|
|
kPad>;
|
|
using Policy = ck_tile::ElementWiseDefaultPolicy;
|
|
ck_tile::ElementWiseKernel<Problem, Policy> ew_kernel;
|
|
|
|
ck_tile::index_t grid_size = (total_m_elements + TestElementWiseShape::kBlockM - 1) /
|
|
TestElementWiseShape::kBlockM;
|
|
dim3 grid(grid_size, 1, 1);
|
|
dim3 block = dim3(ew_kernel.BlockSize());
|
|
constexpr ck_tile::index_t kBlockPerCu = 1;
|
|
|
|
ck_tile::stream_config s{nullptr, false, 0}; // Default stream, no timing, no log
|
|
|
|
ck_tile::launch_kernel(s,
|
|
ck_tile::make_kernel<kBlockPerCu> // MinBlockPerCu
|
|
(ew_kernel,
|
|
grid,
|
|
block,
|
|
0, // actual shared memory
|
|
lens,
|
|
strides, // input strides
|
|
strides, // output strides
|
|
d_x_ptrs_tuple,
|
|
p_y_device));
|
|
};
|
|
|
|
// Problem and Policy
|
|
using BaseProblem = ck_tile::ElementWisePipelineProblem<XDataType,
|
|
ComputeDataType,
|
|
YDataType,
|
|
TestElementWiseShape,
|
|
ElementwiseOpType>;
|
|
if(total_m_elements % BaseProblem::BlockShape::kVectorM)
|
|
{
|
|
run_elementwise_kernel(std::true_type{});
|
|
}
|
|
else
|
|
{
|
|
run_elementwise_kernel(std::false_type{});
|
|
}
|
|
|
|
d_y_mem.FromDevice(h_y.data());
|
|
|
|
// Reference computation on host
|
|
ElementwiseOpType op_host;
|
|
for(ck_tile::index_t i = 0; i < total_m_elements; ++i)
|
|
{
|
|
auto get_host_op_args = [&]<std::size_t... Is>(std::index_sequence<Is...>) {
|
|
return ck_tile::make_tuple(static_cast<ComputeDataType>(h_xs[Is](i))...);
|
|
}(std::make_index_sequence<NumInputs>{});
|
|
|
|
YDataType temp_y_val;
|
|
ck_tile::apply(
|
|
[&](auto&&... host_input_args) {
|
|
op_host(temp_y_val,
|
|
std::forward<decltype(host_input_args)>(host_input_args)...);
|
|
},
|
|
get_host_op_args);
|
|
h_y_ref(i) = temp_y_val;
|
|
}
|
|
|
|
// Check results
|
|
check_err(h_y, h_y_ref, "Error: Incorrect results!", 1e-5, 1e-5);
|
|
}
|
|
};
|
|
|
|
// Shape parameters (can be shared or varied per test type)
|
|
using Shape1_BlockWarps = ck_tile::sequence<1>; // 1D warp arrangement in M
|
|
using Shape1_BlockTile = ck_tile::sequence<256>; // M-dimension of block tile
|
|
using Shape1_WarpTile = ck_tile::sequence<64>; // M-dimension of warp tile
|
|
|
|
// Test configurations
|
|
using TestConfig_F32_Add = std::tuple<float,
|
|
float,
|
|
float,
|
|
ck_tile::element_wise::Add,
|
|
Shape1_BlockWarps,
|
|
Shape1_BlockTile,
|
|
Shape1_WarpTile>;
|
|
|
|
using TestConfig_F32_Relu = std::tuple<float,
|
|
float,
|
|
float,
|
|
ck_tile::element_wise::Relu,
|
|
Shape1_BlockWarps,
|
|
Shape1_BlockTile,
|
|
Shape1_WarpTile>;
|
|
|
|
using TestConfig_F16_Add = std::tuple<ck_tile::half_t,
|
|
ck_tile::half_t,
|
|
float, // Compute in float for half
|
|
ck_tile::element_wise::Add,
|
|
Shape1_BlockWarps,
|
|
Shape1_BlockTile,
|
|
Shape1_WarpTile>;
|
|
|
|
using TestConfig_F32_Neg_Relu =
|
|
std::tuple<float, float, float, NegRelu, Shape1_BlockWarps, Shape1_BlockTile, Shape1_WarpTile>;
|
|
|
|
using TestTypes = ::testing::
|
|
Types<TestConfig_F32_Add, TestConfig_F32_Relu, TestConfig_F16_Add, TestConfig_F32_Neg_Relu>;
|
|
|
|
TYPED_TEST_SUITE(TestCkTileElementwise, TestTypes);
|
|
|
|
TYPED_TEST(TestCkTileElementwise, RunElementwise_1024) { this->RunTest(1024); }
|
|
|
|
TYPED_TEST(TestCkTileElementwise, RunElementwise_513)
|
|
{
|
|
this->RunTest(513); // Test with an input size that's not a multiple of kVectorM
|
|
}
|
|
|
|
TYPED_TEST(TestCkTileElementwise, RunElementwise_516)
|
|
{
|
|
this->RunTest(516); // Test with an input size that's not a multiple of blockM
|
|
}
|
|
|
|
TYPED_TEST(TestCkTileElementwise, RunElementwise_Small_32)
|
|
{
|
|
this->RunTest(32); // Test with a very small size
|
|
}
|