Files
composable_kernel/test/ck_tile/memory_copy/test_copy.cpp
2025-12-23 09:45:40 +00:00

240 lines
9.3 KiB
C++

// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include <algorithm>
#include <gtest/gtest.h>
#include "ck_tile/host.hpp"
#include "ck_tile/core.hpp"
#include "ck_tile/host/kernel_launch.hpp"
#include "test_copy.hpp"
struct MemoryCopyParam
{
MemoryCopyParam(ck_tile::index_t m_, ck_tile::index_t n_, ck_tile::index_t warp_id_)
: m(m_), n(n_), warp_id(warp_id_)
{
}
ck_tile::index_t m;
ck_tile::index_t n;
ck_tile::index_t warp_id;
};
template <typename DataType, bool AsyncCopy = true>
class TestCkTileMemoryCopy : public ::testing::TestWithParam<std::tuple<int, int, int>>
{
protected:
void Run(const MemoryCopyParam& memcpy_params)
{
using XDataType = DataType;
using YDataType = DataType;
ck_tile::index_t m = memcpy_params.m;
ck_tile::index_t n = memcpy_params.n;
ck_tile::index_t warp_id = memcpy_params.warp_id;
constexpr auto dword_bytes = 4;
// if(n % (dword_bytes / sizeof(DataType)) != 0)
// {
// std::cerr << "n size should be multiple of dword_bytes" << std::endl;
// }
ck_tile::HostTensor<XDataType> x_host({m, n});
ck_tile::HostTensor<YDataType> y_host_dev({m, n});
std::cout << "input: " << x_host.mDesc << std::endl;
std::cout << "output: " << y_host_dev.mDesc << std::endl;
ck_tile::index_t value = 1;
int pack_k = 16;
for(int i = 0; i < m; i++)
{
for(int j = 0; j < n; j += pack_k)
{
// for(int k_ = 0; k_ < pack_k; k_++)
// {
// int value = m & 0x3f;
// x_host(i, j).data_[0] = rand();
// x_host(i, j).data_[1] = rand();
// x_host(i, j).data_[2] = rand();
x_host(i, j).data_[0] = i;
x_host(i, j).data_[1] = i;
x_host(i, j).data_[2] = i;
// x_host(i, j).mData[m].data_[0]=i;
// x_host(i, j).mData[m].data_[1]=i;
// x_host(i, j).mData[m].data_[2]=i;
// std::cout << i << " " << j << " " << k_ << std::endl;
// }
}
}
ck_tile::DeviceMem x_buf(x_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem y_buf(y_host_dev.get_element_space_size_in_bytes());
x_buf.ToDevice(x_host.data());
using BlockWaves = ck_tile::sequence<2, 1>;
using BlockTile = ck_tile::sequence<64, 24>;
using WaveTile = ck_tile::sequence<64, 24>;
using Vector = ck_tile::sequence<1, 24>;
ck_tile::index_t kGridSize =
ck_tile::integer_divide_ceil(m, BlockTile::at(ck_tile::number<0>{}));
using Shape = ck_tile::TileCopyShape<BlockWaves, BlockTile, WaveTile, Vector>;
using Problem = ck_tile::TileCopyProblem<int8_t, Shape, AsyncCopy>;
using Kernel = ck_tile::TileCopy<Problem>;
constexpr ck_tile::index_t kBlockSize = 128;
constexpr ck_tile::index_t kBlockPerCu = 1;
auto ms = launch_kernel(
ck_tile::stream_config{nullptr, true},
ck_tile::make_kernel<kBlockPerCu>(Kernel{},
kGridSize,
kBlockSize,
0,
static_cast<int8_t*>(x_buf.GetDeviceBuffer()),
static_cast<int8_t*>(y_buf.GetDeviceBuffer()),
m,
n / 16 * 12,
warp_id));
auto bytes = 2 * m * n * sizeof(DataType);
std::cout << "elapsed: " << ms << " (ms)" << std::endl;
std::cout << (bytes * 1e-6 / ms) << " (GB/s)" << std::endl;
// reference
y_buf.FromDevice(y_host_dev.mData.data());
for(int i = 0; i < m; i++)
{
for(int j = 0; j < n; j += pack_k)
{
// for(int k_ = 0; k_ < pack_k; k_++)
// {
// int value = m & 0x3f;
printf("%d %d\n", i, j);
std::cout << x_host(i, j).data_[0] << " " << x_host(i, j).data_[1] << " "
<< x_host(i, j).data_[2] << std::endl;
printf("==============\n");
std::cout << y_host_dev(i, j).data_[0] << " " << y_host_dev(i, j).data_[1] << " "
<< y_host_dev(i, j).data_[2] << std::endl;
printf("~~~~~~~~~~~~~\n");
}
}
// bool pass = ck_tile::check_err(y_host_dev, x_host);
// EXPECT_TRUE(pass);
}
};
// class TestCkTileMemoryCopyF6x16 : public TestCkTileMemoryCopy<ck_tile::pk_fp4_t, false>
class TestCkTileMemoryCopyF6x16 : public TestCkTileMemoryCopy<ck_tile::f6x16_pk_t, true>
{
};
class TestCkTileMemoryCopyHalfAsync : public TestCkTileMemoryCopy<ck_tile::half_t>
{
};
class TestCkTileMemoryCopyHalfSync : public TestCkTileMemoryCopy<ck_tile::half_t, false>
{
};
class TestCkTileMemoryCopyFloatAsync : public TestCkTileMemoryCopy<float>
{
};
class TestCkTileMemoryCopyFP8Async : public TestCkTileMemoryCopy<ck_tile::fp8_t>
{
};
TEST_P(TestCkTileMemoryCopyF6x16, TestCorrectness)
{
auto [M, N, warp_id] = GetParam();
this->Run({M, N, warp_id});
}
// TEST_P(TestCkTileMemoryCopyHalfAsync, TestCorrectness)
// {
// auto [M, N, warp_id] = GetParam();
// this->Run({M, N, warp_id});
// }
// TEST_P(TestCkTileMemoryCopyHalfSync, TestCorrectness)
// {
// auto [M, N, warp_id] = GetParam();
// this->Run({M, N, warp_id});
// }
// TEST_P(TestCkTileMemoryCopyFloatAsync, TestCorrectness)
// {
// auto [M, N, warp_id] = GetParam();
// this->Run({M, N, warp_id});
// }
// TEST_P(TestCkTileMemoryCopyFP8Async, TestCorrectness)
// {
// auto [M, N, warp_id] = GetParam();
// this->Run({M, N, warp_id});
// }
INSTANTIATE_TEST_SUITE_P(TestCkTileMemCopySuite,
TestCkTileMemoryCopyF6x16,
::testing::Values(std::tuple{128, 32, 0}));
// INSTANTIATE_TEST_SUITE_P(TestCkTileMemCopySuite,
// TestCkTileMemoryCopyHalfAsync,
// ::testing::Values(std::tuple{64, 8, 0},
// std::tuple{63, 8, 0},
// std::tuple{63, 2, 0},
// std::tuple{127, 30, 0},
// std::tuple{64, 8, 1},
// std::tuple{63, 8, 1},
// std::tuple{63, 2, 1},
// std::tuple{127, 30, 1},
// std::tuple{16384, 16384, 0},
// std::tuple{16384, 16384, 1}));
// INSTANTIATE_TEST_SUITE_P(TestCkTileMemCopySuite,
// TestCkTileMemoryCopyHalfSync,
// ::testing::Values(std::tuple{64, 8, 0},
// std::tuple{63, 8, 0},
// std::tuple{63, 2, 0},
// std::tuple{127, 30, 0},
// std::tuple{64, 8, 1},
// std::tuple{63, 8, 1},
// std::tuple{63, 2, 1},
// std::tuple{127, 30, 1},
// std::tuple{16384, 16384, 0},
// std::tuple{16384, 16384, 1}));
// INSTANTIATE_TEST_SUITE_P(TestCkTileMemCopySuite,
// TestCkTileMemoryCopyFloatAsync,
// ::testing::Values(std::tuple{64, 8, 0},
// std::tuple{63, 8, 0},
// std::tuple{63, 2, 0},
// std::tuple{127, 30, 0},
// std::tuple{64, 8, 1},
// std::tuple{63, 8, 1},
// std::tuple{63, 2, 1},
// std::tuple{127, 30, 1},
// std::tuple{16384, 16384, 0},
// std::tuple{16384, 16384, 1}));
// INSTANTIATE_TEST_SUITE_P(TestCkTileMemCopySuite,
// TestCkTileMemoryCopyFP8Async,
// ::testing::Values(std::tuple{64, 8, 0},
// std::tuple{63, 8, 0},
// std::tuple{63, 4, 0},
// std::tuple{127, 20, 0},
// std::tuple{64, 8, 1},
// std::tuple{63, 8, 1},
// std::tuple{63, 4, 1},
// std::tuple{127, 20, 1},
// std::tuple{16384, 16384, 0},
// std::tuple{16384, 16384, 1}));