Skip to content

Commit 452f228

Browse files
tenpercenttenpercent
committed
add gtest
1 parent b056203 commit 452f228

3 files changed

Lines changed: 280 additions & 0 deletions

File tree

‎test/ck_tile/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
add_subdirectory(image_to_column)
55
add_subdirectory(gemm)
6+
add_subdirectory(gemm_persistent_async_input)
67
add_subdirectory(gemm_weight_preshuffle)
78
add_subdirectory(batched_gemm)
89
add_subdirectory(grouped_gemm)

‎test/ck_tile/gemm_persistent_async_input/CMakeLists.txt‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
2+
# SPDX-License-Identifier: MIT
3+
4+
# Test for persistent async input GEMM - currently targeting gfx95
5+
set(PERSISTENT_ASYNC_INPUT_COMPILE_OPTIONS)
6+
if(CK_USE_OCP_FP8)
7+
list(APPEND PERSISTENT_ASYNC_INPUT_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8)
8+
endif()
9+
list(APPEND PERSISTENT_ASYNC_INPUT_COMPILE_OPTIONS
10+
-mllvm
11+
-enable-noalias-to-md-conversion=0
12+
)
13+
14+
if(GPU_TARGETS MATCHES "gfx95")
15+
add_gtest_executable(test_ck_tile_gemm_persistent_async_input test_gemm_persistent_async_input.cpp)
16+
target_compile_options(test_ck_tile_gemm_persistent_async_input PRIVATE ${PERSISTENT_ASYNC_INPUT_COMPILE_OPTIONS})
17+
else()
18+
message(DEBUG "Skipping test_ck_tile_gemm_persistent_async_input for current target - requires gfx95")
19+
endif()

