Skip to content

Commit 6df2d70

Browse files
yungshengtuyungshengtu
authored
Implement device_gemm_universal_preshuffle_instance for RDNA4 (#3429)
* add device_gemm_wmma_cshuffle_v3_b_preshuffle.hpp * add examples * add instances to test * remove duplicate code between examples
1 parent e302079 commit 6df2d70

20 files changed

Lines changed: 1229 additions & 14 deletions

‎example/01_gemm/CMakeLists.txt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,3 +149,7 @@ add_example_executable(example_gemm_wmma_fp16_fp8_v3 gemm_wmma_fp16_fp8_v3.cpp)
149149
add_example_dependencies(example_gemm_wmma example_gemm_wmma_fp16_fp8_v3)
150150
add_example_executable(example_gemm_wmma_fp16_pk_i4_v3_b_scale gemm_wmma_fp16_pk_i4_v3_b_scale.cpp)
151151
add_example_dependencies(example_gemm_wmma example_gemm_wmma_fp16_pk_i4_v3_b_scale)
152+
add_example_executable(example_gemm_wmma_fp8_bpreshuffle gemm_wmma_fp8_bpreshuffle.cpp)
153+
add_example_dependencies(example_gemm_wmma example_gemm_wmma_fp8_bpreshuffle)
154+
add_example_executable(example_gemm_wmma_fp16_bpreshuffle gemm_wmma_fp16_bpreshuffle.cpp)
155+
add_example_dependencies(example_gemm_wmma example_gemm_wmma_fp16_bpreshuffle)

‎example/01_gemm/gemm_wmma_fp16_bpreshuffle.cpp‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
2+
// SPDX-License-Identifier: MIT
3+
4+
#include "common.hpp"
5+
6+
#include "ck/ck.hpp"
7+
#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
8+
#include "ck/library/utility/check_err.hpp"
9+
#include "ck/library/utility/device_memory.hpp"
10+
#include "ck/library/utility/host_tensor.hpp"
11+
#include "ck/library/utility/host_tensor_generator.hpp"
12+
#include "ck/library/utility/literals.hpp"
13+
#include "ck/stream_config.hpp"
14+
#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
15+
#include "ck/tensor_operation/gpu/device/impl/device_gemm_wmma_cshuffle_v3_b_preshuffle.hpp"
16+
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
17+
#include "ck/utility/data_type.hpp"
18+
#include "ck/utility/get_id.hpp"
19+
#include "ck/utility/scheduler_enum.hpp"
20+
21+
#include <cstddef>
22+
#include <iostream>
23+
#include <type_traits>
24+
25+
using F16 = ck::half_t;
26+
using F32 = float;
27+
28+
using ADataType = F16;
29+
using BDataType = F16;
30+
using AccDataType = F32;
31+
using CShuffleDataType = F32;
32+
using CDataType = F16;
33+
using ComputeTypeA = F16;
34+
using ComputeTypeB = F16;
35+
36+
using ALayout = Row;
37+
using BLayout = Col;
38+
using CLayout = Row;
39+
40+
using AElementOp = PassThrough;
41+
using BElementOp = PassThrough;
42+
using CElementOp = PassThrough;
43+
44+
static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
45+
46+
static constexpr bool PermuteA = false;
47+
static constexpr bool PermuteB = false;
48+
static constexpr int KPack = 8; // int4 -> 32, fp8 -> 16, fp16 -> 8
49+
// clang-format off
50+
using DeviceOpInstance =
51+
ck::tensor_operation::device::DeviceGemm_Wmma_CShuffleV3_BPreshuffle<
52+
ALayout, BLayout, CLayout,
53+
ADataType, BDataType, CDataType, AccDataType, CShuffleDataType,
54+
AElementOp, BElementOp, CElementOp, GemmDefault,
55+
128,
56+
32, 128, 128,
57+
8, 8,
58+
16, 16,
59+
2, 2,
60+
S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>,
61+
2, 8, 8, 0,
62+
S<16, 8, 1>, S<1, 0, 2>, S<1, 0, 2>,
63+
2, 8, 8, 0,
64+
1, 1, S<1, 16, 1, 8>, S<4, 4, 1>,
65+
ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, ComputeTypeA, ComputeTypeB, PermuteA, PermuteB>;
66+
// clang-format on
67+
68+
#include "run_gemm_wmma_bpreshuffle_example.inc"
69+
70+
int main(int argc, char* argv[]) { return !run_gemm_splitk_example(argc, argv); }

‎example/01_gemm/gemm_wmma_fp8_bpreshuffle.cpp‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
2+
// SPDX-License-Identifier: MIT
3+
4+
#include "common.hpp"
5+
6+
#include "ck/ck.hpp"
7+
#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
8+
#include "ck/library/utility/check_err.hpp"
9+
#include "ck/library/utility/device_memory.hpp"
10+
#include "ck/library/utility/host_tensor.hpp"
11+
#include "ck/library/utility/host_tensor_generator.hpp"
12+
#include "ck/library/utility/literals.hpp"
13+
#include "ck/stream_config.hpp"
14+
#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
15+
#include "ck/tensor_operation/gpu/device/impl/device_gemm_wmma_cshuffle_v3_b_preshuffle.hpp"
16+
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
17+
#include "ck/utility/amd_ck_fp8.hpp"
18+
#include "ck/utility/data_type.hpp"
19+
#include "ck/utility/get_id.hpp"
20+
#include "ck/utility/scheduler_enum.hpp"
21+
22+
#include <cstddef>
23+
#include <iostream>
24+
#include <type_traits>
25+
26+
using F8 = ck::f8_t;
27+
using F16 = ck::half_t;
28+
using F32 = float;
29+
30+
using ADataType = F8;
31+
using BDataType = F8;
32+
using AccDataType = F32;
33+
using CShuffleDataType = F32;
34+
using CDataType = F16;
35+
using ComputeTypeA = F8;
36+
using ComputeTypeB = F8;
37+
38+
using ALayout = Row;
39+
using BLayout = Col;
40+
using CLayout = Row;
41+
42+
using AElementOp = PassThrough;
43+
using BElementOp = PassThrough;
44+
using CElementOp = PassThrough;
45+
46+
static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
47+
48+
static constexpr bool PermuteA = false;
49+
static constexpr bool PermuteB = false;
50+
static constexpr int KPack = 16; // int4 -> 32, fp8 -> 16, fp16 -> 8
51+
// clang-format off
52+
using DeviceOpInstance =
53+
ck::tensor_operation::device::DeviceGemm_Wmma_CShuffleV3_BPreshuffle<
54+
ALayout, BLayout, CLayout,
55+
ADataType, BDataType, CDataType, AccDataType, CShuffleDataType,
56+
AElementOp, BElementOp, CElementOp, GemmDefault,
57+
256,
58+
32, 128, 256,
59+
16, 16,
60+
16, 16,
61+
2, 1,
62+
S<16, 16, 1>, S<1, 0, 2>, S<1, 0, 2>,
63+
2, 16, 16, 0,
64+
S<16, 16, 1>, S<1, 0, 2>, S<1, 0, 2>,
65+
2, 16, 16, 0,
66+
1, 1, S<1, 16, 1, 16>, S<8, 8, 1>,
67+
ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, ComputeTypeA, ComputeTypeB, PermuteA, PermuteB>;
68+
// clang-format on
69+
70+
#include "run_gemm_wmma_bpreshuffle_example.inc"
71+
72+
int main(int argc, char* argv[]) { return !run_gemm_splitk_example(argc, argv); }

‎example/01_gemm/run_gemm_wmma_bpreshuffle_example.inc‎

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
2+
// SPDX-License-Identifier: MIT
3+
4+
#pragma once
5+
6+
template <typename ProblemType>
7+
bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config)
8+
{
9+
using namespace ck::literals;
10+
11+
auto M = problem_size.M;
12+
auto N = problem_size.N;
13+
auto K = problem_size.K;
14+
auto StrideA = problem_size.StrideA;
15+
auto StrideB = problem_size.StrideB;
16+
auto StrideC = problem_size.StrideC;
17+
auto KBatch = problem_size.KBatch;
18+
19+
auto f_host_tensor_descriptor =
20+
[](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
21+
if constexpr(std::is_same_v<decltype(layout), ck::tensor_layout::gemm::RowMajor>)
22+
{
23+
return HostTensorDescriptor({row, col}, {stride, 1_uz});
24+
}
25+
else
26+
{
27+
return HostTensorDescriptor({row, col}, {1_uz, stride});
28+
}
29+
};
30+
31+
auto f_get_default_stride =
32+
[](std::size_t row, std::size_t col, ck::index_t stride, auto layout) {
33+
if(stride == -1)
34+
{
35+
// give a chance if stride is -1, return a default packed stride
36+
if constexpr(std::is_same_v<decltype(layout), ck::tensor_layout::gemm::RowMajor>)
37+
{
38+
return static_cast<std::size_t>(col);
39+
}
40+
else
41+
{
42+
return static_cast<std::size_t>(row);
43+
}
44+
}
45+
else
46+
return static_cast<std::size_t>(stride);
47+
};
48+
49+
StrideA = f_get_default_stride(M, K, StrideA, ALayout{});
50+
StrideB = f_get_default_stride(K, N, StrideB, BLayout{});
51+
StrideC = f_get_default_stride(M, N, StrideC, CLayout{});
52+
53+
Tensor<ADataType> a_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{}));
54+
Tensor<BDataType> b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
55+
Tensor<BDataType> b_k_n_preshuffled(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
56+
57+
switch(config.init_method)
58+
{
59+
case 0: break;
60+
case 1:
61+
a_m_k.GenerateTensorValue(GeneratorTensor_2<ADataType>{-2, 2});
62+
b_k_n.GenerateTensorValue(GeneratorTensor_2<BDataType>{0, 2});
63+
break;
64+
case 2:
65+
a_m_k.GenerateTensorValue(GeneratorTensor_1<ADataType>{});
66+
b_k_n.GenerateTensorValue(GeneratorTensor_1<BDataType>{});
67+
break;
68+
default:
69+
a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});
70+
b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
71+
}
72+
73+
Tensor<CDataType> c_m_n_host_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
74+
Tensor<CDataType> c_m_n_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
75+
76+
std::cout << "a_m_k: " << a_m_k.mDesc << std::endl;
77+
std::cout << "b_k_n: " << b_k_n.mDesc << std::endl;
78+
std::cout << "b_k_n_preshuffled: " << b_k_n_preshuffled.mDesc << std::endl;
79+
std::cout << "c_m_n: " << c_m_n_host_result.mDesc << std::endl;
80+
81+
DeviceMem a_m_k_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize());
82+
DeviceMem b_k_n_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
83+
DeviceMem c_m_n_device_buf(sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize());
84+
85+
// do GEMM
86+
auto device_op = DeviceOpInstance{};
87+
88+
// weight pre-shuffle
89+
int NPerWmma = device_op.GetPreShuffleParameters();
90+
int KLane = ck::get_warp_size() / NPerWmma;
91+
92+
int K0 = K / (KLane * KPack);
93+
// K -> K0 KLane KPack
94+
// N -> N0 NPerWmma
95+
// N, K -> N0 K0 KLane NPerWmma KPack
96+
int tempk;
97+
for(int n = 0; n < N; ++n)
98+
{
99+
for(int k = 0; k < K; ++k)
100+
{
101+
int n0 = n / NPerWmma;
102+
int n1 = n % NPerWmma;
103+
104+
int k0 = k / (KLane * KPack);
105+
tempk = k % (KLane * KPack);
106+
int k1 = tempk / KPack;
107+
int k2 = tempk % KPack;
108+
109+
int outputIndex = n0 * KPack * NPerWmma * KLane * K0 + k0 * KPack * NPerWmma * KLane +
110+
k1 * KPack * NPerWmma + n1 * KPack + k2;
111+
112+
b_k_n_preshuffled(outputIndex) = b_k_n(n * K + k);
113+
}
114+
}
115+
116+
a_m_k_device_buf.ToDevice(a_m_k.mData.data());
117+
b_k_n_device_buf.ToDevice(b_k_n_preshuffled.mData.data());
118+
c_m_n_device_buf.ToDevice(c_m_n_device_result.mData.data());
119+
120+
auto a_element_op = AElementOp{};
121+
auto b_element_op = BElementOp{};
122+
auto c_element_op = CElementOp{};
123+
124+
auto invoker = device_op.MakeInvoker();
125+
126+
auto argument =
127+
device_op.MakeArgument(static_cast<ADataType*>(a_m_k_device_buf.GetDeviceBuffer()),
128+
static_cast<BDataType*>(b_k_n_device_buf.GetDeviceBuffer()),
129+
static_cast<CDataType*>(c_m_n_device_buf.GetDeviceBuffer()),
130+
M,
131+
N,
132+
K,
133+
StrideA,
134+
StrideB,
135+
StrideC,
136+
KBatch,
137+
a_element_op,
138+
b_element_op,
139+
c_element_op);
140+
141+
if(!device_op.IsSupportedArgument(argument))
142+
{
143+
std::cerr << device_op.GetTypeString() << " does not support this problem" << std::endl;
144+
145+
return true;
146+
}
147+
148+
float ave_time =
149+
invoker.Run(argument, StreamConfig{nullptr, config.time_kernel, 0, 50, 50, false, 1});
150+
151+
bool pass = true;
152+
if(config.do_verification)
153+
{
154+
using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
155+
BDataType,
156+
CDataType,
157+
AccDataType,
158+
PassThrough,
159+
PassThrough,
160+
PassThrough>;
161+
162+
auto ref_gemm = ReferenceGemmInstance{};
163+
auto ref_invoker = ref_gemm.MakeInvoker();
164+
165+
auto ref_argument = ref_gemm.MakeArgument(
166+
a_m_k, b_k_n, c_m_n_host_result, PassThrough{}, PassThrough{}, PassThrough{});
167+
168+
ref_invoker.Run(ref_argument);
169+
170+
invoker.Run(argument, StreamConfig{nullptr, false, 0});
171+
c_m_n_device_buf.FromDevice(c_m_n_device_result.mData.data());
172+
173+
pass &= ck::utils::check_err(c_m_n_device_result,
174+
c_m_n_host_result,
175+
"Error: Incorrect results!",
176+
get_rtol<CDataType>(),
177+
get_atol<CDataType>());
178+
}
179+
180+
if(config.time_kernel)
181+
{
182+
ave_time =
183+
invoker.Run(argument, StreamConfig{nullptr, config.time_kernel, 0, 20, 50, true, 50});
184+
185+
std::size_t flop = 2_uz * M * N * K;
186+
std::size_t num_btype =
187+
sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N;
188+
189+
float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
190+
191+
float gb_per_sec = num_btype / 1.E6 / ave_time;
192+
193+
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec
194+
<< " GB/s, " << device_op.GetTypeString() << std::endl;
195+
}
196+
197+
return pass;
198+
}
199+
200+
bool run_gemm_splitk_example(int argc, char* argv[])
201+
{
202+
ProblemSizeSplitK problem_size{3840, 4096, 4096, 4096, 4096, 4096, 1};
203+
ExecutionConfig config;
204+
205+
return parse_cmd_args(argc, argv, problem_size, config) && run_gemm(problem_size, config);
206+
}

0 commit comments

Comments
 (0)