|
| 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); } |
0 commit comments