‎test/ck_tile/gemm_persistent_async_input/test_gemm_persistent_async_input.cpp‎

Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
1+
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
2+
// SPDX-License-Identifier: MIT
3+
4+
#include "gtest/gtest.h"
5+
#include "ck_tile/host.hpp"
6+
#include "ck_tile/core.hpp"
7+
#include "ck_tile/ops/gemm.hpp"
8+
#include "ck_tile/ops/epilogue.hpp"
9+
#include "ck_tile/host/kernel_launch.hpp"
10+
#include "ck_tile/core/utility/persistent_async_input_scheduler.hpp"
11+
12+
using Row = ck_tile::tensor_layout::gemm::RowMajor;
13+
using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
14+
using F16 = ck_tile::fp16_t;
15+
using F32 = ck_tile::fp32_t;
16+
using Intrawave = ck_tile::integral_constant<ck_tile::GemmPipelineScheduler,
17+
ck_tile::GemmPipelineScheduler::Intrawave>;
18+
19+
template <typename ALayout,
20+
typename BLayout,
21+
typename CLayout,
22+
typename ADataType,
23+
typename BDataType,
24+
typename AccDataType,
25+
typename CDataType>
26+
class TestGemmPersistentAsyncInput : public ::testing::Test
27+
{
28+
protected:
29+
static constexpr ck_tile::index_t M = 512;
30+
static constexpr ck_tile::index_t N = 1024;
31+
static constexpr ck_tile::index_t K = 512;
32+
33+
static constexpr ck_tile::index_t M_Tile = 256;
34+
static constexpr ck_tile::index_t N_Tile = 256;
35+
static constexpr ck_tile::index_t K_Tile = 32;
36+
37+
static constexpr ck_tile::index_t M_Warp_Tile = 32;
38+
static constexpr ck_tile::index_t N_Warp_Tile = 32;
39+
static constexpr ck_tile::index_t K_Warp_Tile = 16;
40+
41+
static constexpr ck_tile::index_t M_Warp = 2;
42+
static constexpr ck_tile::index_t N_Warp = 2;
43+
static constexpr ck_tile::index_t K_Warp = 1;
44+
45+
template <bool IsRowMajor>
46+
static constexpr ck_tile::index_t get_default_stride(ck_tile::index_t row,
47+
ck_tile::index_t col)
48+
{
49+
if constexpr(IsRowMajor)
50+
return col;
51+
else
52+
return row;
53+
}
54+
55+
void Run()
56+
{
57+
constexpr bool is_a_row_major = std::is_same_v<ALayout, Row>;
58+
constexpr bool is_b_row_major = std::is_same_v<BLayout, Row>;
59+
constexpr bool is_c_row_major = std::is_same_v<CLayout, Row>;
60+
61+
ck_tile::index_t stride_A = get_default_stride<is_a_row_major>(M, K);
62+
ck_tile::index_t stride_B = get_default_stride<is_b_row_major>(K, N);
63+
ck_tile::index_t stride_C = get_default_stride<is_c_row_major>(M, N);
64+
65+
ck_tile::HostTensor<ADataType> a_m_k(
66+
ck_tile::host_tensor_descriptor(M, K, stride_A, ck_tile::bool_constant<is_a_row_major>{}));
67+
ck_tile::HostTensor<BDataType> b_k_n(
68+
ck_tile::host_tensor_descriptor(K, N, stride_B, ck_tile::bool_constant<is_b_row_major>{}));
69+
ck_tile::HostTensor<CDataType> c_m_n_dev_result(
70+
ck_tile::host_tensor_descriptor(M, N, stride_C, ck_tile::bool_constant<is_c_row_major>{}));
71+
ck_tile::HostTensor<CDataType> c_m_n_host_ref(
72+
ck_tile::host_tensor_descriptor(M, N, stride_C, ck_tile::bool_constant<is_c_row_major>{}));
73+
74+
// Fill input tensors with random values
75+
ck_tile::FillUniformDistributionIntegerValue<ADataType>{-5, 5, 11939}(a_m_k);
76+
ck_tile::FillUniformDistributionIntegerValue<BDataType>{-5, 5, 11940}(b_k_n);
77+
78+
// Allocate device memory
79+
ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes());
80+
ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes());
81+
ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes());
82+
83+
// Copy input data to device
84+
a_m_k_dev_buf.ToDevice(a_m_k.data());
85+
b_k_n_dev_buf.ToDevice(b_k_n.data());
86+
c_m_n_dev_buf.SetZero();
87+
c_m_n_dev_result.SetZero();
88+
c_m_n_host_ref.SetZero();
89+
90+
// Compute reference result on host
91+
ck_tile::reference_gemm<ADataType, BDataType, AccDataType, CDataType>(
92+
a_m_k, b_k_n, c_m_n_host_ref);
93+
94+
// Setup kernel configuration for persistent async input GEMM
95+
constexpr int kBlockPerCu = 1;
96+
constexpr bool kPadM = true;
97+
constexpr bool kPadN = true;
98+
constexpr bool kPadK = true;
99+
constexpr bool DoubleSmemBuffer = true;
100+
constexpr bool TransposeC = false;
101+
constexpr bool StructuredSparsity = false;
102+
constexpr bool Persistent = true;
103+
constexpr int NumWaveGroup = 1;
104+
constexpr bool Preshuffle = false;
105+
constexpr ck_tile::index_t TilePartitionerGroupNum = 8;
106+
constexpr ck_tile::index_t TilePartitionerM01 = 4;
107+
108+
using GemmShape = ck_tile::TileGemmShape<
109+
ck_tile::sequence<M_Tile, N_Tile, K_Tile>,
110+
ck_tile::sequence<M_Warp, N_Warp, K_Warp>,
111+
ck_tile::sequence<M_Warp_Tile, N_Warp_Tile, K_Warp_Tile>>;
112+
113+
using TilePartitioner = ck_tile::GemmSpatiallyLocalTilePartitioner<
114+
GemmShape, TilePartitionerGroupNum, TilePartitionerM01>;
115+
116+
using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits<
117+
kPadM, kPadN, kPadK, DoubleSmemBuffer,
118+
ALayout, BLayout, CLayout,
119+
TransposeC, StructuredSparsity, Persistent, NumWaveGroup, Preshuffle>;
120+
121+
using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<
122+
ADataType, BDataType, AccDataType,
123+
GemmShape, GemmUniversalTraits, Intrawave::value>;
124+
125+
using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompAsync<UniversalGemmProblem>;
126+
127+
using DsLayout = ck_tile::tuple<>;
128+
using DsDataType = ck_tile::tuple<>;
129+
130+
using GemmEpilogue = ck_tile::CShuffleEpilogue<
131+
ck_tile::CShuffleEpilogueProblem<
132+
ADataType, BDataType, DsDataType, AccDataType, CDataType,
133+
DsLayout, CLayout,
134+
ck_tile::element_wise::PassThrough,
135+
TilePartitioner::MPerBlock, TilePartitioner::NPerBlock,
136+
M_Warp, N_Warp, M_Warp_Tile, N_Warp_Tile, K_Warp_Tile,
137+
UniversalGemmProblem::TransposeC,
138+
1, // kNumWaveGroups_
139+
false, // FixedVectorSize_
140+
1, // VectorSizeC_
141+
false, // TiledMMAPermuteN_
142+
1, // BlockedXDLN_PerWarp_
143+
DoubleSmemBuffer>>;
144+
145+
using Kernel = ck_tile::GemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
146+
147+
// Calculate tiles and chunks for async scheduler
148+
constexpr ck_tile::index_t tiles_per_chunk = 4;
149+
constexpr ck_tile::index_t tile_idx_pivot = 2; // First 2 tiles don't wait
150+
151+
const ck_tile::index_t tiles_m = ck_tile::integer_divide_ceil(M, M_Tile);
152+
const ck_tile::index_t tiles_needing_signals =
153+
(tiles_m > tile_idx_pivot) ? (tiles_m - tile_idx_pivot) : 0;
154+
const ck_tile::index_t num_chunks =
155+
ck_tile::integer_divide_ceil(tiles_needing_signals, tiles_per_chunk);
156+
157+
// Allocate chunk signals (initialized to zero)
158+
ck_tile::DeviceMem signal_buf(num_chunks * sizeof(uint32_t));
159+
signal_buf.SetZero();
160+
uint32_t* d_chunk_signals = static_cast<uint32_t*>(signal_buf.GetDeviceBuffer());
161+
162+
// Setup async input scheduler
163+
ck_tile::PersistentAsyncInputScheduler async_scheduler;
164+
async_scheduler.tiles_per_chunk_m = tiles_per_chunk;
165+
async_scheduler.chunk_signals = d_chunk_signals;
166+
async_scheduler.tile_idx_pivot_m = tile_idx_pivot;
167+
168+
// Create UniversalGemmHostArgs with async scheduler
169+
ck_tile::UniversalGemmHostArgs<1, 1, 0> host_args(
170+
{a_m_k_dev_buf.GetDeviceBuffer()},
171+
{b_k_n_dev_buf.GetDeviceBuffer()},
172+
{},
173+
c_m_n_dev_buf.GetDeviceBuffer(),
174+
1, // k_batch
175+
M, N, K,
176+
{stride_A},
177+
{stride_B},
178+
{},
179+
stride_C,
180+
async_scheduler);
181+
182+
// Create kernel args using UniversalGemmKernel
183+
auto kargs = Kernel::UniversalGemmKernel::MakeKernelArgs(host_args);
184+
185+
// Setup grid and blocks for persistent kernel
186+
ck_tile::stream_config stream_cfg{nullptr, false};
187+
const dim3 grids = Kernel::MaxOccupancyGridSize(stream_cfg);
188+
const dim3 blocks = Kernel::BlockSize();
189+
190+
if(!Kernel::IsSupportedArgument(kargs))
191+
{
192+
GTEST_SKIP() << "Kernel arguments not supported, skipping test";
193+
return;
194+
}
195+
196+
// Launch kernel
197+
ck_tile::ignore = ck_tile::launch_kernel(
198+
stream_cfg,
199+
ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
200+
201+
// Simulate setting chunk signals (would normally be done by producer)
202+
// For testing, we just set all signals to allow kernel completion
203+
std::vector<uint32_t> host_signals(num_chunks, 1);
204+
signal_buf.ToDevice(host_signals.data());
205+
206+
// Wait for kernel completion
207+
HIP_CHECK_ERROR(hipDeviceSynchronize());
208+
209+
// Copy result back to host
210+
c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data());
211+
212+
// Validate results
213+
const float max_accumulated_value =
214+
*std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end());
215+
216+
const auto rtol = ck_tile::get_relative_threshold<ADataType, CDataType, AccDataType>(K);
217+
const auto atol =
218+
ck_tile::get_absolute_threshold<ADataType, CDataType, AccDataType>(
219+
max_accumulated_value, K);
220+
221+
bool pass = ck_tile::check_err(
222+
c_m_n_dev_result, c_m_n_host_ref, "Error: Incorrect results!", rtol, atol);
223+
224+
EXPECT_TRUE(pass);
225+
}
226+
};
227+
228+
// Define test types for different layout combinations
229+
using RowRowRow_F16F16F32F16 =
230+
TestGemmPersistentAsyncInput<Row, Row, Row, F16, F16, F32, F16>;
231+
using RowColRow_F16F16F32F16 =
232+
TestGemmPersistentAsyncInput<Row, Col, Row, F16, F16, F32, F16>;
233+
using ColRowRow_F16F16F32F16 =
234+
TestGemmPersistentAsyncInput<Col, Row, Row, F16, F16, F32, F16>;
235+
using ColColRow_F16F16F32F16 =
236+
TestGemmPersistentAsyncInput<Col, Col, Row, F16, F16, F32, F16>;
237+
238+
// Test case for Row-Row-Row layout
239+
TEST_F(RowRowRow_F16F16F32F16, BasicTest)
240+
{
241+
this->Run();
242+
}
243+
244+
// Test case for Row-Col-Row layout
245+
TEST_F(RowColRow_F16F16F32F16, BasicTest)
246+
{
247+
this->Run();
248+
}
249+
250+
// Test case for Col-Row-Row layout
251+
TEST_F(ColRowRow_F16F16F32F16, BasicTest)
252+
{
253+
this->Run();
254+
}
255+
256+
// Test case for Col-Col-Row layout
257+
TEST_F(ColColRow_F16F16F32F16, BasicTest)
258+
{
259+
this->Run();
260+
}

0 commit comments

Comments
 (0)