mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-03-24 09:07:39 +00:00
302 lines
11 KiB
C++
302 lines
11 KiB
C++
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
#include <algorithm>
|
|
#include <gtest/gtest.h>
|
|
#include <hip/hip_runtime.h>
|
|
|
|
#include "ck_tile/host.hpp"
|
|
#include "ck_tile/core.hpp"
|
|
#include "ck_tile/host/kernel_launch.hpp"
|
|
#include "test_atomic.hpp"
|
|
|
|
struct AtomicKernelParam
|
|
{
|
|
AtomicKernelParam(ck_tile::index_t m_, ck_tile::index_t n_) : m(m_), n(n_) {}
|
|
ck_tile::index_t m;
|
|
ck_tile::index_t n;
|
|
};
|
|
|
|
template <typename DataType_, ck_tile::index_t multiple_>
|
|
class TestAtomicKernel : public ::testing::TestWithParam<std::tuple<int, int>>
|
|
{
|
|
struct AtomicKernelWaveSize
|
|
{
|
|
using BlockWaves = ck_tile::sequence<2, 1>;
|
|
using BlockTile = ck_tile::sequence<128, 8>;
|
|
using WaveTile = ck_tile::sequence<64, 8>;
|
|
};
|
|
|
|
template <typename Config>
|
|
void RunTestImpl_(const AtomicKernelParam& params)
|
|
{
|
|
using XDataType = DataType_;
|
|
|
|
const ck_tile::index_t m = params.m;
|
|
const ck_tile::index_t n = params.n;
|
|
|
|
std::cout << "Input Tensor Dimensions: " << m << ", " << n << std::endl;
|
|
|
|
constexpr int dword_bytes = 4;
|
|
const int base_vec = dword_bytes / static_cast<int>(sizeof(XDataType));
|
|
const int vec = multiple_ * base_vec;
|
|
|
|
ASSERT_EQ(n % vec, 0) << " Row dimension must be divisible by vector width: n=" << n
|
|
<< " vec=" << vec << " (multiple=" << multiple_
|
|
<< ", base_vec=" << base_vec << ")";
|
|
|
|
// host tensors
|
|
ck_tile::HostTensor<XDataType> x_host_ref({m, n});
|
|
ck_tile::HostTensor<XDataType> x_host_dev({m, n});
|
|
|
|
// device buffers
|
|
ck_tile::DeviceMem x_dev_input(x_host_dev.get_element_space_size_in_bytes());
|
|
x_dev_input.SetZero();
|
|
x_host_ref.SetZero();
|
|
|
|
using BlockWaves = typename Config::BlockWaves;
|
|
using BlockTile = typename Config::BlockTile;
|
|
using WaveTile = typename Config::WaveTile;
|
|
using Vector = ck_tile::sequence<1, vec>;
|
|
|
|
// Compile-time sanity: BlockTile == WaveTile * BlockWaves
|
|
static_assert(BlockTile::at(ck_tile::number<0>{}) ==
|
|
WaveTile::at(ck_tile::number<0>{}) * BlockWaves::at(ck_tile::number<0>{}),
|
|
"BlockTile.M must equal WaveTile.M * BlockWaves.M");
|
|
static_assert(BlockTile::at(ck_tile::number<1>{}) ==
|
|
WaveTile::at(ck_tile::number<1>{}) * BlockWaves::at(ck_tile::number<1>{}),
|
|
"BlockTile.N must equal WaveTile.N * BlockWaves.N");
|
|
|
|
std::cout << "Vector per thread = " << vec
|
|
<< " BlockWaves=" << BlockWaves::at(ck_tile::number<0>{}) << "x"
|
|
<< BlockWaves::at(ck_tile::number<1>{})
|
|
<< " WaveTile=" << WaveTile::at(ck_tile::number<0>{}) << "x"
|
|
<< WaveTile::at(ck_tile::number<1>{})
|
|
<< " BlockTile=" << BlockTile::at(ck_tile::number<0>{}) << "x"
|
|
<< BlockTile::at(ck_tile::number<1>{}) << std::endl;
|
|
|
|
const ck_tile::index_t kGridSize =
|
|
ck_tile::integer_divide_ceil(m, BlockTile::at(ck_tile::number<0>{}));
|
|
|
|
using Shape = ck_tile::AtomicKernelShape<BlockWaves, BlockTile, WaveTile, Vector>;
|
|
using Problem = ck_tile::AtomicKernelProblem<XDataType, Shape>;
|
|
using Kernel = ck_tile::AtomicKernel<Problem>;
|
|
|
|
const ck_tile::index_t kBlockSize = Kernel::BlockSize();
|
|
constexpr ck_tile::index_t kBlockPerCu = 1;
|
|
|
|
(void)hipGetLastError(); // clear sticky
|
|
|
|
launch_kernel(ck_tile::stream_config{nullptr, false, 0, 0, 1},
|
|
ck_tile::make_kernel<kBlockPerCu>(
|
|
Kernel{},
|
|
kGridSize,
|
|
kBlockSize,
|
|
0,
|
|
static_cast<XDataType*>(x_dev_input.GetDeviceBuffer()),
|
|
m,
|
|
n));
|
|
|
|
ASSERT_EQ(hipPeekAtLastError(), hipSuccess)
|
|
<< "hipPeekAtLastError: " << hipGetErrorString(hipGetLastError());
|
|
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess) << "hipDeviceSynchronize failed";
|
|
|
|
// host reference computation
|
|
x_dev_input.FromDevice(x_host_dev.mData.data());
|
|
for(int i = 0; i < m; ++i)
|
|
for(int j = 0; j < n; ++j)
|
|
x_host_ref(i, j) = static_cast<XDataType>(1);
|
|
|
|
const bool pass = ck_tile::check_err(x_host_dev, x_host_ref);
|
|
EXPECT_TRUE(pass);
|
|
}
|
|
|
|
protected:
|
|
void RunTest(const AtomicKernelParam& params) { RunTestImpl_<AtomicKernelWaveSize>(params); }
|
|
};
|
|
|
|
class TestAtomicKernelHalf_1 : public TestAtomicKernel<ck_tile::half_t, 1>
|
|
{
|
|
};
|
|
class TestAtomicKernelHalf_2 : public TestAtomicKernel<ck_tile::half_t, 2>
|
|
{
|
|
};
|
|
class TestAtomicKernelHalf_4 : public TestAtomicKernel<ck_tile::half_t, 4>
|
|
{
|
|
};
|
|
class TestAtomicKernelBF16_1 : public TestAtomicKernel<ck_tile::bf16_t, 1>
|
|
{
|
|
};
|
|
class TestAtomicKernelBF16_2 : public TestAtomicKernel<ck_tile::bf16_t, 2>
|
|
{
|
|
};
|
|
class TestAtomicKernelBF16_4 : public TestAtomicKernel<ck_tile::bf16_t, 4>
|
|
{
|
|
};
|
|
class TestAtomicKernelBF8_1 : public TestAtomicKernel<ck_tile::bf8_t, 1>
|
|
{
|
|
};
|
|
class TestAtomicKernelBF8_2 : public TestAtomicKernel<ck_tile::bf8_t, 2>
|
|
{
|
|
};
|
|
class TestAtomicKernelFP8_1 : public TestAtomicKernel<ck_tile::fp8_t, 1>
|
|
{
|
|
};
|
|
class TestAtomicKernelFP8_2 : public TestAtomicKernel<ck_tile::fp8_t, 2>
|
|
{
|
|
};
|
|
class TestAtomicKernelFloat_1 : public TestAtomicKernel<float, 1>
|
|
{
|
|
};
|
|
class TestAtomicKernelFloat_2 : public TestAtomicKernel<float, 2>
|
|
{
|
|
};
|
|
class TestAtomicKernelFloat_4 : public TestAtomicKernel<float, 4>
|
|
{
|
|
};
|
|
|
|
TEST_P(TestAtomicKernelHalf_1, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelHalf_2, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelHalf_4, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelBF16_1, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelBF16_2, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelBF16_4, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelBF8_1, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelBF8_2, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelFP8_1, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelFP8_2, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelFloat_1, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelFloat_2, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
TEST_P(TestAtomicKernelFloat_4, TestCorrectness)
|
|
{
|
|
auto [M, N] = GetParam();
|
|
this->RunTest({M, N});
|
|
}
|
|
|
|
// Common parameter lists
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelHalf_1,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelHalf_2,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelHalf_4,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelBF16_1,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelBF16_2,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelBF16_4,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelBF8_1,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelBF8_2,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelFP8_1,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelFP8_2,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelFloat_1,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelFloat_2,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|
|
|
|
INSTANTIATE_TEST_SUITE_P(TestAtomicKernelSuite,
|
|
TestAtomicKernelFloat_4,
|
|
::testing::Values(std::tuple{64, 8},
|
|
std::tuple{64, 16},
|
|
std::tuple{64, 32}));
|