mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-18 01:28:27 +00:00
initial commit
This commit is contained in:
aska-0096
9
example/91_tile_program/CMakeLists.txt
Normal file
9
example/91_tile_program/CMakeLists.txt
Normal file
@@ -0,0 +1,9 @@
|
||||
add_example_executable(example_im2col im2col.cpp)
|
||||
add_example_executable(example_gemm gemm.cpp)
|
||||
add_example_executable(example_gemm_gemm gemm_gemm.cpp)
|
||||
add_example_executable(example_reduce reduce.cpp)
|
||||
add_example_executable(example_softmax softmax.cpp)
|
||||
add_example_executable(example_gemm_softmax_gemm gemm_softmax_gemm.cpp)
|
||||
add_example_executable(example_batched_gemm_softmax_gemm batched_gemm_softmax_gemm.cpp)
|
||||
add_example_executable(example_fmha_fwd fmha_fwd.cpp)
|
||||
add_example_executable(example_flash_attention_fwd flash_attention_fwd.cpp)
|
||||
166
example/91_tile_program/batched_gemm_softmax_gemm.cpp
Normal file
166
example/91_tile_program/batched_gemm_softmax_gemm.cpp
Normal file
@@ -0,0 +1,166 @@
|
||||
#include <cstring>
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/cluster_descriptor.hpp"
|
||||
#include "ck/tensor/tensor_view.hpp"
|
||||
#include "ck/host_utility/device_prop.hpp"
|
||||
#include "ck/host_utility/kernel_launch.hpp"
|
||||
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/fill.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
|
||||
#include "reference_batched_gemm.hpp"
|
||||
#include "reference_batched_softmax.hpp"
|
||||
#include "batched_gemm_softmax_gemm.hpp"
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
using QDataType = ck::half_t;
|
||||
using KDataType = ck::half_t;
|
||||
using VDataType = ck::half_t;
|
||||
using SaccDataType = float;
|
||||
using SMPLComputeDataType = float;
|
||||
using PDataType = ck::half_t;
|
||||
using OaccDataType = float;
|
||||
using ODataType = ck::half_t;
|
||||
|
||||
ck::index_t Batch = 16;
|
||||
ck::index_t M0 = 3328;
|
||||
ck::index_t N0 = 4096;
|
||||
ck::index_t K0 = 128;
|
||||
ck::index_t N1 = 128;
|
||||
|
||||
if(argc == 6)
|
||||
{
|
||||
Batch = std::stoi(argv[1]);
|
||||
M0 = std::stoi(argv[2]);
|
||||
N0 = std::stoi(argv[3]);
|
||||
K0 = std::stoi(argv[4]);
|
||||
N1 = std::stoi(argv[5]);
|
||||
}
|
||||
|
||||
std::array<ck::index_t, 3> q_lengths{Batch, M0, K0};
|
||||
std::array<ck::index_t, 3> q_strides{M0 * K0, K0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> k_lengths{Batch, N0, K0};
|
||||
std::array<ck::index_t, 3> k_strides{N0 * K0, K0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> v_lengths{Batch, N1, N0};
|
||||
std::array<ck::index_t, 3> v_strides{N1 * N0, N0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> s_lengths{Batch, M0, N0};
|
||||
std::array<ck::index_t, 3> s_strides{M0 * N0, N0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> p_lengths{Batch, M0, N0};
|
||||
std::array<ck::index_t, 3> p_strides{M0 * N0, N0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> o_lengths{Batch, M0, N1};
|
||||
std::array<ck::index_t, 3> o_strides{M0 * N1, N1, 1};
|
||||
|
||||
// host verify
|
||||
Tensor<QDataType> q_host(q_lengths, q_strides);
|
||||
Tensor<KDataType> k_host(k_lengths, k_strides);
|
||||
Tensor<VDataType> v_host(v_lengths, v_strides);
|
||||
Tensor<SMPLComputeDataType> s_host_ref(s_lengths, s_strides);
|
||||
Tensor<PDataType> p_host_ref(p_lengths, p_strides);
|
||||
Tensor<ODataType> o_host_ref(o_lengths, o_strides);
|
||||
Tensor<ODataType> o_host_dev(o_lengths, o_strides);
|
||||
|
||||
#if 0
|
||||
ck::utils::FillUniformDistributionIntegerValue<QDataType>{-3.f, 3.f}(q_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<KDataType>{-3.f, 3.f}(k_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<VDataType>{-3.f, 3.f}(v_host);
|
||||
#else
|
||||
ck::utils::FillUniformDistribution<QDataType>{-3.f, 3.f}(q_host);
|
||||
ck::utils::FillUniformDistribution<KDataType>{-3.f, 3.f}(k_host);
|
||||
ck::utils::FillUniformDistribution<VDataType>{-3.f, 3.f}(v_host);
|
||||
#endif
|
||||
|
||||
// reference
|
||||
reference_batched_gemm<QDataType, KDataType, SaccDataType, SMPLComputeDataType>(
|
||||
q_host, k_host, s_host_ref);
|
||||
reference_batched_softmax<SMPLComputeDataType, SMPLComputeDataType, PDataType>(s_host_ref,
|
||||
p_host_ref);
|
||||
reference_batched_gemm<PDataType, VDataType, OaccDataType, ODataType>(
|
||||
p_host_ref, v_host, o_host_ref);
|
||||
|
||||
DeviceMem q_buf(sizeof(QDataType) * q_host.GetElementSpaceSize());
|
||||
DeviceMem k_buf(sizeof(KDataType) * k_host.GetElementSpaceSize());
|
||||
DeviceMem v_buf(sizeof(VDataType) * v_host.GetElementSpaceSize());
|
||||
DeviceMem o_buf(sizeof(ODataType) * o_host_ref.GetElementSpaceSize());
|
||||
|
||||
q_buf.ToDevice(q_host.mData.data());
|
||||
k_buf.ToDevice(k_host.mData.data());
|
||||
v_buf.ToDevice(v_host.mData.data());
|
||||
|
||||
constexpr ck::index_t kM0PerBlock = 128;
|
||||
constexpr ck::index_t kN0PerBlock = 128;
|
||||
constexpr ck::index_t kK0PerBlock = 32;
|
||||
constexpr ck::index_t kN1PerBlock = 128;
|
||||
|
||||
constexpr ck::index_t kBlockSize = 256;
|
||||
ck::index_t kGridSize = Batch * (M0 / kM0PerBlock) * (N1 / kN1PerBlock);
|
||||
|
||||
std::cout << "grid size " << kGridSize << std::endl;
|
||||
|
||||
constexpr ck::index_t kWarpPerCu = 8; // 2 warps per SIMD
|
||||
constexpr ck::index_t kWarpPerBlock = kBlockSize / warpSize;
|
||||
constexpr ck::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock;
|
||||
|
||||
float ave_time =
|
||||
launch_kernel<kBlockSize, kBlockPerCu>(StreamConfig{nullptr, true},
|
||||
BatchedGemmSoftmaxGemm<QDataType,
|
||||
KDataType,
|
||||
VDataType,
|
||||
SaccDataType,
|
||||
SMPLComputeDataType,
|
||||
PDataType,
|
||||
OaccDataType,
|
||||
ODataType,
|
||||
kBlockSize,
|
||||
kM0PerBlock,
|
||||
kN0PerBlock,
|
||||
kK0PerBlock,
|
||||
kN1PerBlock>{},
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0,
|
||||
static_cast<QDataType*>(q_buf.GetDeviceBuffer()),
|
||||
static_cast<KDataType*>(k_buf.GetDeviceBuffer()),
|
||||
static_cast<VDataType*>(v_buf.GetDeviceBuffer()),
|
||||
static_cast<ODataType*>(o_buf.GetDeviceBuffer()),
|
||||
M0,
|
||||
N0,
|
||||
K0,
|
||||
N1,
|
||||
Batch,
|
||||
K0, // StrideQ
|
||||
K0, // StrideK
|
||||
N0, // StrideV
|
||||
N1, // StrideO
|
||||
M0 * K0, // BatchStrideQ
|
||||
N0 * K0, // BatchStrideK
|
||||
N1 * N0, // BatchStrideV
|
||||
M0 * N1); // BatchStrideO
|
||||
|
||||
o_buf.FromDevice(o_host_dev.mData.data());
|
||||
|
||||
std::size_t flop =
|
||||
std::size_t(2) * Batch * M0 * N0 * K0 + std::size_t(2) * Batch * M0 * N1 * N0;
|
||||
std::size_t num_btype =
|
||||
sizeof(QDataType) * Batch * M0 * K0 + sizeof(KDataType) * Batch * N0 * K0 +
|
||||
sizeof(VDataType) * Batch * N1 * N0 + sizeof(ODataType) * Batch * M0 * N1;
|
||||
|
||||
float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
|
||||
|
||||
float gb_per_sec = num_btype / 1.E6 / ave_time;
|
||||
|
||||
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
|
||||
<< std::endl;
|
||||
|
||||
return !ck::utils::check_err(o_host_dev, o_host_ref);
|
||||
}
|
||||
109
example/91_tile_program/batched_gemm_softmax_gemm.hpp
Normal file
109
example/91_tile_program/batched_gemm_softmax_gemm.hpp
Normal file
@@ -0,0 +1,109 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/tensor_adaptor.hpp"
|
||||
|
||||
#include "ck/tile_program/tile/tile_distribution.hpp"
|
||||
#include "ck/tile_program/tile/tile_elementwise.hpp"
|
||||
#include "ck/tile_program/tile/tile_gemm_shape.hpp"
|
||||
#include "ck/tile_program/warp_tile/warp_gemm.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_problem.hpp"
|
||||
#include "ck/tile_program/block_tile/block_gemm_areg_bsmem_creg_v1.hpp"
|
||||
#include "ck/tile_program/block_tile/block_reduce.hpp"
|
||||
|
||||
#include "gemm_softmax_gemm_impl.hpp"
|
||||
|
||||
// S[M0, N0] = Q[M0, K0] * K[N0, K0]
|
||||
// P[M0, N0] = Softmax(S[M0, N0])
|
||||
// O[M0, N1] = P[M0, N0] * V[N1, N0]
|
||||
template <typename QDataType,
|
||||
typename KDataType,
|
||||
typename VDataType,
|
||||
typename SaccDataType,
|
||||
typename SMPLComputeDataType,
|
||||
typename PDataType,
|
||||
typename OaccDataType,
|
||||
typename ODataType,
|
||||
ck::index_t kBlockSize,
|
||||
ck::index_t kM0PerBlock,
|
||||
ck::index_t kN0PerBlock,
|
||||
ck::index_t kK0PerBlock,
|
||||
ck::index_t kN1PerBlock>
|
||||
struct BatchedGemmSoftmaxGemm
|
||||
{
|
||||
__device__ void operator()(const QDataType* q_ptr,
|
||||
const KDataType* k_ptr,
|
||||
const VDataType* v_ptr,
|
||||
ODataType* o_ptr,
|
||||
const ck::index_t M0,
|
||||
const ck::index_t N0,
|
||||
const ck::index_t K0,
|
||||
const ck::index_t N1,
|
||||
const ck::index_t /* Batch */,
|
||||
const ck::index_t StrideQ,
|
||||
const ck::index_t StrideK,
|
||||
const ck::index_t StrideV,
|
||||
const ck::index_t StrideO,
|
||||
const ck::index_t BatchStrideQ,
|
||||
const ck::index_t BatchStrideK,
|
||||
const ck::index_t BatchStrideV,
|
||||
const ck::index_t BatchStrideO) const
|
||||
{
|
||||
using namespace ck;
|
||||
|
||||
// divide problem
|
||||
const index_t num_tile_m0 = M0 / kM0PerBlock;
|
||||
const index_t num_tile_n1 = N1 / kN1PerBlock;
|
||||
|
||||
const index_t id_block = get_block_id();
|
||||
|
||||
const auto f = [](index_t dividend, index_t divisor) {
|
||||
index_t quotient = dividend / divisor;
|
||||
index_t modulus = dividend - quotient * divisor;
|
||||
|
||||
return ck::make_tuple(quotient, modulus);
|
||||
};
|
||||
|
||||
const auto [itmp, id_tile_n] = f(id_block, num_tile_n1);
|
||||
const auto [id_tile_batch, id_tile_m] = f(itmp, num_tile_m0);
|
||||
|
||||
const index_t iBatch = __builtin_amdgcn_readfirstlane(id_tile_batch);
|
||||
const index_t iM0 = __builtin_amdgcn_readfirstlane(id_tile_m * kM0PerBlock);
|
||||
const index_t iN1 = __builtin_amdgcn_readfirstlane(id_tile_n * kN1PerBlock);
|
||||
|
||||
const auto kernel_impl = GemmSoftmaxGemmImpl<QDataType,
|
||||
KDataType,
|
||||
VDataType,
|
||||
SaccDataType,
|
||||
SMPLComputeDataType,
|
||||
PDataType,
|
||||
OaccDataType,
|
||||
ODataType,
|
||||
kBlockSize,
|
||||
kM0PerBlock,
|
||||
kN0PerBlock,
|
||||
kK0PerBlock,
|
||||
kN1PerBlock>{};
|
||||
|
||||
kernel_impl(q_ptr + iBatch * BatchStrideQ,
|
||||
k_ptr + iBatch * BatchStrideK,
|
||||
v_ptr + iBatch * BatchStrideV,
|
||||
o_ptr + iBatch * BatchStrideO,
|
||||
M0,
|
||||
N0,
|
||||
K0,
|
||||
N1,
|
||||
StrideQ,
|
||||
StrideK,
|
||||
StrideV,
|
||||
StrideO,
|
||||
iM0,
|
||||
iN1);
|
||||
}
|
||||
};
|
||||
209
example/91_tile_program/flash_attention_fwd.cpp
Normal file
209
example/91_tile_program/flash_attention_fwd.cpp
Normal file
@@ -0,0 +1,209 @@
|
||||
#include <cstring>
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/cluster_descriptor.hpp"
|
||||
#include "ck/tensor/tensor_view.hpp"
|
||||
#include "ck/host_utility/device_prop.hpp"
|
||||
#include "ck/host_utility/kernel_launch.hpp"
|
||||
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/fill.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
|
||||
#include "reference_batched_gemm.hpp"
|
||||
#include "reference_batched_softmax.hpp"
|
||||
#include "flash_attention_fwd.hpp"
|
||||
|
||||
/*
|
||||
* Toy code of flash attention forward pass
|
||||
* Assume simplest case.
|
||||
* Q [Batch, HeadNum, SeqenceLengthQ, HeadDim]
|
||||
* K [Batch, HeadNum, SeqenceLengthK, HeadDim]
|
||||
* V [Batch, HeadNum, HeadDim, SeqenceLengthK]
|
||||
* O [Batch, HeadNum, SeqenceLengthQ, HeadDim]
|
||||
*/
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
using QDataType = ck::half_t;
|
||||
using KDataType = ck::half_t;
|
||||
using VDataType = ck::half_t;
|
||||
using SaccDataType = float;
|
||||
using SMPLComputeDataType = float;
|
||||
using PDataType = ck::half_t;
|
||||
using OaccDataType = float;
|
||||
using ODataType = ck::half_t;
|
||||
|
||||
ck::index_t Batch = 64; // Batch Number * Head Number
|
||||
ck::index_t M0 = 4096; // SequenceLengthQ
|
||||
ck::index_t N0 = 4096; // SequencelengthK
|
||||
ck::index_t K0 = 128; // HeadDim
|
||||
ck::index_t N1 = 128; // HeadDim
|
||||
ck::index_t verification = 0;
|
||||
ck::index_t init_method = 1;
|
||||
ck::index_t time_kernel = 0;
|
||||
|
||||
if(argc == 4)
|
||||
{
|
||||
init_method = std::stoi(argv[1]);
|
||||
time_kernel = std::stoi(argv[2]);
|
||||
verification = std::stoi(argv[3]);
|
||||
}
|
||||
|
||||
if(argc == 9)
|
||||
{
|
||||
init_method = std::stoi(argv[1]);
|
||||
time_kernel = std::stoi(argv[2]);
|
||||
verification = std::stoi(argv[3]);
|
||||
Batch = std::stoi(argv[4]);
|
||||
M0 = std::stoi(argv[5]);
|
||||
N0 = std::stoi(argv[6]);
|
||||
K0 = std::stoi(argv[7]);
|
||||
N1 = std::stoi(argv[8]);
|
||||
}
|
||||
|
||||
std::array<ck::index_t, 3> q_lengths{Batch, M0, K0};
|
||||
std::array<ck::index_t, 3> q_strides{M0 * K0, K0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> k_lengths{Batch, N0, K0};
|
||||
std::array<ck::index_t, 3> k_strides{N0 * K0, K0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> v_lengths{Batch, N1, N0};
|
||||
std::array<ck::index_t, 3> v_strides{N1 * N0, N0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> s_lengths{Batch, M0, N0};
|
||||
std::array<ck::index_t, 3> s_strides{M0 * N0, N0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> p_lengths{Batch, M0, N0};
|
||||
std::array<ck::index_t, 3> p_strides{M0 * N0, N0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> o_lengths{Batch, M0, N1};
|
||||
std::array<ck::index_t, 3> o_strides{M0 * N1, N1, 1};
|
||||
|
||||
// host verify
|
||||
Tensor<QDataType> q_host(q_lengths, q_strides);
|
||||
Tensor<KDataType> k_host(k_lengths, k_strides);
|
||||
Tensor<VDataType> v_host(v_lengths, v_strides);
|
||||
Tensor<ODataType> o_host_dev(o_lengths, o_strides);
|
||||
|
||||
switch(init_method)
|
||||
{
|
||||
case 0: break;
|
||||
case 1:
|
||||
ck::utils::FillUniformDistributionIntegerValue<QDataType>{-3.f, 3.f}(q_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<KDataType>{-3.f, 3.f}(k_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<VDataType>{-3.f, 3.f}(v_host);
|
||||
break;
|
||||
case 2:
|
||||
ck::utils::FillUniformDistribution<QDataType>{-3.f, 3.f}(q_host);
|
||||
ck::utils::FillUniformDistribution<KDataType>{-3.f, 3.f}(k_host);
|
||||
ck::utils::FillUniformDistribution<VDataType>{-3.f, 3.f}(v_host);
|
||||
break;
|
||||
default:
|
||||
ck::utils::FillUniformDistributionIntegerValue<QDataType>{-2.f, 2.f}(q_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<KDataType>{-2.f, 2.f}(k_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<VDataType>{-2.f, 2.f}(v_host);
|
||||
}
|
||||
|
||||
DeviceMem q_buf(sizeof(QDataType) * q_host.GetElementSpaceSize());
|
||||
DeviceMem k_buf(sizeof(KDataType) * k_host.GetElementSpaceSize());
|
||||
DeviceMem v_buf(sizeof(VDataType) * v_host.GetElementSpaceSize());
|
||||
DeviceMem o_buf(sizeof(ODataType) * o_host_dev.GetElementSpaceSize());
|
||||
|
||||
q_buf.ToDevice(q_host.mData.data());
|
||||
k_buf.ToDevice(k_host.mData.data());
|
||||
v_buf.ToDevice(v_host.mData.data());
|
||||
|
||||
constexpr ck::index_t kM0PerBlock = 128;
|
||||
constexpr ck::index_t kN0PerBlock = 128;
|
||||
constexpr ck::index_t kK0PerBlock = 32;
|
||||
constexpr ck::index_t kN1PerBlock = 128;
|
||||
constexpr ck::index_t kK1PerBlock = 32;
|
||||
|
||||
constexpr ck::index_t kBlockSize = 256;
|
||||
constexpr ck::index_t kHeadDim = 128;
|
||||
|
||||
ck::index_t kGridSize = Batch * (M0 / kM0PerBlock) * (N1 / kN1PerBlock);
|
||||
|
||||
std::cout << "grid size " << kGridSize << std::endl;
|
||||
|
||||
constexpr ck::index_t kWarpPerCu = 8; // 2 warps per SIMD
|
||||
constexpr ck::index_t kWarpPerBlock = kBlockSize / warpSize;
|
||||
constexpr ck::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock;
|
||||
|
||||
float ave_time = launch_kernel<kBlockSize, kBlockPerCu>(
|
||||
StreamConfig{nullptr, static_cast<bool>(time_kernel)},
|
||||
FlashAttentionFwd<QDataType,
|
||||
KDataType,
|
||||
VDataType,
|
||||
SaccDataType,
|
||||
SMPLComputeDataType,
|
||||
PDataType,
|
||||
OaccDataType,
|
||||
ODataType,
|
||||
kBlockSize,
|
||||
kHeadDim,
|
||||
kM0PerBlock,
|
||||
kN0PerBlock,
|
||||
kK0PerBlock,
|
||||
kN1PerBlock,
|
||||
kK1PerBlock>{},
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0,
|
||||
static_cast<QDataType*>(q_buf.GetDeviceBuffer()),
|
||||
static_cast<KDataType*>(k_buf.GetDeviceBuffer()),
|
||||
static_cast<VDataType*>(v_buf.GetDeviceBuffer()),
|
||||
static_cast<ODataType*>(o_buf.GetDeviceBuffer()),
|
||||
M0,
|
||||
N0,
|
||||
K0,
|
||||
N1,
|
||||
Batch,
|
||||
K0, // StrideQ
|
||||
K0, // StrideK
|
||||
N0, // StrideV
|
||||
N1, // StrideO
|
||||
M0 * K0, // BatchStrideQ
|
||||
N0 * K0, // BatchStrideK
|
||||
N1 * N0, // BatchStrideV
|
||||
M0 * N1); // BatchStrideO
|
||||
|
||||
// reference
|
||||
auto pass = true;
|
||||
if(verification)
|
||||
{
|
||||
o_buf.FromDevice(o_host_dev.mData.data());
|
||||
|
||||
Tensor<SMPLComputeDataType> s_host_ref(s_lengths, s_strides);
|
||||
Tensor<PDataType> p_host_ref(p_lengths, p_strides);
|
||||
Tensor<ODataType> o_host_ref(o_lengths, o_strides);
|
||||
|
||||
reference_batched_gemm<QDataType, KDataType, SaccDataType, SMPLComputeDataType>(
|
||||
q_host, k_host, s_host_ref);
|
||||
reference_batched_softmax<SMPLComputeDataType, SMPLComputeDataType, PDataType>(s_host_ref,
|
||||
p_host_ref);
|
||||
reference_batched_gemm<PDataType, VDataType, OaccDataType, ODataType>(
|
||||
p_host_ref, v_host, o_host_ref);
|
||||
|
||||
pass &= ck::utils::check_err(o_host_dev, o_host_ref);
|
||||
}
|
||||
|
||||
std::size_t flop =
|
||||
std::size_t(2) * Batch * M0 * N0 * K0 + std::size_t(2) * Batch * M0 * N1 * N0;
|
||||
std::size_t num_btype =
|
||||
sizeof(QDataType) * Batch * M0 * K0 + sizeof(KDataType) * Batch * N0 * K0 +
|
||||
sizeof(VDataType) * Batch * N1 * N0 + sizeof(ODataType) * Batch * M0 * N1;
|
||||
|
||||
float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
|
||||
|
||||
float gb_per_sec = num_btype / 1.E6 / ave_time;
|
||||
|
||||
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
|
||||
<< std::endl;
|
||||
|
||||
return !pass;
|
||||
}
|
||||
113
example/91_tile_program/flash_attention_fwd.hpp
Normal file
113
example/91_tile_program/flash_attention_fwd.hpp
Normal file
@@ -0,0 +1,113 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/tensor_adaptor.hpp"
|
||||
|
||||
#include "ck/tile_program/tile/tile_distribution.hpp"
|
||||
#include "ck/tile_program/tile/tile_elementwise.hpp"
|
||||
#include "ck/tile_program/tile/tile_gemm_shape.hpp"
|
||||
#include "ck/tile_program/warp_tile/warp_gemm.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_problem.hpp"
|
||||
#include "ck/tile_program/block_tile/block_gemm_areg_bsmem_creg_v1.hpp"
|
||||
#include "ck/tile_program/block_tile/block_reduce.hpp"
|
||||
|
||||
#include "flash_attention_fwd_impl.hpp"
|
||||
|
||||
// S[M0, N0] = Q[M0, K0] * K[N0, K0]
|
||||
// P[M0, N0] = Softmax(S[M0, N0])
|
||||
// O[M0, N1] = P[M0, N0] * V[N1, N0]
|
||||
template <typename QDataType,
|
||||
typename KDataType,
|
||||
typename VDataType,
|
||||
typename SaccDataType,
|
||||
typename SMPLComputeDataType,
|
||||
typename PDataType,
|
||||
typename OaccDataType,
|
||||
typename ODataType,
|
||||
ck::index_t kBlockSize,
|
||||
ck::index_t kHeadDim,
|
||||
ck::index_t kM0PerBlock,
|
||||
ck::index_t kN0PerBlock,
|
||||
ck::index_t kK0PerBlock,
|
||||
ck::index_t kN1PerBlock,
|
||||
ck::index_t kK1PerBlock>
|
||||
struct FlashAttentionFwd
|
||||
{
|
||||
__device__ void operator()(const QDataType* q_ptr,
|
||||
const KDataType* k_ptr,
|
||||
const VDataType* v_ptr,
|
||||
ODataType* o_ptr,
|
||||
const ck::index_t M0,
|
||||
const ck::index_t N0,
|
||||
const ck::index_t K0,
|
||||
const ck::index_t N1,
|
||||
const ck::index_t /* Batch */,
|
||||
const ck::index_t StrideQ,
|
||||
const ck::index_t StrideK,
|
||||
const ck::index_t StrideV,
|
||||
const ck::index_t StrideO,
|
||||
const ck::index_t BatchStrideQ,
|
||||
const ck::index_t BatchStrideK,
|
||||
const ck::index_t BatchStrideV,
|
||||
const ck::index_t BatchStrideO) const
|
||||
{
|
||||
using namespace ck;
|
||||
|
||||
// divide problem
|
||||
const index_t num_tile_m0 = M0 / kM0PerBlock;
|
||||
const index_t num_tile_n1 = N1 / kN1PerBlock;
|
||||
|
||||
const index_t id_block = get_block_id();
|
||||
|
||||
const auto f = [](index_t dividend, index_t divisor) {
|
||||
index_t quotient = dividend / divisor;
|
||||
index_t modulus = dividend - quotient * divisor;
|
||||
|
||||
return ck::make_tuple(quotient, modulus);
|
||||
};
|
||||
|
||||
const auto [itmp, id_tile_n] = f(id_block, num_tile_n1);
|
||||
const auto [id_tile_batch, id_tile_m] = f(itmp, num_tile_m0);
|
||||
|
||||
const index_t iBatch = __builtin_amdgcn_readfirstlane(id_tile_batch);
|
||||
const index_t iM0 = __builtin_amdgcn_readfirstlane(id_tile_m * kM0PerBlock);
|
||||
const index_t iN1 = __builtin_amdgcn_readfirstlane(id_tile_n * kN1PerBlock);
|
||||
|
||||
const auto kernel_impl = FlashAttentionFwdImpl<QDataType,
|
||||
KDataType,
|
||||
VDataType,
|
||||
SaccDataType,
|
||||
SMPLComputeDataType,
|
||||
PDataType,
|
||||
OaccDataType,
|
||||
ODataType,
|
||||
kBlockSize,
|
||||
kHeadDim,
|
||||
kM0PerBlock,
|
||||
kN0PerBlock,
|
||||
kK0PerBlock,
|
||||
kN1PerBlock,
|
||||
kK1PerBlock>{};
|
||||
|
||||
kernel_impl(q_ptr + iBatch * BatchStrideQ,
|
||||
k_ptr + iBatch * BatchStrideK,
|
||||
v_ptr + iBatch * BatchStrideV,
|
||||
o_ptr + iBatch * BatchStrideO,
|
||||
M0,
|
||||
N0,
|
||||
K0,
|
||||
N1,
|
||||
StrideQ,
|
||||
StrideK,
|
||||
StrideV,
|
||||
StrideO,
|
||||
iM0,
|
||||
iN1);
|
||||
}
|
||||
};
|
||||
368
example/91_tile_program/flash_attention_fwd_impl.hpp
Normal file
368
example/91_tile_program/flash_attention_fwd_impl.hpp
Normal file
@@ -0,0 +1,368 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/tensor_adaptor.hpp"
|
||||
|
||||
#include "ck/tile_program/tile/tile_distribution.hpp"
|
||||
#include "ck/tile_program/tile/tile_elementwise.hpp"
|
||||
#include "ck/tile_program/tile/tile_gemm_shape.hpp"
|
||||
#include "ck/tile_program/tile/slice_tile.hpp"
|
||||
#include "ck/tile_program/warp_tile/warp_gemm.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2_askiplds.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_problem.hpp"
|
||||
#include "ck/tile_program/block_tile/block_gemm_areg_bsmem_creg_v1.hpp"
|
||||
#include "ck/tile_program/block_tile/block_reduce.hpp"
|
||||
|
||||
// S[M0, N0] = Q[M0, K0] * K[N0, K0]
|
||||
// P[M0, N0] = Softmax(S[M0, N0])
|
||||
// O[M0, N1] = P[M0, N0] * V[N1, N0]
|
||||
template <typename QDataType,
|
||||
typename KDataType,
|
||||
typename VDataType,
|
||||
typename SaccDataType,
|
||||
typename SMPLComputeDataType,
|
||||
typename PDataType,
|
||||
typename OaccDataType,
|
||||
typename ODataType,
|
||||
ck::index_t kBlockSize,
|
||||
ck::index_t kHeadDim,
|
||||
ck::index_t kM0PerBlock,
|
||||
ck::index_t kN0PerBlock,
|
||||
ck::index_t kK0PerBlock,
|
||||
ck::index_t kN1PerBlock,
|
||||
ck::index_t kK1PerBlock>
|
||||
struct FlashAttentionFwdImpl
|
||||
{
|
||||
// block gemm0 pipeline
|
||||
using BlockGemm0Problem = ck::tile_program::block::BlockGemmPipelineProblem<
|
||||
QDataType,
|
||||
KDataType,
|
||||
SaccDataType,
|
||||
kBlockSize,
|
||||
ck::tile_program::TileGemmShape<kM0PerBlock, kN0PerBlock, kK0PerBlock>>;
|
||||
|
||||
using BlockGemm0Policy =
|
||||
ck::tile_program::block::BlockGemmPipelineAGmemBGmemCRegV2SkipALdsPersistentQRegCachePolicy<
|
||||
kHeadDim>;
|
||||
|
||||
using BlockGemm0Pipeline =
|
||||
ck::tile_program::block::BlockGemmPipelineAGmemBGmemCRegV2<BlockGemm0Problem,
|
||||
BlockGemm0Policy>;
|
||||
|
||||
// block gemm1
|
||||
using BlockGemm1 = ck::tile_program::block::BlockGemmARegBSmemCRegV1<
|
||||
ck::tile_program::block::BlockGemmARegBSmemCRegProblem<
|
||||
PDataType,
|
||||
VDataType,
|
||||
OaccDataType,
|
||||
kBlockSize,
|
||||
ck::tile_program::TileGemmShape<kM0PerBlock, kN1PerBlock, kK1PerBlock>>,
|
||||
ck::tile_program::block::BlockGemmARegBSmemCRegV1DefaultPolicy>;
|
||||
|
||||
// 3d, with padding
|
||||
__device__ static constexpr auto MakeVLdsBlockDescriptor()
|
||||
{
|
||||
using namespace ck;
|
||||
|
||||
// using BDataType = B1DataType;
|
||||
|
||||
constexpr index_t kNPerBlock = kN1PerBlock;
|
||||
constexpr index_t kKPerBlock = kK1PerBlock;
|
||||
constexpr index_t kPad = 1;
|
||||
// 2% faster than use kK1 = 8
|
||||
constexpr index_t kK1 = 4;
|
||||
|
||||
constexpr auto b_lds_block_desc_0 = make_naive_tensor_descriptor(
|
||||
make_tuple(Number<kKPerBlock / kK1>{}, Number<kNPerBlock>{}, Number<kK1>{}),
|
||||
make_tuple(Number<(kNPerBlock + kPad) * kK1>{}, Number<kK1>{}, Number<1>{}),
|
||||
Number<kK1>{},
|
||||
Number<1>{});
|
||||
|
||||
constexpr auto b_lds_block_desc = transform_tensor_descriptor(
|
||||
b_lds_block_desc_0,
|
||||
make_tuple(make_pass_through_transform(kNPerBlock),
|
||||
make_merge_transform(make_tuple(Number<kKPerBlock / kK1>{}, Number<kK1>{}))),
|
||||
make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
|
||||
make_tuple(Sequence<0>{}, Sequence<1>{}));
|
||||
|
||||
return b_lds_block_desc;
|
||||
}
|
||||
|
||||
__device__ static constexpr auto MakeVDramTileDistribution()
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
|
||||
using BDataType = VDataType;
|
||||
|
||||
constexpr index_t kNPerBlock = kN1PerBlock;
|
||||
constexpr index_t kKPerBlock = kK1PerBlock;
|
||||
|
||||
constexpr index_t K1 = 16 / sizeof(BDataType);
|
||||
constexpr index_t K0 = kKPerBlock / K1;
|
||||
constexpr index_t N2 = get_warp_size() / K0;
|
||||
constexpr index_t N1 = kBlockSize / get_warp_size();
|
||||
constexpr index_t N0 = kNPerBlock / (N2 * N1);
|
||||
|
||||
return make_static_tile_distribution(
|
||||
StaticTileDistributionEncoding<Sequence<1>,
|
||||
Tuple<Sequence<N0, N1, N2>, Sequence<K0, K1>>,
|
||||
Tuple<Sequence<1>, Sequence<1, 2>>,
|
||||
Tuple<Sequence<1>, Sequence<2, 0>>,
|
||||
Sequence<1, 2>,
|
||||
Sequence<0, 1>>{});
|
||||
}
|
||||
|
||||
__device__ static constexpr ck::index_t GetStaticLdsSize()
|
||||
{
|
||||
using namespace ck;
|
||||
|
||||
return math::max(BlockGemm0Pipeline::GetStaticLdsSize(),
|
||||
static_cast<index_t>(MakeVLdsBlockDescriptor().GetElementSpaceSize() *
|
||||
sizeof(VDataType)));
|
||||
}
|
||||
|
||||
__device__ void operator()(const QDataType* q_ptr,
|
||||
const KDataType* k_ptr,
|
||||
const VDataType* v_ptr,
|
||||
ODataType* o_ptr,
|
||||
const ck::index_t M0,
|
||||
const ck::index_t N0,
|
||||
const ck::index_t K0,
|
||||
const ck::index_t N1,
|
||||
const ck::index_t StrideQ,
|
||||
const ck::index_t StrideK,
|
||||
const ck::index_t StrideV,
|
||||
const ck::index_t StrideO,
|
||||
const ck::index_t iM0,
|
||||
const ck::index_t iN1) const
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
using namespace ck::tile_program::block;
|
||||
|
||||
constexpr auto I0 = Number<0>{};
|
||||
constexpr auto I1 = Number<1>{};
|
||||
|
||||
// allocate LDS
|
||||
__shared__ char smem_ptr[GetStaticLdsSize()];
|
||||
|
||||
// Q/K/V DRAM and DRAM window
|
||||
const auto q_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
q_ptr, make_tuple(M0, K0), make_tuple(StrideQ, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
const auto k_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
k_ptr, make_tuple(N0, K0), make_tuple(StrideK, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
const auto v_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
v_ptr, make_tuple(N1, N0), make_tuple(StrideV, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
auto q_dram_window = make_tile_window(
|
||||
q_dram, make_tuple(Number<kM0PerBlock>{}, Number<kK0PerBlock>{}), {iM0, 0});
|
||||
|
||||
auto k_dram_window = make_tile_window(
|
||||
k_dram, make_tuple(Number<kN0PerBlock>{}, Number<kK0PerBlock>{}), {0, 0});
|
||||
|
||||
auto v_dram_window =
|
||||
make_tile_window(v_dram,
|
||||
make_tuple(Number<kN1PerBlock>{}, Number<kK1PerBlock>{}),
|
||||
{iN1, 0},
|
||||
MakeVDramTileDistribution());
|
||||
|
||||
// Q in Register
|
||||
auto q_reg_tensor = make_static_distributed_tensor<QDataType>(
|
||||
BlockGemm0Policy::template MakeARegBlockDescriptor<BlockGemm0Problem>());
|
||||
|
||||
// V LDS and LDS window
|
||||
// V LDS occupies the same LDS allocation Q/K LDS
|
||||
auto v_lds = make_tensor_view<AddressSpaceEnum::Lds>(reinterpret_cast<VDataType*>(smem_ptr),
|
||||
MakeVLdsBlockDescriptor());
|
||||
|
||||
auto v_lds_window = make_tile_window(
|
||||
v_lds, make_tuple(Number<kN1PerBlock>{}, Number<kK1PerBlock>{}), {0, 0});
|
||||
|
||||
// Block GEMM0 pipeline and Block GEMM1
|
||||
constexpr auto gemm0_pipeline = BlockGemm0Pipeline{};
|
||||
constexpr auto gemm1 = BlockGemm1{};
|
||||
|
||||
// reduction function for softmax
|
||||
const auto f_max = [](auto e0, auto e1) { return max(e0, e1); };
|
||||
const auto f_sum = [](auto e0, auto e1) { return e0 + e1; };
|
||||
|
||||
// infer Sacc, S, P, M, L, Oacc type
|
||||
using SaccBlockTileType =
|
||||
decltype(gemm0_pipeline(q_dram_window, k_dram_window, q_reg_tensor, nullptr));
|
||||
|
||||
using SBlockTileType = decltype(tile_elementwise_in(
|
||||
type_convert<SMPLComputeDataType, SaccDataType>, SaccBlockTileType{}));
|
||||
|
||||
using PBlockTileType = decltype(tile_elementwise_in(type_convert<PDataType, SaccDataType>,
|
||||
SaccBlockTileType{}));
|
||||
|
||||
using MLBlockTileType = decltype(block_tile_reduce<SMPLComputeDataType>(
|
||||
SBlockTileType{}, Sequence<1>{}, f_max, SMPLComputeDataType{0}));
|
||||
|
||||
using OaccBlockTileType = decltype(gemm1(
|
||||
get_slice_tile(
|
||||
PBlockTileType{}, Sequence<0, 0>{}, Sequence<kM0PerBlock, kK1PerBlock>{}),
|
||||
v_dram_window));
|
||||
|
||||
// init Sacc, Oacc, M, L
|
||||
auto s_acc = SaccBlockTileType{};
|
||||
auto o_acc = OaccBlockTileType{};
|
||||
auto m = MLBlockTileType{};
|
||||
auto l = MLBlockTileType{};
|
||||
|
||||
tile_elementwise_inout([](auto& e) { e = 0; }, o_acc);
|
||||
tile_elementwise_inout([](auto& e) { e = NumericLimits<SMPLComputeDataType>::Lowest(); },
|
||||
m);
|
||||
tile_elementwise_inout([](auto& e) { e = 0; }, l);
|
||||
|
||||
// loop over Column of S (J loop)
|
||||
index_t iN0 = 0;
|
||||
|
||||
// Cold Q_Reg_Cache
|
||||
s_acc = gemm0_pipeline(q_dram_window, k_dram_window, q_reg_tensor, smem_ptr);
|
||||
do
|
||||
{
|
||||
// Hot Q_Reg_Cache
|
||||
if(iN0 > 0)
|
||||
{
|
||||
s_acc = gemm0_pipeline(k_dram_window, q_reg_tensor, smem_ptr);
|
||||
}
|
||||
// S{j}
|
||||
const auto s =
|
||||
tile_elementwise_in(type_convert<SMPLComputeDataType, SaccDataType>, s_acc);
|
||||
|
||||
// prefetch load v tile
|
||||
const auto v_prefetch = load_tile(v_dram_window);
|
||||
|
||||
// m_local = rowmax(S{j})
|
||||
auto m_local = block_tile_reduce<SMPLComputeDataType>(
|
||||
s, Sequence<1>{}, f_max, NumericLimits<SMPLComputeDataType>::Lowest());
|
||||
|
||||
block_tile_reduce_sync(m_local, f_max);
|
||||
|
||||
// m{j-1}
|
||||
const auto m_old = m;
|
||||
|
||||
// m{j}
|
||||
tile_elementwise_inout(
|
||||
[](auto& e0, auto e1, auto e2) { e0 = max(e1, e2); }, m, m_old, m_local);
|
||||
|
||||
// Pcompute{j}
|
||||
auto p_compute =
|
||||
make_static_distributed_tensor<SMPLComputeDataType>(s.GetTileDistribution());
|
||||
|
||||
constexpr auto p_spans = decltype(p_compute)::GetDistributedSpans();
|
||||
|
||||
sweep_tile_span(p_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
sweep_tile_span(p_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
p_compute(i_j_idx) = math::exp(s[i_j_idx] - m[i_idx]);
|
||||
});
|
||||
});
|
||||
|
||||
// rowsum(Pcompute{j})
|
||||
auto rowsum_p = block_tile_reduce<SMPLComputeDataType>(
|
||||
p_compute, Sequence<1>{}, f_sum, SMPLComputeDataType{0});
|
||||
|
||||
block_tile_reduce_sync(rowsum_p, f_sum);
|
||||
|
||||
// l{j}, Oacc{j}
|
||||
sweep_tile_span(p_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
const auto tmp = math::exp(m_old[i_idx] - m[i_idx]);
|
||||
|
||||
l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx];
|
||||
|
||||
sweep_tile_span(p_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
o_acc(i_j_idx) *= tmp;
|
||||
});
|
||||
});
|
||||
|
||||
block_sync_lds();
|
||||
store_tile(v_lds_window, v_prefetch);
|
||||
move_tile_window(v_dram_window, {0, kK1PerBlock});
|
||||
|
||||
// type cast Pcompute{j} into P{j}
|
||||
const auto p =
|
||||
tile_elementwise_in(type_convert<PDataType, SMPLComputeDataType>, p_compute);
|
||||
|
||||
// Oacc{j}
|
||||
constexpr index_t k1_loops = kN0PerBlock / kK1PerBlock;
|
||||
|
||||
if constexpr(k1_loops > 1)
|
||||
{
|
||||
static_for<0, k1_loops - 1, 1>{}([&](auto i_k1) {
|
||||
const auto v = load_tile(v_dram_window); // load next v
|
||||
block_sync_lds();
|
||||
gemm1(o_acc,
|
||||
get_slice_tile(p,
|
||||
Sequence<0, i_k1 * kK1PerBlock>{},
|
||||
Sequence<kM0PerBlock, (i_k1 + 1) * kK1PerBlock>{}),
|
||||
v_lds_window);
|
||||
block_sync_lds();
|
||||
store_tile(v_lds_window, v);
|
||||
move_tile_window(v_dram_window, {0, kK1PerBlock});
|
||||
});
|
||||
}
|
||||
// tail
|
||||
{
|
||||
block_sync_lds();
|
||||
gemm1(o_acc,
|
||||
get_slice_tile(p,
|
||||
Sequence<0, (k1_loops - 1) * kK1PerBlock>{},
|
||||
Sequence<kM0PerBlock, kN0PerBlock>{}),
|
||||
v_lds_window);
|
||||
block_sync_lds();
|
||||
}
|
||||
// move tile windows
|
||||
move_tile_window(k_dram_window, {kN0PerBlock, 0});
|
||||
iN0 += kN0PerBlock;
|
||||
} while(iN0 < N0);
|
||||
|
||||
// Oacc
|
||||
constexpr auto o_spans = decltype(o_acc)::GetDistributedSpans();
|
||||
|
||||
sweep_tile_span(o_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
const auto tmp = 1 / l[i_idx];
|
||||
|
||||
sweep_tile_span(o_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
o_acc(i_j_idx) *= tmp;
|
||||
});
|
||||
});
|
||||
|
||||
// type cast Oacc into O
|
||||
const auto o = tile_elementwise_in(type_convert<ODataType, OaccDataType>, o_acc);
|
||||
|
||||
// O DRAM and O DRAM window
|
||||
auto o_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
o_ptr, make_tuple(M0, N1), make_tuple(StrideO, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
auto o_dram_window =
|
||||
make_tile_window(o_dram,
|
||||
make_tuple(Number<kM0PerBlock>{}, Number<kN1PerBlock>{}),
|
||||
{iM0, iN1},
|
||||
o.GetTileDistribution());
|
||||
|
||||
// store O
|
||||
store_tile(o_dram_window, o);
|
||||
}
|
||||
};
|
||||
176
example/91_tile_program/fmha_fwd.cpp
Normal file
176
example/91_tile_program/fmha_fwd.cpp
Normal file
@@ -0,0 +1,176 @@
|
||||
#include <cstring>
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/cluster_descriptor.hpp"
|
||||
#include "ck/tensor/tensor_view.hpp"
|
||||
#include "ck/host_utility/device_prop.hpp"
|
||||
#include "ck/host_utility/kernel_launch.hpp"
|
||||
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/fill.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
|
||||
#include "ck/tile_program/block_tile_pipeline/block_fmha_pipeline_qkvs.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_fmha_pipeline_qkvs_default_policy.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_fmha_pipeline_problem.hpp"
|
||||
#include "ck/tile_program/tile/tile_fmha_shape.hpp"
|
||||
|
||||
#include "reference_batched_gemm.hpp"
|
||||
#include "reference_batched_softmax.hpp"
|
||||
#include "fmha_fwd_kernel.hpp"
|
||||
#include "fmha_fwd_epilogue.hpp"
|
||||
|
||||
using QDataType = ck::half_t;
|
||||
using KDataType = ck::half_t;
|
||||
using VDataType = ck::half_t;
|
||||
using SaccDataType = float; // data type for first gemm accumulation
|
||||
using SMPLComputeDataType = float; // data type for reduction, softmax
|
||||
using PDataType = ck::half_t; // data type for A matrix of second gemm
|
||||
using OaccDataType = float; // data type for second gemm accumulation
|
||||
using ODataType = ck::half_t;
|
||||
|
||||
using FmhaShape =
|
||||
ck::tile_program::TileFmhaShape<128 /*M0*/, 128 /*N0*/, 32 /*K0*/, 128 /*N1*/, 32 /*K1*/>;
|
||||
|
||||
using FmhaPipelineProblem = ck::tile_program::block::BlockFmhaPipelineProblem<QDataType,
|
||||
KDataType,
|
||||
VDataType,
|
||||
SaccDataType,
|
||||
SMPLComputeDataType,
|
||||
PDataType,
|
||||
OaccDataType,
|
||||
ODataType,
|
||||
256, // BlockSize
|
||||
FmhaShape>;
|
||||
using FmhaPipeline = ck::tile_program::block::BlockFmhaPipelineQKVS<FmhaPipelineProblem>;
|
||||
|
||||
using FmhaEpilogue = FmhaFwdEpilogue<FmhaFwdEpilogueProblem<OaccDataType, ODataType>>;
|
||||
using FmhaKernel = FmhaFwdKernel<FmhaPipeline, FmhaEpilogue>;
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
ck::index_t Batch = 16; // batch * nheads
|
||||
ck::index_t M0 = 3328; // seqlen_q
|
||||
ck::index_t N0 = 4096; // seqlen_k
|
||||
ck::index_t K0 = 128; // hdim_q
|
||||
ck::index_t N1 = 128; // hdim_v
|
||||
|
||||
if(argc == 6)
|
||||
{
|
||||
Batch = std::stoi(argv[1]);
|
||||
M0 = std::stoi(argv[2]);
|
||||
N0 = std::stoi(argv[3]);
|
||||
K0 = std::stoi(argv[4]);
|
||||
N1 = std::stoi(argv[5]);
|
||||
}
|
||||
|
||||
std::array<ck::index_t, 3> q_lengths{Batch, M0, K0};
|
||||
std::array<ck::index_t, 3> q_strides{M0 * K0, K0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> k_lengths{Batch, N0, K0};
|
||||
std::array<ck::index_t, 3> k_strides{N0 * K0, K0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> v_lengths{Batch, N1, N0};
|
||||
std::array<ck::index_t, 3> v_strides{N1 * N0, N0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> s_lengths{Batch, M0, N0};
|
||||
std::array<ck::index_t, 3> s_strides{M0 * N0, N0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> p_lengths{Batch, M0, N0};
|
||||
std::array<ck::index_t, 3> p_strides{M0 * N0, N0, 1};
|
||||
|
||||
std::array<ck::index_t, 3> o_lengths{Batch, M0, N1};
|
||||
std::array<ck::index_t, 3> o_strides{M0 * N1, N1, 1};
|
||||
|
||||
// host verify
|
||||
Tensor<QDataType> q_host(q_lengths, q_strides);
|
||||
Tensor<KDataType> k_host(k_lengths, k_strides);
|
||||
Tensor<VDataType> v_host(v_lengths, v_strides);
|
||||
Tensor<SMPLComputeDataType> s_host_ref(s_lengths, s_strides);
|
||||
Tensor<PDataType> p_host_ref(p_lengths, p_strides);
|
||||
Tensor<ODataType> o_host_ref(o_lengths, o_strides);
|
||||
Tensor<ODataType> o_host_dev(o_lengths, o_strides);
|
||||
|
||||
#if 0
|
||||
ck::utils::FillUniformDistributionIntegerValue<QDataType>{-3.f, 3.f}(q_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<KDataType>{-3.f, 3.f}(k_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<VDataType>{-3.f, 3.f}(v_host);
|
||||
#else
|
||||
ck::utils::FillUniformDistribution<QDataType>{-3.f, 3.f}(q_host);
|
||||
ck::utils::FillUniformDistribution<KDataType>{-3.f, 3.f}(k_host);
|
||||
ck::utils::FillUniformDistribution<VDataType>{-3.f, 3.f}(v_host);
|
||||
#endif
|
||||
|
||||
// reference
|
||||
reference_batched_gemm<QDataType, KDataType, SaccDataType, SMPLComputeDataType>(
|
||||
q_host, k_host, s_host_ref);
|
||||
reference_batched_softmax<SMPLComputeDataType, SMPLComputeDataType, PDataType>(s_host_ref,
|
||||
p_host_ref);
|
||||
reference_batched_gemm<PDataType, VDataType, OaccDataType, ODataType>(
|
||||
p_host_ref, v_host, o_host_ref);
|
||||
|
||||
DeviceMem q_buf(sizeof(QDataType) * q_host.GetElementSpaceSize());
|
||||
DeviceMem k_buf(sizeof(KDataType) * k_host.GetElementSpaceSize());
|
||||
DeviceMem v_buf(sizeof(VDataType) * v_host.GetElementSpaceSize());
|
||||
DeviceMem o_buf(sizeof(ODataType) * o_host_ref.GetElementSpaceSize());
|
||||
|
||||
q_buf.ToDevice(q_host.mData.data());
|
||||
k_buf.ToDevice(k_host.mData.data());
|
||||
v_buf.ToDevice(v_host.mData.data());
|
||||
|
||||
dim3 kGridSize = FmhaKernel::GridSize(Batch, M0, N1);
|
||||
constexpr dim3 kBlockSize = FmhaKernel::BlockSize();
|
||||
|
||||
std::cout << "batch:" << Batch << ", seqlen_q:" << M0 << ", seqlen_k:" << N0
|
||||
<< ", hdim_q:" << K0 << ", hdim_v:" << N1 << ", grid_size " << kGridSize.x
|
||||
<< std::endl;
|
||||
|
||||
constexpr ck::index_t kWarpPerCu = 8; // 2 warps per SIMD
|
||||
constexpr ck::index_t kWarpPerBlock = kBlockSize.x / warpSize;
|
||||
constexpr ck::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock;
|
||||
|
||||
auto kargs = FmhaKernel::MakeKargs(q_buf.GetDeviceBuffer(),
|
||||
k_buf.GetDeviceBuffer(),
|
||||
v_buf.GetDeviceBuffer(),
|
||||
o_buf.GetDeviceBuffer(),
|
||||
M0, // seqlen_q
|
||||
N0, // seqlen_k
|
||||
K0, // hdim_q
|
||||
N1, // hdim_v
|
||||
K0, // stride_q
|
||||
K0, // stride_k
|
||||
N0, // stride_v
|
||||
N1, // stride_o
|
||||
M0 * K0, // batch_stride_q
|
||||
N0 * K0, // batch_stride_k
|
||||
N1 * N0, // batch_stride_v
|
||||
M0 * N1); // batch_stride_o
|
||||
|
||||
float ave_time = launch_kernel<kBlockSize.x, kBlockPerCu>(StreamConfig{nullptr, true},
|
||||
FmhaKernel{},
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0,
|
||||
kargs); // BatchStrideO
|
||||
|
||||
o_buf.FromDevice(o_host_dev.mData.data());
|
||||
|
||||
std::size_t flop =
|
||||
std::size_t(2) * Batch * M0 * N0 * K0 + std::size_t(2) * Batch * M0 * N1 * N0;
|
||||
|
||||
std::size_t num_btype =
|
||||
sizeof(QDataType) * Batch * M0 * K0 + sizeof(KDataType) * Batch * N0 * K0 +
|
||||
sizeof(VDataType) * Batch * N1 * N0 + sizeof(ODataType) * Batch * M0 * N1;
|
||||
|
||||
float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
|
||||
|
||||
float gb_per_sec = num_btype / 1.E6 / ave_time;
|
||||
|
||||
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
|
||||
<< std::endl;
|
||||
|
||||
return !ck::utils::check_err(o_host_dev, o_host_ref);
|
||||
}
|
||||
35
example/91_tile_program/fmha_fwd_epilogue.hpp
Normal file
35
example/91_tile_program/fmha_fwd_epilogue.hpp
Normal file
@@ -0,0 +1,35 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tile_program/tile/store_tile.hpp"
|
||||
#include "ck/tile_program/tile/tile_elementwise.hpp"
|
||||
|
||||
template <typename OaccDataType_, typename ODataType_>
|
||||
struct FmhaFwdEpilogueProblem
|
||||
{
|
||||
using OaccDataType = ck::remove_cvref_t<OaccDataType_>;
|
||||
using ODataType = ck::remove_cvref_t<ODataType_>;
|
||||
};
|
||||
|
||||
template <typename Problem_, typename Policy_ = void>
|
||||
struct FmhaFwdEpilogue
|
||||
{
|
||||
using Problem = ck::remove_cvref_t<Problem_>;
|
||||
using OaccDataType = ck::remove_cvref_t<typename Problem::OaccDataType>;
|
||||
using ODataType = ck::remove_cvref_t<typename Problem::ODataType>;
|
||||
|
||||
__host__ __device__ static constexpr ck::index_t GetSmemSize() { return 0; }
|
||||
|
||||
template <typename ODramWindowTmp, typename OAccTile>
|
||||
__device__ auto operator()(ODramWindowTmp& o_dram_window_tmp, const OAccTile& o_acc_tile)
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
|
||||
const auto o = tile_elementwise_in(type_convert<ODataType, OaccDataType>, o_acc_tile);
|
||||
store_tile(o_dram_window_tmp, o);
|
||||
}
|
||||
};
|
||||
191
example/91_tile_program/fmha_fwd_kernel.hpp
Normal file
191
example/91_tile_program/fmha_fwd_kernel.hpp
Normal file
@@ -0,0 +1,191 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor/tensor_view.hpp"
|
||||
#include "ck/tile_program/tile/tile_window.hpp"
|
||||
|
||||
// S[seqlen_q, seqlen_k] = Q[seqlen_q, hdim_q] * K[seqlen_k, hdim_q]
|
||||
// P[seqlen_q, seqlen_k] = Softmax(S[seqlen_q, seqlen_k])
|
||||
// O[seqlen_q, hdim_v] = P[seqlen_q, seqlen_k] * V[hdim_v, seqlen_k]
|
||||
|
||||
template <typename FmhaPipeline_, typename EpiloguePipeline_>
|
||||
struct FmhaFwdKernel
|
||||
{
|
||||
using FmhaPipeline = ck::remove_cvref_t<FmhaPipeline_>;
|
||||
using EpiloguePipeline = ck::remove_cvref_t<EpiloguePipeline_>;
|
||||
static constexpr ck::index_t kBlockSize = FmhaPipeline::kBlockSize;
|
||||
|
||||
using QDataType = ck::remove_cvref_t<typename FmhaPipeline::QDataType>;
|
||||
using KDataType = ck::remove_cvref_t<typename FmhaPipeline::KDataType>;
|
||||
using VDataType = ck::remove_cvref_t<typename FmhaPipeline::VDataType>;
|
||||
using ODataType = ck::remove_cvref_t<typename FmhaPipeline::ODataType>;
|
||||
|
||||
struct Kargs
|
||||
{
|
||||
const void* q_ptr;
|
||||
const void* k_ptr;
|
||||
const void* v_ptr;
|
||||
void* o_ptr;
|
||||
ck::index_t seqlen_q;
|
||||
ck::index_t seqlen_k;
|
||||
ck::index_t hdim_q;
|
||||
ck::index_t hdim_v;
|
||||
ck::index_t stride_q;
|
||||
ck::index_t stride_k;
|
||||
ck::index_t stride_v;
|
||||
ck::index_t stride_o;
|
||||
ck::index_t batch_stride_q;
|
||||
ck::index_t batch_stride_k;
|
||||
ck::index_t batch_stride_v;
|
||||
ck::index_t batch_stride_o;
|
||||
};
|
||||
|
||||
__host__ static constexpr Kargs MakeKargs(const void* q_ptr,
|
||||
const void* k_ptr,
|
||||
const void* v_ptr,
|
||||
void* o_ptr,
|
||||
ck::index_t seqlen_q,
|
||||
ck::index_t seqlen_k,
|
||||
ck::index_t hdim_q,
|
||||
ck::index_t hdim_v,
|
||||
ck::index_t stride_q,
|
||||
ck::index_t stride_k,
|
||||
ck::index_t stride_v,
|
||||
ck::index_t stride_o,
|
||||
ck::index_t batch_stride_q,
|
||||
ck::index_t batch_stride_k,
|
||||
ck::index_t batch_stride_v,
|
||||
ck::index_t batch_stride_o)
|
||||
{
|
||||
return Kargs{q_ptr,
|
||||
k_ptr,
|
||||
v_ptr,
|
||||
o_ptr,
|
||||
seqlen_q,
|
||||
seqlen_k,
|
||||
hdim_q,
|
||||
hdim_v,
|
||||
stride_q,
|
||||
stride_k,
|
||||
stride_v,
|
||||
stride_o,
|
||||
batch_stride_q,
|
||||
batch_stride_k,
|
||||
batch_stride_v,
|
||||
batch_stride_o};
|
||||
}
|
||||
|
||||
__host__ static constexpr auto
|
||||
GridSize(ck::index_t batch_size_, ck::index_t seqlen_q_, ck::index_t hdim_v_)
|
||||
{
|
||||
return dim3(batch_size_ * (seqlen_q_ / FmhaPipeline::kM0) * (hdim_v_ / FmhaPipeline::kN1));
|
||||
}
|
||||
|
||||
__host__ static constexpr auto BlockSize() { return dim3(kBlockSize); }
|
||||
|
||||
__host__ __device__ static constexpr ck::index_t GetSmemSize()
|
||||
{
|
||||
return ck::math::max(FmhaPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize());
|
||||
}
|
||||
|
||||
__device__ void operator()(Kargs kargs) const
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
using namespace ck::tile_program::block;
|
||||
|
||||
// allocate LDS
|
||||
__shared__ char smem_ptr[GetSmemSize()];
|
||||
|
||||
// divide problem
|
||||
const index_t num_tile_m0 = kargs.seqlen_q / FmhaPipeline::kM0;
|
||||
const index_t num_tile_n1 = kargs.hdim_v / FmhaPipeline::kN1;
|
||||
|
||||
const index_t id_block = ck::get_block_id();
|
||||
|
||||
const auto f = [](index_t dividend, index_t divisor) {
|
||||
index_t quotient = dividend / divisor;
|
||||
index_t modulus = dividend - quotient * divisor;
|
||||
|
||||
return ck::make_tuple(quotient, modulus);
|
||||
};
|
||||
|
||||
const auto [itmp, id_tile_n] = f(id_block, num_tile_n1);
|
||||
const auto [id_tile_batch, id_tile_m] = f(itmp, num_tile_m0);
|
||||
|
||||
const index_t i_batch = __builtin_amdgcn_readfirstlane(id_tile_batch);
|
||||
const index_t i_m0 = __builtin_amdgcn_readfirstlane(id_tile_m * FmhaPipeline::kM0);
|
||||
const index_t i_n1 = __builtin_amdgcn_readfirstlane(id_tile_n * FmhaPipeline::kN1);
|
||||
|
||||
// for simplicity, batch stride we just modify the pointer
|
||||
const QDataType* q_ptr =
|
||||
reinterpret_cast<const QDataType*>(kargs.q_ptr) + i_batch * kargs.batch_stride_q;
|
||||
const KDataType* k_ptr =
|
||||
reinterpret_cast<const KDataType*>(kargs.k_ptr) + i_batch * kargs.batch_stride_k;
|
||||
const VDataType* v_ptr =
|
||||
reinterpret_cast<const VDataType*>(kargs.v_ptr) + i_batch * kargs.batch_stride_v;
|
||||
ODataType* o_ptr =
|
||||
reinterpret_cast<ODataType*>(kargs.o_ptr) + i_batch * kargs.batch_stride_o;
|
||||
|
||||
// Q/K/V DRAM and DRAM window
|
||||
// FIXME: assume layout Q[seqlen_q, hdim_q], K[seqlen_k, hdim_q], V[hdim_v, seqlen_k],
|
||||
const auto q_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
q_ptr,
|
||||
make_tuple(kargs.seqlen_q, kargs.hdim_q),
|
||||
make_tuple(kargs.stride_q, 1),
|
||||
Number<32>{},
|
||||
Number<1>{});
|
||||
|
||||
const auto k_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
k_ptr,
|
||||
make_tuple(kargs.seqlen_k, kargs.hdim_q),
|
||||
make_tuple(kargs.stride_k, 1),
|
||||
Number<32>{},
|
||||
Number<1>{});
|
||||
|
||||
const auto v_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
v_ptr,
|
||||
make_tuple(kargs.hdim_v, kargs.seqlen_k),
|
||||
make_tuple(kargs.stride_v, 1),
|
||||
Number<32>{},
|
||||
Number<1>{});
|
||||
|
||||
auto q_dram_window =
|
||||
make_tile_window(q_dram,
|
||||
make_tuple(Number<FmhaPipeline::kM0>{}, Number<FmhaPipeline::kK0>{}),
|
||||
{i_m0, 0});
|
||||
|
||||
auto k_dram_window = make_tile_window(
|
||||
k_dram, make_tuple(Number<FmhaPipeline::kN0>{}, Number<FmhaPipeline::kK0>{}), {0, 0});
|
||||
|
||||
auto v_dram_window =
|
||||
make_tile_window(v_dram,
|
||||
make_tuple(Number<FmhaPipeline::kN1>{}, Number<FmhaPipeline::kK1>{}),
|
||||
{i_n1, 0});
|
||||
|
||||
auto o_acc_tile = FmhaPipeline{}(q_dram_window,
|
||||
k_dram_window,
|
||||
v_dram_window,
|
||||
kargs.seqlen_k / FmhaPipeline::kN0,
|
||||
kargs.hdim_q / FmhaPipeline::kK0,
|
||||
smem_ptr);
|
||||
|
||||
// O DRAM and O DRAM window
|
||||
auto o_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
o_ptr,
|
||||
make_tuple(kargs.seqlen_q, kargs.hdim_v),
|
||||
make_tuple(kargs.stride_o, 1),
|
||||
Number<32>{},
|
||||
Number<1>{});
|
||||
|
||||
auto o_dram_window =
|
||||
make_tile_window(o_dram,
|
||||
make_tuple(Number<FmhaPipeline::kM0>{}, Number<FmhaPipeline::kN1>{}),
|
||||
{i_m0, i_n1});
|
||||
|
||||
EpiloguePipeline{}(o_dram_window, o_acc_tile);
|
||||
}
|
||||
};
|
||||
195
example/91_tile_program/gemm.cpp
Normal file
195
example/91_tile_program/gemm.cpp
Normal file
@@ -0,0 +1,195 @@
|
||||
#include <cstring>
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/cluster_descriptor.hpp"
|
||||
#include "ck/tensor/tensor_view.hpp"
|
||||
#include "ck/host_utility/device_prop.hpp"
|
||||
#include "ck/host_utility/kernel_launch.hpp"
|
||||
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/fill.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
|
||||
#include "reference_gemm.hpp"
|
||||
#include "gemm.hpp"
|
||||
|
||||
/*
|
||||
* Toy code of GEMM
|
||||
* Assume simplest case.
|
||||
* A [M, K]
|
||||
* B [N, K]
|
||||
* C [M, N]
|
||||
*/
|
||||
|
||||
// elementwise lambda
|
||||
struct AElementFunction
|
||||
{
|
||||
template <typename X>
|
||||
__host__ __device__ auto operator()(const X& x) const
|
||||
{
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct BElementFunction
|
||||
{
|
||||
template <typename X>
|
||||
__host__ __device__ auto operator()(const X& x) const
|
||||
{
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct CElementFunction
|
||||
{
|
||||
template <typename X>
|
||||
__host__ __device__ auto operator()(const X& x) const
|
||||
{
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
using ADataType = ck::half_t;
|
||||
using BDataType = ck::half_t;
|
||||
using AccDataType = float;
|
||||
using CDataType = ck::half_t;
|
||||
|
||||
using ALayout = ck::tensor_layout::gemm::RowMajor;
|
||||
using BLayout = ck::tensor_layout::gemm::ColumnMajor;
|
||||
using CLayout = ck::tensor_layout::gemm::RowMajor;
|
||||
|
||||
ck::index_t verification = 0;
|
||||
ck::index_t M = 3328;
|
||||
ck::index_t N = 4096;
|
||||
ck::index_t K = 4096;
|
||||
|
||||
if(argc == 2)
|
||||
{
|
||||
verification = std::stoi(argv[1]);
|
||||
}
|
||||
if(argc == 5)
|
||||
{
|
||||
verification = std::stoi(argv[1]);
|
||||
M = std::stoi(argv[1]);
|
||||
N = std::stoi(argv[2]);
|
||||
K = std::stoi(argv[3]);
|
||||
}
|
||||
|
||||
const ck::index_t Lda = std::is_same_v<ALayout, ck::tensor_layout::gemm::RowMajor> ? K : M;
|
||||
const ck::index_t Ldb = std::is_same_v<BLayout, ck::tensor_layout::gemm::ColumnMajor> ? K : N;
|
||||
const ck::index_t Ldc = std::is_same_v<CLayout, ck::tensor_layout::gemm::RowMajor> ? N : M;
|
||||
|
||||
const auto a_lengths = std::array<ck::index_t, 2>{M, K};
|
||||
const auto a_strides = std::is_same_v<ALayout, ck::tensor_layout::gemm::RowMajor>
|
||||
? std::array<ck::index_t, 2>{Lda, 1}
|
||||
: std::array<ck::index_t, 2>{1, Lda};
|
||||
|
||||
const auto b_lengths = std::array<ck::index_t, 2>{N, K};
|
||||
const auto b_strides = std::is_same_v<BLayout, ck::tensor_layout::gemm::ColumnMajor>
|
||||
? std::array<ck::index_t, 2>{Ldb, 1}
|
||||
: std::array<ck::index_t, 2>{1, Ldb};
|
||||
|
||||
const auto c_lengths = std::array<ck::index_t, 2>{M, N};
|
||||
const auto c_strides = std::is_same_v<CLayout, ck::tensor_layout::gemm::RowMajor>
|
||||
? std::array<ck::index_t, 2>{Ldc, 1}
|
||||
: std::array<ck::index_t, 2>{1, Ldc};
|
||||
|
||||
// host verify
|
||||
Tensor<ADataType> a_host(a_lengths, a_strides);
|
||||
Tensor<BDataType> b_host(b_lengths, b_strides);
|
||||
Tensor<CDataType> c_host_dev(c_lengths, c_strides);
|
||||
|
||||
ck::utils::FillUniformDistributionIntegerValue<ADataType>{-5.f, 5.f}(a_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<BDataType>{-5.f, 5.f}(b_host);
|
||||
|
||||
DeviceMem a_buf(sizeof(ADataType) * a_host.GetElementSpaceSize());
|
||||
DeviceMem b_buf(sizeof(BDataType) * b_host.GetElementSpaceSize());
|
||||
DeviceMem c_buf(sizeof(CDataType) * c_host_dev.GetElementSpaceSize());
|
||||
|
||||
a_buf.ToDevice(a_host.mData.data());
|
||||
b_buf.ToDevice(b_host.mData.data());
|
||||
|
||||
// Alignment
|
||||
constexpr ck::index_t kAAlignment = 32;
|
||||
constexpr ck::index_t kBAlignment = 32;
|
||||
constexpr ck::index_t kCAlignment = 32;
|
||||
|
||||
constexpr ck::index_t kBlockSize = 256;
|
||||
|
||||
constexpr ck::index_t kGemmMPerBlock = 256;
|
||||
constexpr ck::index_t kGemmNPerBlock = 128;
|
||||
constexpr ck::index_t kGemmKPerBlock = 32;
|
||||
|
||||
ck::index_t kGridSize = (M / kGemmMPerBlock) * (N / kGemmNPerBlock);
|
||||
|
||||
std::cout << "grid size " << kGridSize << std::endl;
|
||||
|
||||
constexpr ck::index_t kWarpPerCu = 8; // 2 warps per SIMD
|
||||
constexpr ck::index_t kWarpPerBlock = kBlockSize / warpSize;
|
||||
constexpr ck::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock;
|
||||
|
||||
const auto gemm_kernel = Gemm<ADataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
CDataType,
|
||||
ALayout,
|
||||
BLayout,
|
||||
CLayout,
|
||||
AElementFunction,
|
||||
BElementFunction,
|
||||
CElementFunction,
|
||||
kAAlignment,
|
||||
kBAlignment,
|
||||
kCAlignment,
|
||||
kBlockSize,
|
||||
kGemmMPerBlock,
|
||||
kGemmNPerBlock,
|
||||
kGemmKPerBlock>{};
|
||||
|
||||
float ave_time =
|
||||
launch_kernel<kBlockSize, kBlockPerCu>(StreamConfig{nullptr, true},
|
||||
gemm_kernel,
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0,
|
||||
static_cast<ADataType*>(a_buf.GetDeviceBuffer()),
|
||||
static_cast<BDataType*>(b_buf.GetDeviceBuffer()),
|
||||
static_cast<CDataType*>(c_buf.GetDeviceBuffer()),
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
Lda,
|
||||
Ldb,
|
||||
Ldc,
|
||||
AElementFunction{},
|
||||
BElementFunction{},
|
||||
CElementFunction{});
|
||||
auto pass = true;
|
||||
|
||||
if(verification)
|
||||
{
|
||||
// reference gemm
|
||||
Tensor<CDataType> c_host_ref(c_lengths, c_strides);
|
||||
reference_gemm<ADataType, ADataType, AccDataType, CDataType>(a_host, b_host, c_host_ref);
|
||||
c_buf.FromDevice(c_host_dev.mData.data());
|
||||
pass &= ck::utils::check_err(c_host_dev, c_host_ref);
|
||||
}
|
||||
|
||||
std::size_t flop = std::size_t(2) * M * N * K;
|
||||
std::size_t num_btype =
|
||||
sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N;
|
||||
|
||||
float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
|
||||
|
||||
float gb_per_sec = num_btype / 1.E6 / ave_time;
|
||||
|
||||
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
|
||||
<< std::endl;
|
||||
|
||||
return !pass;
|
||||
}
|
||||
171
example/91_tile_program/gemm.hpp
Normal file
171
example/91_tile_program/gemm.hpp
Normal file
@@ -0,0 +1,171 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/tensor_adaptor.hpp"
|
||||
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
|
||||
|
||||
#include "ck/tile_program/tile/tile_distribution.hpp"
|
||||
#include "ck/tile_program/tile/tile_elementwise.hpp"
|
||||
#include "ck/tile_program/tile/tile_gemm_shape.hpp"
|
||||
#include "ck/tile_program/warp_tile/warp_gemm.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v1.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp"
|
||||
#include "ck/tile_program/grid/grid_gemm_problem.hpp"
|
||||
#include "ck/tile_program/grid/grid_gemm_v1.hpp"
|
||||
#include "ck/tile_program/grid/grid_gemm_v1_default_policy.hpp"
|
||||
|
||||
// C = A * B
|
||||
template <typename ADataType,
|
||||
typename BDataType,
|
||||
typename AccDataType,
|
||||
typename CDataType,
|
||||
typename ALayout,
|
||||
typename BLayout,
|
||||
typename CLayout,
|
||||
typename AElementFunction,
|
||||
typename BElementFunction,
|
||||
typename CElementFunction,
|
||||
ck::index_t kAAlignment,
|
||||
ck::index_t kBAlignment,
|
||||
ck::index_t kCAlignment,
|
||||
ck::index_t kBlockSize_,
|
||||
ck::index_t kMPerBlock_,
|
||||
ck::index_t kNPerBlock_,
|
||||
ck::index_t kKPerBlock_>
|
||||
struct Gemm
|
||||
{
|
||||
using GridGemmProblem = ck::tile_program::grid::GridGemmProblem<ADataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
CDataType,
|
||||
AElementFunction,
|
||||
BElementFunction,
|
||||
CElementFunction>;
|
||||
|
||||
struct GridGemmPolicy
|
||||
{
|
||||
static constexpr ck::index_t kBlockSize = kBlockSize_;
|
||||
static constexpr ck::index_t kMPerBlock = kMPerBlock_;
|
||||
static constexpr ck::index_t kNPerBlock = kNPerBlock_;
|
||||
static constexpr ck::index_t kKPerBlock = kKPerBlock_;
|
||||
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto MakeBlock2TileMap(ck::index_t NumTilesM,
|
||||
ck::index_t NumTilesN)
|
||||
{
|
||||
using namespace ck;
|
||||
|
||||
const auto unmerge = make_merge_transform(make_tuple(NumTilesN, NumTilesM));
|
||||
|
||||
return [unmerge](index_t block_id) {
|
||||
MultiIndex<2> unmerged;
|
||||
unmerge.CalculateLowerIndex(unmerged, make_multi_index(block_id));
|
||||
|
||||
return make_multi_index(unmerged.At<1>(), unmerged.At<0>());
|
||||
};
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
__host__ __device__ static constexpr auto GetBlockGemmPipeline()
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
using namespace ck::tile_program::block;
|
||||
|
||||
using BlockGemmPipelineProblem_ =
|
||||
BlockGemmPipelineProblem<ADataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
kBlockSize,
|
||||
TileGemmShape<kMPerBlock, kNPerBlock, kKPerBlock>>;
|
||||
|
||||
return BlockGemmPipelineAGmemBGmemCRegV2<
|
||||
BlockGemmPipelineProblem_,
|
||||
BlockGemmPipelineAGmemBGmemCRegV2DefaultPolicy>{};
|
||||
}
|
||||
};
|
||||
|
||||
using GridGemm = ck::GridGemmV1<GridGemmProblem, GridGemmPolicy>;
|
||||
|
||||
__device__ void operator()(const ADataType* p_a,
|
||||
const BDataType* p_b,
|
||||
CDataType* p_c,
|
||||
const ck::index_t M,
|
||||
const ck::index_t N,
|
||||
const ck::index_t K,
|
||||
const ck::index_t Lda,
|
||||
const ck::index_t Ldb,
|
||||
const ck::index_t Ldc,
|
||||
const AElementFunction& a_element_func,
|
||||
const BElementFunction& b_element_func,
|
||||
const CElementFunction& c_element_func) const
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
using namespace ck::tile_program::block;
|
||||
|
||||
const auto a_dram = [&] {
|
||||
if constexpr(is_same_v<ALayout, ck::tensor_layout::gemm::RowMajor>)
|
||||
{
|
||||
return make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_a, make_tuple(M, K), make_tuple(Lda, 1), Number<kAAlignment>{}, Number<1>{});
|
||||
}
|
||||
else
|
||||
{
|
||||
const auto a_k_m_desc = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_a, make_tuple(K, M), make_tuple(Lda, 1), Number<kAAlignment>{}, Number<1>{});
|
||||
|
||||
return transform_tensor_view(
|
||||
a_k_m_desc,
|
||||
make_tuple(make_pass_through_transform(M), make_pass_through_transform(K)),
|
||||
make_tuple(Sequence<1>{}, Sequence<0>{}),
|
||||
make_tuple(Sequence<0>{}, Sequence<1>{}));
|
||||
}
|
||||
}();
|
||||
|
||||
const auto b_dram = [&] {
|
||||
if constexpr(is_same_v<BLayout, ck::tensor_layout::gemm::ColumnMajor>)
|
||||
{
|
||||
return make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_b, make_tuple(N, K), make_tuple(Ldb, 1), Number<kBAlignment>{}, Number<1>{});
|
||||
}
|
||||
else
|
||||
{
|
||||
const auto b_k_n_desc = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_b, make_tuple(K, N), make_tuple(Ldb, 1), Number<kBAlignment>{}, Number<1>{});
|
||||
|
||||
return transform_tensor_view(
|
||||
b_k_n_desc,
|
||||
make_tuple(make_pass_through_transform(N), make_pass_through_transform(K)),
|
||||
make_tuple(Sequence<1>{}, Sequence<0>{}),
|
||||
make_tuple(Sequence<0>{}, Sequence<1>{}));
|
||||
}
|
||||
}();
|
||||
|
||||
const auto c_dram = [&] {
|
||||
if constexpr(is_same_v<CLayout, ck::tensor_layout::gemm::RowMajor>)
|
||||
{
|
||||
return make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_c, make_tuple(M, N), make_tuple(Ldc, 1), Number<kCAlignment>{}, Number<1>{});
|
||||
}
|
||||
else
|
||||
{
|
||||
const auto c_n_m_desc = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_c, make_tuple(N, M), make_tuple(Ldc, 1), Number<kCAlignment>{}, Number<1>{});
|
||||
|
||||
return transform_tensor_view(
|
||||
c_n_m_desc,
|
||||
make_tuple(make_pass_through_transform(M), make_pass_through_transform(N)),
|
||||
make_tuple(Sequence<1>{}, Sequence<0>{}),
|
||||
make_tuple(Sequence<0>{}, Sequence<1>{}));
|
||||
}
|
||||
}();
|
||||
|
||||
GridGemm{}(a_dram, b_dram, c_dram, a_element_func, b_element_func, c_element_func);
|
||||
}
|
||||
};
|
||||
143
example/91_tile_program/gemm_gemm.cpp
Normal file
143
example/91_tile_program/gemm_gemm.cpp
Normal file
@@ -0,0 +1,143 @@
|
||||
#include <cstring>
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/cluster_descriptor.hpp"
|
||||
#include "ck/tensor/tensor_view.hpp"
|
||||
#include "ck/host_utility/device_prop.hpp"
|
||||
#include "ck/host_utility/kernel_launch.hpp"
|
||||
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/fill.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
|
||||
#include "reference_gemm.hpp"
|
||||
#include "gemm_gemm.hpp"
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
using A0DataType = ck::half_t;
|
||||
using B0DataType = ck::half_t;
|
||||
using B1DataType = ck::half_t;
|
||||
using Acc0DataType = float;
|
||||
using C0DataType = ck::half_t;
|
||||
using Acc1DataType = float;
|
||||
using C1DataType = ck::half_t;
|
||||
|
||||
ck::index_t M0 = 13312;
|
||||
ck::index_t N0 = 4096;
|
||||
ck::index_t K0 = 128;
|
||||
ck::index_t N1 = 128;
|
||||
|
||||
if(argc == 5)
|
||||
{
|
||||
M0 = std::stoi(argv[1]);
|
||||
N0 = std::stoi(argv[2]);
|
||||
K0 = std::stoi(argv[3]);
|
||||
N1 = std::stoi(argv[4]);
|
||||
}
|
||||
|
||||
std::array<ck::index_t, 2> a0_lengths{M0, K0};
|
||||
std::array<ck::index_t, 2> a0_strides{K0, 1};
|
||||
|
||||
std::array<ck::index_t, 2> b0_lengths{N0, K0};
|
||||
std::array<ck::index_t, 2> b0_strides{K0, 1};
|
||||
|
||||
std::array<ck::index_t, 2> c0_lengths{M0, N0};
|
||||
std::array<ck::index_t, 2> c0_strides{N0, 1};
|
||||
|
||||
std::array<ck::index_t, 2> b1_lengths{N1, N0};
|
||||
std::array<ck::index_t, 2> b1_strides{N0, 1};
|
||||
|
||||
std::array<ck::index_t, 2> c1_lengths{M0, N1};
|
||||
std::array<ck::index_t, 2> c1_strides{N1, 1};
|
||||
|
||||
// host verify
|
||||
Tensor<A0DataType> a0_host(a0_lengths, a0_strides);
|
||||
Tensor<B0DataType> b0_host(b0_lengths, b0_strides);
|
||||
Tensor<B1DataType> b1_host(b1_lengths, b1_strides);
|
||||
Tensor<C0DataType> c0_host_ref(c0_lengths, c0_strides);
|
||||
Tensor<C1DataType> c1_host_ref(c1_lengths, c1_strides);
|
||||
Tensor<C1DataType> c1_host_dev(c1_lengths, c1_strides);
|
||||
|
||||
ck::utils::FillUniformDistributionIntegerValue<A0DataType>{-3.f, 3.f}(a0_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<B0DataType>{-3.f, 3.f}(b0_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<B1DataType>{-3.f, 3.f}(b1_host);
|
||||
|
||||
// reference gemm
|
||||
reference_gemm<A0DataType, B0DataType, Acc0DataType, C0DataType>(a0_host, b0_host, c0_host_ref);
|
||||
reference_gemm<C0DataType, B1DataType, Acc1DataType, C1DataType>(
|
||||
c0_host_ref, b1_host, c1_host_ref);
|
||||
|
||||
DeviceMem a0_buf(sizeof(A0DataType) * a0_host.GetElementSpaceSize());
|
||||
DeviceMem b0_buf(sizeof(B0DataType) * b0_host.GetElementSpaceSize());
|
||||
DeviceMem b1_buf(sizeof(B1DataType) * b1_host.GetElementSpaceSize());
|
||||
DeviceMem c1_buf(sizeof(C1DataType) * c1_host_ref.GetElementSpaceSize());
|
||||
|
||||
a0_buf.ToDevice(a0_host.mData.data());
|
||||
b0_buf.ToDevice(b0_host.mData.data());
|
||||
b1_buf.ToDevice(b1_host.mData.data());
|
||||
|
||||
constexpr ck::index_t kM0PerBlock = 128;
|
||||
constexpr ck::index_t kN0PerBlock = 128;
|
||||
constexpr ck::index_t kK0PerBlock = 32;
|
||||
constexpr ck::index_t kN1PerBlock = 128;
|
||||
constexpr ck::index_t kK1PerBlock = 32;
|
||||
|
||||
constexpr ck::index_t kBlockSize = 256;
|
||||
ck::index_t kGridSize = (M0 / kM0PerBlock) * (N1 / kN1PerBlock);
|
||||
|
||||
std::cout << "grid size " << kGridSize << std::endl;
|
||||
|
||||
constexpr ck::index_t kWarpPerCu = 8; // 2 warps per SIMD
|
||||
constexpr ck::index_t kWarpPerBlock = kBlockSize / warpSize;
|
||||
constexpr ck::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock;
|
||||
|
||||
float ave_time =
|
||||
launch_kernel<kBlockSize, kBlockPerCu>(StreamConfig{nullptr, true},
|
||||
GemmGemm<A0DataType,
|
||||
B0DataType,
|
||||
B1DataType,
|
||||
Acc0DataType,
|
||||
C0DataType,
|
||||
Acc1DataType,
|
||||
C1DataType,
|
||||
kBlockSize,
|
||||
kM0PerBlock,
|
||||
kN0PerBlock,
|
||||
kK0PerBlock,
|
||||
kN1PerBlock,
|
||||
kK1PerBlock>{},
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0,
|
||||
static_cast<A0DataType*>(a0_buf.GetDeviceBuffer()),
|
||||
static_cast<B0DataType*>(b0_buf.GetDeviceBuffer()),
|
||||
static_cast<B1DataType*>(b1_buf.GetDeviceBuffer()),
|
||||
static_cast<C1DataType*>(c1_buf.GetDeviceBuffer()),
|
||||
M0,
|
||||
N0,
|
||||
K0,
|
||||
N1,
|
||||
K0, // Lda0
|
||||
K0, // Ldb0
|
||||
N0, // Ldb1
|
||||
N1); // Ldc1
|
||||
|
||||
c1_buf.FromDevice(c1_host_dev.mData.data());
|
||||
|
||||
std::size_t flop = std::size_t(2) * M0 * N0 * K0 + std::size_t(2) * M0 * N1 * N0;
|
||||
std::size_t num_btype = sizeof(A0DataType) * M0 * K0 + sizeof(B0DataType) * N0 * K0 +
|
||||
sizeof(B1DataType) * N1 * N0 + sizeof(C1DataType) * M0 * N1;
|
||||
|
||||
float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
|
||||
|
||||
float gb_per_sec = num_btype / 1.E6 / ave_time;
|
||||
|
||||
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
|
||||
<< std::endl;
|
||||
|
||||
return !ck::utils::check_err(c1_host_dev, c1_host_ref);
|
||||
}
|
||||
326
example/91_tile_program/gemm_gemm.hpp
Normal file
326
example/91_tile_program/gemm_gemm.hpp
Normal file
@@ -0,0 +1,326 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/tensor_adaptor.hpp"
|
||||
|
||||
#include "ck/tile_program/tile/tile_distribution.hpp"
|
||||
#include "ck/tile_program/tile/tile_elementwise.hpp"
|
||||
#include "ck/tile_program/tile/tile_gemm_shape.hpp"
|
||||
#include "ck/tile_program/tile/slice_tile.hpp"
|
||||
#include "ck/tile_program/warp_tile/warp_gemm.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_problem.hpp"
|
||||
#include "ck/tile_program/block_tile/block_gemm_areg_bsmem_creg_v1.hpp"
|
||||
|
||||
// C0 = A0 * B0
|
||||
// C1 = C0 * B1
|
||||
template <typename A0DataType,
|
||||
typename B0DataType,
|
||||
typename B1DataType,
|
||||
typename Acc0DataType,
|
||||
typename C0DataType,
|
||||
typename Acc1DataType,
|
||||
typename C1DataType,
|
||||
ck::index_t kBlockSize,
|
||||
ck::index_t kM0PerBlock,
|
||||
ck::index_t kN0PerBlock,
|
||||
ck::index_t kK0PerBlock,
|
||||
ck::index_t kN1PerBlock,
|
||||
ck::index_t kK1PerBlock>
|
||||
struct GemmGemm
|
||||
{
|
||||
static constexpr auto I0 = ck::Number<0>{};
|
||||
static constexpr auto BlockSize = ck::Number<kBlockSize>{};
|
||||
static constexpr auto M0PerBlock = ck::Number<kM0PerBlock>{};
|
||||
static constexpr auto N0PerBlock = ck::Number<kN0PerBlock>{};
|
||||
static constexpr auto K0PerBlock = ck::Number<kK0PerBlock>{};
|
||||
static constexpr auto N1PerBlock = ck::Number<kN1PerBlock>{};
|
||||
static constexpr auto K1PerBlock = ck::Number<kK1PerBlock>{};
|
||||
|
||||
// block gemm0 pipeline
|
||||
using BlockGemm0Pipeline = ck::tile_program::block::BlockGemmPipelineAGmemBGmemCRegV2<
|
||||
ck::tile_program::block::BlockGemmPipelineProblem<
|
||||
A0DataType,
|
||||
B0DataType,
|
||||
Acc0DataType,
|
||||
kBlockSize,
|
||||
ck::tile_program::TileGemmShape<kM0PerBlock, kN0PerBlock, kK0PerBlock>>,
|
||||
ck::tile_program::block::BlockGemmPipelineAGmemBGmemCRegV2DefaultPolicy>;
|
||||
|
||||
// block gemm1
|
||||
using BlockGemm1 = ck::tile_program::block::BlockGemmARegBSmemCRegV1<
|
||||
ck::tile_program::block::BlockGemmARegBSmemCRegV1Problem<
|
||||
C0DataType,
|
||||
B1DataType,
|
||||
Acc1DataType,
|
||||
kBlockSize,
|
||||
ck::tile_program::TileGemmShape<kM0PerBlock, kN1PerBlock, kK1PerBlock>>,
|
||||
ck::tile_program::block::BlockGemmARegBSmemCRegV1DefaultPolicy>;
|
||||
|
||||
#if 0
|
||||
// 2d
|
||||
__device__ static constexpr auto MakeB1LdsBlockDescriptor()
|
||||
{
|
||||
using namespace ck;
|
||||
|
||||
constexpr index_t kNPerBlock = kN1PerBlock;
|
||||
constexpr index_t kKPerBlock = kK1PerBlock;
|
||||
|
||||
constexpr auto b_lds_block_desc =
|
||||
make_naive_tensor_descriptor_packed(make_tuple(kNPerBlock, kKPerBlock), Number<32>{});
|
||||
|
||||
return b_lds_block_desc;
|
||||
}
|
||||
#elif 1
|
||||
// 3d, with padding
|
||||
__device__ static constexpr auto MakeB1LdsBlockDescriptor()
|
||||
{
|
||||
using namespace ck;
|
||||
|
||||
// using BDataType = B1DataType;
|
||||
|
||||
constexpr index_t kNPerBlock = kN1PerBlock;
|
||||
constexpr index_t kKPerBlock = kK1PerBlock;
|
||||
constexpr index_t kPad = 1;
|
||||
constexpr index_t kK1 = 8;
|
||||
|
||||
constexpr auto b_lds_block_desc_0 = make_naive_tensor_descriptor(
|
||||
make_tuple(Number<kKPerBlock / kK1>{}, Number<kNPerBlock>{}, Number<kK1>{}),
|
||||
make_tuple(Number<(kNPerBlock + kPad) * kK1>{}, Number<kK1>{}, Number<1>{}),
|
||||
Number<kK1>{},
|
||||
Number<1>{});
|
||||
|
||||
constexpr auto b_lds_block_desc = transform_tensor_descriptor(
|
||||
b_lds_block_desc_0,
|
||||
make_tuple(make_pass_through_transform(kNPerBlock),
|
||||
make_merge_transform(make_tuple(Number<kKPerBlock / kK1>{}, Number<kK1>{}))),
|
||||
make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
|
||||
make_tuple(Sequence<0>{}, Sequence<1>{}));
|
||||
|
||||
return b_lds_block_desc;
|
||||
}
|
||||
#else
|
||||
// fake XOR
|
||||
__host__ __device__ static constexpr auto MakeB1LdsBlockDescriptor()
|
||||
{
|
||||
using namespace ck;
|
||||
|
||||
using BDataType = B1DataType;
|
||||
|
||||
constexpr index_t kNPerBlock = kN1PerBlock;
|
||||
constexpr index_t kKPerBlock = kK1PerBlock;
|
||||
|
||||
constexpr auto b_lds_block_desc_d1_d2_d3 = make_naive_tensor_descriptor_packed(
|
||||
make_tuple(kNPerBlock / 2, 2, kKPerBlock), Number<kKPerBlock>{});
|
||||
|
||||
constexpr index_t kK1 = 16 / sizeof(BDataType);
|
||||
|
||||
constexpr auto b_lds_block_desc_d4_d5_d6 = transform_tensor_descriptor(
|
||||
b_lds_block_desc_d1_d2_d3,
|
||||
make_tuple(make_xor_transform(make_tuple(kNPerBlock / 2, kKPerBlock), kK1),
|
||||
make_pass_through_transform(2)),
|
||||
make_tuple(Sequence<0, 2>{}, Sequence<1>{}),
|
||||
make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
|
||||
|
||||
constexpr auto b_lds_block_desc_n_k = transform_tensor_descriptor(
|
||||
b_lds_block_desc_d4_d5_d6,
|
||||
make_tuple(make_merge_transform(make_tuple(kNPerBlock / 2, 2)),
|
||||
make_pass_through_transform(kKPerBlock)),
|
||||
make_tuple(Sequence<0, 1>{}, Sequence<2>{}),
|
||||
make_tuple(Sequence<0>{}, Sequence<1>{}));
|
||||
|
||||
return b_lds_block_desc_n_k;
|
||||
}
|
||||
#endif
|
||||
|
||||
__device__ static constexpr auto MakeB1DramTileDistribution()
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
|
||||
using BDataType = B1DataType;
|
||||
|
||||
constexpr index_t kNPerBlock = kN1PerBlock;
|
||||
constexpr index_t kKPerBlock = kK1PerBlock;
|
||||
|
||||
constexpr index_t K1 = 16 / sizeof(BDataType);
|
||||
constexpr index_t K0 = kKPerBlock / K1;
|
||||
constexpr index_t N2 = get_warp_size() / K0;
|
||||
constexpr index_t N1 = kBlockSize / get_warp_size();
|
||||
constexpr index_t N0 = kNPerBlock / (N2 * N1);
|
||||
|
||||
return make_static_tile_distribution(
|
||||
StaticTileDistributionEncoding<Sequence<1>,
|
||||
Tuple<Sequence<N0, N1, N2>, Sequence<K0, K1>>,
|
||||
Tuple<Sequence<1>, Sequence<1, 2>>,
|
||||
Tuple<Sequence<1>, Sequence<2, 0>>,
|
||||
Sequence<1, 2>,
|
||||
Sequence<0, 1>>{});
|
||||
}
|
||||
|
||||
__device__ static constexpr ck::index_t GetStaticLdsSize()
|
||||
{
|
||||
using namespace ck;
|
||||
|
||||
return math::max(BlockGemm0Pipeline::GetStaticLdsSize(),
|
||||
static_cast<index_t>(MakeB1LdsBlockDescriptor().GetElementSpaceSize() *
|
||||
sizeof(B1DataType)));
|
||||
}
|
||||
|
||||
__device__ void operator()(const A0DataType* p_a0,
|
||||
const B0DataType* p_b0,
|
||||
const B1DataType* p_b1,
|
||||
C1DataType* p_c1,
|
||||
const ck::index_t M0,
|
||||
const ck::index_t N0,
|
||||
const ck::index_t K0,
|
||||
const ck::index_t N1,
|
||||
const ck::index_t Lda0,
|
||||
const ck::index_t Ldb0,
|
||||
const ck::index_t Ldb1,
|
||||
const ck::index_t Ldc1)
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
using namespace ck::tile_program::block;
|
||||
|
||||
// FIXME: assume layout A0[M0, K0], B0[N0, K0], B1[N1, N0], C1[M0, N1]
|
||||
const auto a0_dram_grid = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_a0, make_tuple(M0, K0), make_tuple(Lda0, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
const auto b0_dram_grid = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_b0, make_tuple(N0, K0), make_tuple(Ldb0, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
const auto b1_dram_grid = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_b1, make_tuple(N1, N0), make_tuple(Ldb1, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
// divide problem
|
||||
const auto id_block = get_block_id();
|
||||
|
||||
const auto num_tile_m0 = M0 / kM0PerBlock;
|
||||
const auto num_tile_n1 = N1 / kN1PerBlock;
|
||||
|
||||
const auto block2tile = make_cluster_descriptor(make_tuple(num_tile_m0, num_tile_n1));
|
||||
|
||||
const auto id_tile = block2tile.CalculateBottomIndex(make_tuple(id_block));
|
||||
|
||||
const auto iM0 = __builtin_amdgcn_readfirstlane(id_tile.At<0>() * kM0PerBlock);
|
||||
const auto iN1 = __builtin_amdgcn_readfirstlane(id_tile.At<1>() * kN1PerBlock);
|
||||
|
||||
__shared__ char p_smem_char[GetStaticLdsSize()];
|
||||
|
||||
// A0 DRAM block window
|
||||
auto a0_dram_block_window =
|
||||
make_tile_window(a0_dram_grid, make_tuple(M0PerBlock, K0PerBlock), {iM0, 0});
|
||||
|
||||
// B0 DRAM block window
|
||||
auto b0_dram_block_window =
|
||||
make_tile_window(b0_dram_grid, make_tuple(N0PerBlock, K0PerBlock), {0, 0});
|
||||
|
||||
// Block GEMM0 pipeline
|
||||
constexpr auto block_gemm0_pipeline = BlockGemm0Pipeline{};
|
||||
|
||||
// B1 DRAM window
|
||||
auto b1_dram_block_window = make_tile_window(b1_dram_grid,
|
||||
make_tuple(N1PerBlock, K1PerBlock),
|
||||
{iN1, 0},
|
||||
MakeB1DramTileDistribution());
|
||||
|
||||
// B1 LDS tensor view: occupies the same LDS allocation as block_gemm0_pipeline
|
||||
auto b1_lds_block = make_tensor_view<AddressSpaceEnum::Lds>(
|
||||
reinterpret_cast<B1DataType*>(p_smem_char), MakeB1LdsBlockDescriptor());
|
||||
|
||||
auto b1_lds_block_window =
|
||||
make_tile_window(b1_lds_block, make_tuple(N1PerBlock, K1PerBlock), {0, 0});
|
||||
|
||||
// Bock GEMM1
|
||||
constexpr auto block_gemm1 = BlockGemm1{};
|
||||
|
||||
// Acc1 tile
|
||||
auto acc1_block_tile = decltype(block_gemm1(
|
||||
get_slice_tile(
|
||||
tile_elementwise_in(
|
||||
type_convert<C0DataType, Acc0DataType>,
|
||||
block_gemm0_pipeline(a0_dram_block_window, b0_dram_block_window, 0, nullptr)),
|
||||
Sequence<0, 0>{},
|
||||
Sequence<kM0PerBlock, kK1PerBlock>{}),
|
||||
b1_dram_block_window)){};
|
||||
|
||||
// init Acc1
|
||||
tile_elementwise_inout([](auto& acc1) { acc1 = 0; }, acc1_block_tile);
|
||||
|
||||
index_t iN0 = 0;
|
||||
|
||||
do
|
||||
{
|
||||
// Block GEMM0 pipeline: acc0 = a0 * b0
|
||||
const auto acc0_block_tile = block_gemm0_pipeline(
|
||||
a0_dram_block_window, b0_dram_block_window, K0 / kK0PerBlock, p_smem_char);
|
||||
|
||||
// type cast acc0 into c0
|
||||
const auto c0_block_tile =
|
||||
tile_elementwise_in(type_convert<C0DataType, Acc0DataType>, acc0_block_tile);
|
||||
|
||||
// prefetch load b1
|
||||
const auto b1_block_tile = load_tile(b1_dram_block_window);
|
||||
move_tile_window(b1_dram_block_window, {0, kK1PerBlock});
|
||||
|
||||
block_sync_lds();
|
||||
|
||||
store_tile(b1_lds_block_window, b1_block_tile);
|
||||
|
||||
constexpr index_t k1_loops = kN0PerBlock / kK1PerBlock;
|
||||
|
||||
if constexpr(k1_loops > 1)
|
||||
{
|
||||
static_for<0, k1_loops - 1, 1>{}([&](auto i) {
|
||||
// acc1 += c0 * b1
|
||||
const auto b1_block_tile_1 = load_tile(b1_dram_block_window);
|
||||
block_sync_lds();
|
||||
block_gemm1(acc1_block_tile,
|
||||
get_slice_tile(c0_block_tile,
|
||||
Sequence<0, i * kK1PerBlock>{},
|
||||
Sequence<kM0PerBlock, (i + 1) * kK1PerBlock>{}),
|
||||
b1_lds_block_window);
|
||||
block_sync_lds();
|
||||
move_tile_window(b1_dram_block_window, {0, kK1PerBlock});
|
||||
store_tile(b1_lds_block_window, b1_block_tile_1);
|
||||
});
|
||||
}
|
||||
// tail
|
||||
{
|
||||
block_sync_lds();
|
||||
block_gemm1(acc1_block_tile,
|
||||
get_slice_tile(c0_block_tile,
|
||||
Sequence<0, (k1_loops - 1) * kK1PerBlock>{},
|
||||
Sequence<kM0PerBlock, kN0PerBlock>{}),
|
||||
b1_lds_block_window);
|
||||
}
|
||||
|
||||
move_tile_window(b0_dram_block_window, {kN0PerBlock, 0});
|
||||
block_sync_lds();
|
||||
iN0 += kN0PerBlock;
|
||||
|
||||
} while(iN0 < N0);
|
||||
|
||||
// type cast acc1 into c1
|
||||
const auto c1_block_tile =
|
||||
tile_elementwise_in(type_convert<C1DataType, Acc1DataType>, acc1_block_tile);
|
||||
|
||||
// store c1
|
||||
auto c1_dram_grid = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_c1, make_tuple(M0, N1), make_tuple(Ldc1, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
auto c1_dram_window = make_tile_window(c1_dram_grid,
|
||||
make_tuple(M0PerBlock, N1PerBlock),
|
||||
{iM0, iN1},
|
||||
c1_block_tile.GetTileDistribution());
|
||||
|
||||
store_tile(c1_dram_window, c1_block_tile);
|
||||
}
|
||||
};
|
||||
155
example/91_tile_program/gemm_softmax_gemm.cpp
Normal file
155
example/91_tile_program/gemm_softmax_gemm.cpp
Normal file
@@ -0,0 +1,155 @@
|
||||
#include <cstring>
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/cluster_descriptor.hpp"
|
||||
#include "ck/tensor/tensor_view.hpp"
|
||||
#include "ck/host_utility/device_prop.hpp"
|
||||
#include "ck/host_utility/kernel_launch.hpp"
|
||||
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/fill.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
|
||||
#include "reference_gemm.hpp"
|
||||
#include "reference_softmax.hpp"
|
||||
#include "gemm_softmax_gemm.hpp"
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
using QDataType = ck::half_t;
|
||||
using KDataType = ck::half_t;
|
||||
using VDataType = ck::half_t;
|
||||
using SaccDataType = float;
|
||||
using SMPLComputeDataType = float;
|
||||
using PDataType = ck::half_t;
|
||||
using OaccDataType = float;
|
||||
using ODataType = ck::half_t;
|
||||
|
||||
ck::index_t M0 = 13312;
|
||||
ck::index_t N0 = 4096;
|
||||
ck::index_t K0 = 128;
|
||||
ck::index_t N1 = 128;
|
||||
|
||||
if(argc == 5)
|
||||
{
|
||||
M0 = std::stoi(argv[1]);
|
||||
N0 = std::stoi(argv[2]);
|
||||
K0 = std::stoi(argv[3]);
|
||||
N1 = std::stoi(argv[4]);
|
||||
}
|
||||
|
||||
std::array<ck::index_t, 2> q_lengths{M0, K0};
|
||||
std::array<ck::index_t, 2> q_strides{K0, 1};
|
||||
|
||||
std::array<ck::index_t, 2> k_lengths{N0, K0};
|
||||
std::array<ck::index_t, 2> k_strides{K0, 1};
|
||||
|
||||
std::array<ck::index_t, 2> v_lengths{N1, N0};
|
||||
std::array<ck::index_t, 2> v_strides{N0, 1};
|
||||
|
||||
std::array<ck::index_t, 2> s_lengths{M0, N0};
|
||||
std::array<ck::index_t, 2> s_strides{N0, 1};
|
||||
|
||||
std::array<ck::index_t, 2> p_lengths{M0, N0};
|
||||
std::array<ck::index_t, 2> p_strides{N0, 1};
|
||||
|
||||
std::array<ck::index_t, 2> o_lengths{M0, N1};
|
||||
std::array<ck::index_t, 2> o_strides{N1, 1};
|
||||
|
||||
// host verify
|
||||
Tensor<QDataType> q_host(q_lengths, q_strides);
|
||||
Tensor<KDataType> k_host(k_lengths, k_strides);
|
||||
Tensor<VDataType> v_host(v_lengths, v_strides);
|
||||
Tensor<SMPLComputeDataType> s_host_ref(s_lengths, s_strides);
|
||||
Tensor<PDataType> p_host_ref(p_lengths, p_strides);
|
||||
Tensor<ODataType> o_host_ref(o_lengths, o_strides);
|
||||
Tensor<ODataType> o_host_dev(o_lengths, o_strides);
|
||||
|
||||
#if 0
|
||||
ck::utils::FillUniformDistributionIntegerValue<QDataType>{-3.f, 3.f}(q_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<KDataType>{-3.f, 3.f}(k_host);
|
||||
ck::utils::FillUniformDistributionIntegerValue<VDataType>{-3.f, 3.f}(v_host);
|
||||
#else
|
||||
ck::utils::FillUniformDistribution<QDataType>{-3.f, 3.f}(q_host);
|
||||
ck::utils::FillUniformDistribution<KDataType>{-3.f, 3.f}(k_host);
|
||||
ck::utils::FillUniformDistribution<VDataType>{-3.f, 3.f}(v_host);
|
||||
#endif
|
||||
|
||||
// reference
|
||||
reference_gemm<QDataType, KDataType, SaccDataType, SMPLComputeDataType>(
|
||||
q_host, k_host, s_host_ref);
|
||||
reference_softmax<SMPLComputeDataType, SMPLComputeDataType, PDataType>(s_host_ref, p_host_ref);
|
||||
reference_gemm<PDataType, VDataType, OaccDataType, ODataType>(p_host_ref, v_host, o_host_ref);
|
||||
|
||||
DeviceMem q_buf(sizeof(QDataType) * q_host.GetElementSpaceSize());
|
||||
DeviceMem k_buf(sizeof(KDataType) * k_host.GetElementSpaceSize());
|
||||
DeviceMem v_buf(sizeof(VDataType) * v_host.GetElementSpaceSize());
|
||||
DeviceMem o_buf(sizeof(ODataType) * o_host_ref.GetElementSpaceSize());
|
||||
|
||||
q_buf.ToDevice(q_host.mData.data());
|
||||
k_buf.ToDevice(k_host.mData.data());
|
||||
v_buf.ToDevice(v_host.mData.data());
|
||||
|
||||
constexpr ck::index_t kM0PerBlock = 128;
|
||||
constexpr ck::index_t kN0PerBlock = 128;
|
||||
constexpr ck::index_t kK0PerBlock = 32;
|
||||
constexpr ck::index_t kN1PerBlock = 128;
|
||||
|
||||
constexpr ck::index_t kBlockSize = 256;
|
||||
ck::index_t kGridSize = (M0 / kM0PerBlock) * (N1 / kN1PerBlock);
|
||||
|
||||
std::cout << "grid size " << kGridSize << std::endl;
|
||||
|
||||
constexpr ck::index_t kWarpPerCu = 8; // 2 warps per SIMD
|
||||
constexpr ck::index_t kWarpPerBlock = kBlockSize / warpSize;
|
||||
constexpr ck::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock;
|
||||
|
||||
float ave_time =
|
||||
launch_kernel<kBlockSize, kBlockPerCu>(StreamConfig{nullptr, true},
|
||||
GemmSoftmaxGemm<QDataType,
|
||||
KDataType,
|
||||
VDataType,
|
||||
SaccDataType,
|
||||
SMPLComputeDataType,
|
||||
PDataType,
|
||||
OaccDataType,
|
||||
ODataType,
|
||||
kBlockSize,
|
||||
kM0PerBlock,
|
||||
kN0PerBlock,
|
||||
kK0PerBlock,
|
||||
kN1PerBlock>{},
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0,
|
||||
static_cast<QDataType*>(q_buf.GetDeviceBuffer()),
|
||||
static_cast<KDataType*>(k_buf.GetDeviceBuffer()),
|
||||
static_cast<VDataType*>(v_buf.GetDeviceBuffer()),
|
||||
static_cast<ODataType*>(o_buf.GetDeviceBuffer()),
|
||||
M0,
|
||||
N0,
|
||||
K0,
|
||||
N1,
|
||||
K0, // StrideQ
|
||||
K0, // StrideK
|
||||
N0, // StrideV
|
||||
N1); // StrideO
|
||||
|
||||
o_buf.FromDevice(o_host_dev.mData.data());
|
||||
|
||||
std::size_t flop = std::size_t(2) * M0 * N0 * K0 + std::size_t(2) * M0 * N1 * N0;
|
||||
std::size_t num_btype = sizeof(QDataType) * M0 * K0 + sizeof(KDataType) * N0 * K0 +
|
||||
sizeof(VDataType) * N1 * N0 + sizeof(ODataType) * M0 * N1;
|
||||
|
||||
float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
|
||||
|
||||
float gb_per_sec = num_btype / 1.E6 / ave_time;
|
||||
|
||||
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
|
||||
<< std::endl;
|
||||
|
||||
return !ck::utils::check_err(o_host_dev, o_host_ref);
|
||||
}
|
||||
95
example/91_tile_program/gemm_softmax_gemm.hpp
Normal file
95
example/91_tile_program/gemm_softmax_gemm.hpp
Normal file
@@ -0,0 +1,95 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/tensor_adaptor.hpp"
|
||||
|
||||
#include "ck/tile_program/tile/tile_distribution.hpp"
|
||||
#include "ck/tile_program/tile/tile_elementwise.hpp"
|
||||
#include "ck/tile_program/tile/tile_gemm_shape.hpp"
|
||||
#include "ck/tile_program/warp_tile/warp_gemm.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_problem.hpp"
|
||||
#include "ck/tile_program/block_tile/block_gemm_areg_bsmem_creg_v1.hpp"
|
||||
#include "ck/tile_program/block_tile/block_reduce.hpp"
|
||||
|
||||
#include "gemm_softmax_gemm_impl.hpp"
|
||||
|
||||
// S[M0, N0] = Q[M0, K0] * K[N0, K0]
|
||||
// P[M0, N0] = Softmax(S[M0, N0])
|
||||
// O[M0, N1] = P[M0, N0] * V[N1, N0]
|
||||
template <typename QDataType,
|
||||
typename KDataType,
|
||||
typename VDataType,
|
||||
typename SaccDataType,
|
||||
typename SMPLComputeDataType,
|
||||
typename PDataType,
|
||||
typename OaccDataType,
|
||||
typename ODataType,
|
||||
ck::index_t kBlockSize,
|
||||
ck::index_t kM0PerBlock,
|
||||
ck::index_t kN0PerBlock,
|
||||
ck::index_t kK0PerBlock,
|
||||
ck::index_t kN1PerBlock>
|
||||
struct GemmSoftmaxGemm
|
||||
{
|
||||
__device__ void operator()(const QDataType* q_ptr,
|
||||
const KDataType* k_ptr,
|
||||
const VDataType* v_ptr,
|
||||
ODataType* o_ptr,
|
||||
const ck::index_t M0,
|
||||
const ck::index_t N0,
|
||||
const ck::index_t K0,
|
||||
const ck::index_t N1,
|
||||
const ck::index_t StrideQ,
|
||||
const ck::index_t StrideK,
|
||||
const ck::index_t StrideV,
|
||||
const ck::index_t StrideO) const
|
||||
{
|
||||
using namespace ck;
|
||||
|
||||
// divide problem
|
||||
const auto num_tile_n1 = N1 / kN1PerBlock;
|
||||
|
||||
const auto id_block = get_block_id();
|
||||
|
||||
const auto id_tile_m = id_block / num_tile_n1;
|
||||
const auto id_tile_n = id_block - id_tile_m * num_tile_n1;
|
||||
|
||||
const auto iM0 = __builtin_amdgcn_readfirstlane(id_tile_m * kM0PerBlock);
|
||||
const auto iN1 = __builtin_amdgcn_readfirstlane(id_tile_n * kN1PerBlock);
|
||||
|
||||
const auto kernel_impl = GemmSoftmaxGemmImpl<QDataType,
|
||||
KDataType,
|
||||
VDataType,
|
||||
SaccDataType,
|
||||
SMPLComputeDataType,
|
||||
PDataType,
|
||||
OaccDataType,
|
||||
ODataType,
|
||||
kBlockSize,
|
||||
kM0PerBlock,
|
||||
kN0PerBlock,
|
||||
kK0PerBlock,
|
||||
kN1PerBlock>{};
|
||||
|
||||
kernel_impl(q_ptr,
|
||||
k_ptr,
|
||||
v_ptr,
|
||||
o_ptr,
|
||||
M0,
|
||||
N0,
|
||||
K0,
|
||||
N1,
|
||||
StrideQ,
|
||||
StrideK,
|
||||
StrideV,
|
||||
StrideO,
|
||||
iM0,
|
||||
iN1);
|
||||
}
|
||||
};
|
||||
263
example/91_tile_program/gemm_softmax_gemm_impl.hpp
Normal file
263
example/91_tile_program/gemm_softmax_gemm_impl.hpp
Normal file
@@ -0,0 +1,263 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/tensor_adaptor.hpp"
|
||||
|
||||
#include "ck/tile_program/tile/tile_distribution.hpp"
|
||||
#include "ck/tile_program/tile/tile_gemm_shape.hpp"
|
||||
#include "ck/tile_program/tile/tile_elementwise.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_problem.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp"
|
||||
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp"
|
||||
#include "ck/tile_program/block_tile/block_gemm_areg_bgmem_creg_problem.hpp"
|
||||
#include "ck/tile_program/block_tile/block_gemm_areg_bgmem_creg_v1.hpp"
|
||||
#include "ck/tile_program/block_tile/block_gemm_areg_bsmem_creg_v1_default_policy.hpp"
|
||||
#include "ck/tile_program/block_tile/block_reduce.hpp"
|
||||
|
||||
// S[M0, N0] = Q[M0, K0] * K[N0, K0]
|
||||
// P[M0, N0] = Softmax(S[M0, N0])
|
||||
// O[M0, N1] = P[M0, N0] * V[N1, N0]
|
||||
template <typename QDataType,
|
||||
typename KDataType,
|
||||
typename VDataType,
|
||||
typename SaccDataType,
|
||||
typename SMPLComputeDataType,
|
||||
typename PDataType,
|
||||
typename OaccDataType,
|
||||
typename ODataType,
|
||||
ck::index_t kBlockSize,
|
||||
ck::index_t kM0PerBlock,
|
||||
ck::index_t kN0PerBlock,
|
||||
ck::index_t kK0PerBlock,
|
||||
ck::index_t kN1PerBlock>
|
||||
struct GemmSoftmaxGemmImpl
|
||||
{
|
||||
// block gemm0 pipeline
|
||||
using BlockGemm0Pipeline = ck::tile_program::block::BlockGemmPipelineAGmemBGmemCRegV2<
|
||||
ck::tile_program::block::BlockGemmPipelineProblem<
|
||||
QDataType,
|
||||
KDataType,
|
||||
SaccDataType,
|
||||
kBlockSize,
|
||||
ck::tile_program::TileGemmShape<kM0PerBlock, kN0PerBlock, kK0PerBlock>>,
|
||||
ck::tile_program::block::BlockGemmPipelineAGmemBGmemCRegV2DefaultPolicy>;
|
||||
|
||||
// block gemm1
|
||||
using BlockGemm1 = ck::tile_program::block::BlockGemmARegBGmemCRegV1<
|
||||
ck::tile_program::block::BlockGemmARegBGmemCRegProblem<
|
||||
PDataType,
|
||||
VDataType,
|
||||
OaccDataType,
|
||||
kBlockSize,
|
||||
ck::tile_program::TileGemmShape<kM0PerBlock, kN1PerBlock, kN0PerBlock>>,
|
||||
ck::tile_program::block::BlockGemmARegBGmemCRegV1DefaultPolicy>;
|
||||
|
||||
__device__ static constexpr ck::index_t GetStaticLdsSize()
|
||||
{
|
||||
return ck::math::max(BlockGemm0Pipeline::GetStaticLdsSize(),
|
||||
BlockGemm1::GetStaticLdsSize());
|
||||
}
|
||||
|
||||
__device__ void operator()(const QDataType* q_ptr,
|
||||
const KDataType* k_ptr,
|
||||
const VDataType* v_ptr,
|
||||
ODataType* o_ptr,
|
||||
const ck::index_t M0,
|
||||
const ck::index_t N0,
|
||||
const ck::index_t K0,
|
||||
const ck::index_t N1,
|
||||
const ck::index_t StrideQ,
|
||||
const ck::index_t StrideK,
|
||||
const ck::index_t StrideV,
|
||||
const ck::index_t StrideO,
|
||||
const ck::index_t iM0,
|
||||
const ck::index_t iN1) const
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
using namespace ck::tile_program::block;
|
||||
|
||||
constexpr auto I0 = Number<0>{};
|
||||
constexpr auto I1 = Number<1>{};
|
||||
|
||||
// allocate LDS
|
||||
__shared__ char smem_ptr[GetStaticLdsSize()];
|
||||
|
||||
// Q/K/V DRAM
|
||||
// FIXME: assume layout Q[M0, K0], K[N0, K0], V[N1, N0], O[M0, N1]
|
||||
const auto q_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
q_ptr, make_tuple(M0, K0), make_tuple(StrideQ, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
const auto k_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
k_ptr, make_tuple(N0, K0), make_tuple(StrideK, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
const auto v_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
v_ptr, make_tuple(N1, N0), make_tuple(StrideV, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
// Q/K/V DRAM window
|
||||
auto q_dram_window = make_tile_window(
|
||||
q_dram, make_tuple(Number<kM0PerBlock>{}, Number<kK0PerBlock>{}), {iM0, 0});
|
||||
|
||||
auto k_dram_window = make_tile_window(
|
||||
k_dram, make_tuple(Number<kN0PerBlock>{}, Number<kK0PerBlock>{}), {0, 0});
|
||||
|
||||
auto v_dram_window = make_tile_window(
|
||||
v_dram, make_tuple(Number<kN1PerBlock>{}, Number<kN0PerBlock>{}), {iN1, 0});
|
||||
|
||||
// Block GEMM0 pipeline and Block GEMM1
|
||||
constexpr auto gemm0_pipeline = BlockGemm0Pipeline{};
|
||||
constexpr auto gemm1 = BlockGemm1{};
|
||||
|
||||
// reduction function for softmax
|
||||
const auto f_max = [](auto e0, auto e1) { return max(e0, e1); };
|
||||
const auto f_sum = [](auto e0, auto e1) { return e0 + e1; };
|
||||
|
||||
// infer Sacc, S, P, M, L, Oacc type
|
||||
using SaccBlockTileType =
|
||||
decltype(gemm0_pipeline(q_dram_window, k_dram_window, 0, nullptr));
|
||||
|
||||
using SBlockTileType = decltype(tile_elementwise_in(
|
||||
type_convert<SMPLComputeDataType, SaccDataType>, SaccBlockTileType{}));
|
||||
|
||||
using PBlockTileType = decltype(tile_elementwise_in(type_convert<PDataType, SaccDataType>,
|
||||
SaccBlockTileType{}));
|
||||
|
||||
using MLBlockTileType = decltype(block_tile_reduce<SMPLComputeDataType>(
|
||||
SBlockTileType{}, Sequence<1>{}, f_max, SMPLComputeDataType{0}));
|
||||
|
||||
using OaccBlockTileType = decltype(gemm1(PBlockTileType{}, v_dram_window, smem_ptr));
|
||||
|
||||
// init Oacc, M, L
|
||||
auto o_acc = OaccBlockTileType{};
|
||||
auto m = MLBlockTileType{};
|
||||
auto l = MLBlockTileType{};
|
||||
|
||||
tile_elementwise_inout([](auto& e) { e = 0; }, o_acc);
|
||||
tile_elementwise_inout([](auto& e) { e = NumericLimits<SMPLComputeDataType>::Lowest(); },
|
||||
m);
|
||||
tile_elementwise_inout([](auto& e) { e = 0; }, l);
|
||||
|
||||
// loop over Column of S (J loop)
|
||||
index_t iN0 = 0;
|
||||
|
||||
do
|
||||
{
|
||||
// Sacc{j} = Q * K{j}
|
||||
const auto s_acc =
|
||||
gemm0_pipeline(q_dram_window, k_dram_window, K0 / kK0PerBlock, smem_ptr);
|
||||
|
||||
// S{j}
|
||||
const auto s =
|
||||
tile_elementwise_in(type_convert<SMPLComputeDataType, SaccDataType>, s_acc);
|
||||
|
||||
// m_local = rowmax(S{j})
|
||||
auto m_local = block_tile_reduce<SMPLComputeDataType>(
|
||||
s, Sequence<1>{}, f_max, NumericLimits<SMPLComputeDataType>::Lowest());
|
||||
|
||||
block_tile_reduce_sync(m_local, f_max);
|
||||
|
||||
// m{j-1}
|
||||
const auto m_old = m;
|
||||
|
||||
// m{j}
|
||||
tile_elementwise_inout(
|
||||
[](auto& e0, auto e1, auto e2) { e0 = max(e1, e2); }, m, m_old, m_local);
|
||||
|
||||
// Pcompute{j}
|
||||
auto p_compute =
|
||||
make_static_distributed_tensor<SMPLComputeDataType>(s.GetTileDistribution());
|
||||
|
||||
constexpr auto p_spans = decltype(p_compute)::GetDistributedSpans();
|
||||
|
||||
sweep_tile_span(p_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
sweep_tile_span(p_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
p_compute(i_j_idx) = math::exp(s[i_j_idx] - m[i_idx]);
|
||||
});
|
||||
});
|
||||
|
||||
// rowsum(Pcompute{j})
|
||||
auto rowsum_p = block_tile_reduce<SMPLComputeDataType>(
|
||||
p_compute, Sequence<1>{}, f_sum, SMPLComputeDataType{0});
|
||||
|
||||
block_tile_reduce_sync(rowsum_p, f_sum);
|
||||
|
||||
// l{j}, Oacc{j}
|
||||
sweep_tile_span(p_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
const auto tmp = math::exp(m_old[i_idx] - m[i_idx]);
|
||||
|
||||
l(i_idx) = tmp * l[i_idx] + rowsum_p[i_idx];
|
||||
|
||||
sweep_tile_span(p_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
// FIXME: this use different equation from FA v2 paper,
|
||||
// but produce correct result.
|
||||
// Is the equation wrong?
|
||||
o_acc(i_j_idx) *= tmp;
|
||||
});
|
||||
});
|
||||
|
||||
// type cast Pcompute{j} into P{j}
|
||||
const auto p =
|
||||
tile_elementwise_in(type_convert<PDataType, SMPLComputeDataType>, p_compute);
|
||||
|
||||
// wait for gemm0 pipeline to finish reading Lds
|
||||
block_sync_lds();
|
||||
|
||||
// Block GEMM1: Oacc{j} += P{j} * V{j}
|
||||
gemm1(o_acc, p, v_dram_window, smem_ptr);
|
||||
|
||||
// move K/V tile windows for next iteration (J loop)
|
||||
move_tile_window(k_dram_window, {kN0PerBlock, 0});
|
||||
move_tile_window(v_dram_window, {0, kN0PerBlock});
|
||||
|
||||
// wait for gemm1 to finish reading Lds, before next iteration (J loop)
|
||||
block_sync_lds();
|
||||
|
||||
iN0 += kN0PerBlock;
|
||||
|
||||
} while(iN0 < N0);
|
||||
|
||||
// O
|
||||
constexpr auto o_spans = decltype(o_acc)::GetDistributedSpans();
|
||||
|
||||
sweep_tile_span(o_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
const auto tmp = 1 / l[i_idx];
|
||||
|
||||
sweep_tile_span(o_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
o_acc(i_j_idx) *= tmp;
|
||||
});
|
||||
});
|
||||
|
||||
// type cast Oacc into O
|
||||
const auto o = tile_elementwise_in(type_convert<ODataType, OaccDataType>, o_acc);
|
||||
|
||||
// O DRAM and O DRAM window
|
||||
auto o_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
o_ptr, make_tuple(M0, N1), make_tuple(StrideO, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
auto o_dram_window =
|
||||
make_tile_window(o_dram,
|
||||
make_tuple(Number<kM0PerBlock>{}, Number<kN1PerBlock>{}),
|
||||
{iM0, iN1},
|
||||
o.GetTileDistribution());
|
||||
|
||||
// store O
|
||||
store_tile(o_dram_window, o);
|
||||
}
|
||||
};
|
||||
356
example/91_tile_program/im2col.cpp
Normal file
356
example/91_tile_program/im2col.cpp
Normal file
@@ -0,0 +1,356 @@
|
||||
#include <string_view>
|
||||
#include <tuple>
|
||||
#include <array>
|
||||
#include <utility>
|
||||
#include <type_traits>
|
||||
#include <cstring>
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/cluster_descriptor.hpp"
|
||||
#include "ck/tensor/tensor_view.hpp"
|
||||
#include "ck/host_utility/device_prop.hpp"
|
||||
#include "ck/host_utility/kernel_launch.hpp"
|
||||
|
||||
#include "ck/tile_program/tile/tile_distribution.hpp"
|
||||
#include "ck/tile_program/tile/tile_window.hpp"
|
||||
#include "ck/tile_program/tile/load_tile.hpp"
|
||||
#include "ck/tile_program/tile/store_tile.hpp"
|
||||
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/fill.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
#include "ck/library/utility/literals.hpp"
|
||||
|
||||
template <typename T>
|
||||
void reference_im2col(Tensor<T>& in_mtx_host_ref,
|
||||
const Tensor<T>& in_host,
|
||||
int /*N*/,
|
||||
int /*K*/,
|
||||
int C,
|
||||
int /*Y*/,
|
||||
int X,
|
||||
int Hi,
|
||||
int Wi,
|
||||
int Ho,
|
||||
int Wo,
|
||||
int ConvStrideH,
|
||||
int ConvStrideW,
|
||||
int ConvDilationH,
|
||||
int ConvDilationW,
|
||||
int InLeftPadH,
|
||||
int InLeftPadW,
|
||||
int /*InRightPadH*/,
|
||||
int /*InRightPadW*/)
|
||||
{
|
||||
int GemmM = in_mtx_host_ref.GetLengths()[0];
|
||||
int GemmK = in_mtx_host_ref.GetLengths()[1];
|
||||
|
||||
for(int gemm_m = 0; gemm_m < GemmM; ++gemm_m)
|
||||
{
|
||||
int mtmp = gemm_m;
|
||||
int n = mtmp / (Ho * Wo);
|
||||
mtmp -= n * Ho * Wo;
|
||||
int ho = mtmp / Wo;
|
||||
int wo = mtmp - ho * Wo;
|
||||
|
||||
for(int gemm_k = 0; gemm_k < GemmK; ++gemm_k)
|
||||
{
|
||||
int ktmp = gemm_k;
|
||||
int y = ktmp / (X * C);
|
||||
ktmp -= y * X * C;
|
||||
int x = ktmp / C;
|
||||
int c = ktmp - x * C;
|
||||
|
||||
int hi = y * ConvDilationH + ho * ConvStrideH - InLeftPadH;
|
||||
int wi = x * ConvDilationW + wo * ConvStrideW - InLeftPadW;
|
||||
|
||||
bool inbound = (hi >= 0 && hi < Hi && wi >= 0 && wi < Wi);
|
||||
|
||||
in_mtx_host_ref(gemm_m, gemm_k) = inbound ? in_host(n, hi, wi, c) : 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <ck::index_t NDimSpatial,
|
||||
typename T,
|
||||
ck::index_t kBlockSize,
|
||||
ck::index_t kMPerBlock,
|
||||
ck::index_t kKPerBlock>
|
||||
struct Im2Col
|
||||
{
|
||||
__host__ __device__ static constexpr auto MakeBlockCopyTileDistribution()
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
|
||||
constexpr index_t NumWarp = kBlockSize / get_warp_size();
|
||||
|
||||
constexpr index_t K1 = 16 / sizeof(T);
|
||||
constexpr index_t K0 = kKPerBlock / K1;
|
||||
|
||||
constexpr index_t M2 = get_warp_size() / K0;
|
||||
constexpr index_t M1 = NumWarp;
|
||||
constexpr index_t M0 = kMPerBlock / (M1 * M2);
|
||||
|
||||
return make_static_tile_distribution(
|
||||
StaticTileDistributionEncoding<Sequence<>,
|
||||
Tuple<Sequence<M0, M1, M2>, Sequence<K0, K1>>,
|
||||
Tuple<Sequence<1>, Sequence<1, 2>>,
|
||||
Tuple<Sequence<1>, Sequence<2, 0>>,
|
||||
Sequence<1, 2>,
|
||||
Sequence<0, 1>>{});
|
||||
}
|
||||
|
||||
__host__ __device__ void
|
||||
operator()(const std::array<ck::index_t, NDimSpatial + 2>& a_n_wis_c_lengths,
|
||||
const std::array<ck::index_t, NDimSpatial + 2>& /* a_n_wis_c_strides */,
|
||||
const std::array<ck::index_t, NDimSpatial + 2>& b_k_xs_c_lengths,
|
||||
const std::array<ck::index_t, NDimSpatial + 2>& /* b_k_xs_c_strides */,
|
||||
const std::array<ck::index_t, NDimSpatial + 2>& c_n_wos_k_lengths,
|
||||
const std::array<ck::index_t, NDimSpatial + 2>& /* c_n_wos_k_strides */,
|
||||
const std::array<ck::index_t, NDimSpatial>& conv_filter_strides,
|
||||
const std::array<ck::index_t, NDimSpatial>& conv_filter_dilations,
|
||||
const std::array<ck::index_t, NDimSpatial>& input_left_pads,
|
||||
const std::array<ck::index_t, NDimSpatial>& input_right_pads,
|
||||
const std::array<ck::index_t, 2> a_gemmm_gemmk_lengths,
|
||||
const std::array<ck::index_t, 2> a_gemmm_gemmk_strides,
|
||||
const T* p_a_img,
|
||||
T* p_a_mtx)
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
|
||||
const index_t N = a_n_wis_c_lengths[0];
|
||||
const index_t C = a_n_wis_c_lengths[3];
|
||||
|
||||
const index_t Hi = a_n_wis_c_lengths[1];
|
||||
const index_t Wi = a_n_wis_c_lengths[2];
|
||||
|
||||
const index_t Ho = c_n_wos_k_lengths[1];
|
||||
const index_t Wo = c_n_wos_k_lengths[2];
|
||||
|
||||
const index_t Y = b_k_xs_c_lengths[1];
|
||||
const index_t X = b_k_xs_c_lengths[2];
|
||||
|
||||
const index_t ConvStrideH = conv_filter_strides[0];
|
||||
const index_t ConvStrideW = conv_filter_strides[1];
|
||||
|
||||
const index_t ConvDilationH = conv_filter_dilations[0];
|
||||
const index_t ConvDilationW = conv_filter_dilations[1];
|
||||
|
||||
const index_t InLeftPadH = input_left_pads[0];
|
||||
const index_t InLeftPadW = input_left_pads[1];
|
||||
|
||||
const index_t InRightPadH = input_right_pads[0];
|
||||
const index_t InRightPadW = input_right_pads[1];
|
||||
|
||||
const auto a_n_hi_wi_c = make_naive_tensor_view_packed<AddressSpaceEnum::Global>(
|
||||
p_a_img, make_tuple(N, Hi, Wi, C), Number<32>{});
|
||||
|
||||
const auto a_n_hip_wip_c = transform_tensor_view(
|
||||
a_n_hi_wi_c,
|
||||
make_tuple(make_pass_through_transform(N),
|
||||
make_pad_transform(Hi, InLeftPadH, InRightPadH),
|
||||
make_pad_transform(Wi, InLeftPadW, InRightPadW),
|
||||
make_pass_through_transform(C)),
|
||||
make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
|
||||
make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
|
||||
|
||||
const auto a_n_y_ho_x_wo_c = transform_tensor_view(
|
||||
a_n_hip_wip_c,
|
||||
make_tuple(
|
||||
make_pass_through_transform(N),
|
||||
make_embed_transform(make_tuple(Y, Ho), make_tuple(ConvDilationH, ConvStrideH)),
|
||||
make_embed_transform(make_tuple(X, Wo), make_tuple(ConvDilationW, ConvStrideW)),
|
||||
make_pass_through_transform(C)),
|
||||
make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
|
||||
make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
|
||||
|
||||
const auto src_gemmm_gemmk =
|
||||
transform_tensor_view(a_n_y_ho_x_wo_c,
|
||||
make_tuple(make_merge_transform(make_tuple(N, Ho, Wo)),
|
||||
make_merge_transform(make_tuple(Y, X, C))),
|
||||
make_tuple(Sequence<0, 2, 4>{}, Sequence<1, 3, 5>{}),
|
||||
make_tuple(Sequence<0>{}, Sequence<1>{}));
|
||||
|
||||
auto dst_gemmm_gemmk = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_a_mtx,
|
||||
make_tuple(a_gemmm_gemmk_lengths[0], a_gemmm_gemmk_lengths[1]),
|
||||
make_tuple(a_gemmm_gemmk_strides[0], a_gemmm_gemmk_strides[1]),
|
||||
Number<32>{},
|
||||
Number<1>{});
|
||||
|
||||
const auto numGemmM = a_gemmm_gemmk_lengths[0];
|
||||
const auto numGemmK = a_gemmm_gemmk_lengths[1];
|
||||
|
||||
const auto id_block = get_block_id();
|
||||
|
||||
const auto num_tile_m = __builtin_amdgcn_readfirstlane(numGemmM / kMPerBlock);
|
||||
|
||||
const auto block2tile = make_cluster_descriptor(make_tuple(num_tile_m));
|
||||
|
||||
const auto i_gemmm_gemmk = block2tile.CalculateBottomIndex(make_multi_index(id_block));
|
||||
|
||||
const auto iGemmM = __builtin_amdgcn_readfirstlane(i_gemmm_gemmk[0]) * kMPerBlock;
|
||||
|
||||
// src window
|
||||
auto src_block_window =
|
||||
make_tile_window(src_gemmm_gemmk,
|
||||
make_tuple(Number<kMPerBlock>{}, Number<kKPerBlock>{}),
|
||||
{iGemmM, 0},
|
||||
MakeBlockCopyTileDistribution());
|
||||
|
||||
// dst window
|
||||
auto dst_block_window = make_tile_window(
|
||||
dst_gemmm_gemmk, make_tuple(Number<kMPerBlock>{}, Number<kKPerBlock>{}), {iGemmM, 0});
|
||||
|
||||
index_t iGemmK = 0;
|
||||
|
||||
do
|
||||
{
|
||||
const auto block_tile = load_tile(src_block_window);
|
||||
|
||||
store_tile(dst_block_window, block_tile);
|
||||
|
||||
move_tile_window(src_block_window, {0, kKPerBlock});
|
||||
move_tile_window(dst_block_window, {0, kKPerBlock});
|
||||
|
||||
iGemmK += kKPerBlock;
|
||||
|
||||
} while(iGemmK < numGemmK);
|
||||
}
|
||||
};
|
||||
|
||||
int main()
|
||||
{
|
||||
using DataType = ck::half_t;
|
||||
|
||||
constexpr ck::index_t NumDimSpatial = 2;
|
||||
|
||||
ck::index_t N = 128;
|
||||
ck::index_t K = 1;
|
||||
ck::index_t C = 256;
|
||||
ck::index_t Y = 3;
|
||||
ck::index_t X = 3;
|
||||
ck::index_t Hi = 28;
|
||||
ck::index_t Wi = 28;
|
||||
ck::index_t Ho = 14;
|
||||
ck::index_t Wo = 14;
|
||||
|
||||
std::array<ck::index_t, NumDimSpatial + 2> in_lengths{N, Hi, Wi, C};
|
||||
std::array<ck::index_t, NumDimSpatial + 2> in_strides{0, 0, 0, 1};
|
||||
|
||||
std::array<ck::index_t, NumDimSpatial + 2> wei_lengths{K, Y, X, C};
|
||||
std::array<ck::index_t, NumDimSpatial + 2> wei_strides{0, 0, 0, 1};
|
||||
|
||||
std::array<ck::index_t, NumDimSpatial + 2> out_lengths{N, Ho, Wo, K};
|
||||
std::array<ck::index_t, NumDimSpatial + 2> out_strides{0, 0, 0, 1};
|
||||
|
||||
std::partial_sum(rbegin(in_lengths),
|
||||
std::prev(rend(in_lengths)),
|
||||
std::next(rbegin(in_strides)),
|
||||
std::multiplies<>{});
|
||||
std::partial_sum(rbegin(wei_lengths),
|
||||
std::prev(rend(wei_lengths)),
|
||||
std::next(rbegin(wei_strides)),
|
||||
std::multiplies<>{});
|
||||
std::partial_sum(rbegin(out_lengths),
|
||||
std::prev(rend(out_lengths)),
|
||||
std::next(rbegin(out_strides)),
|
||||
std::multiplies<>{});
|
||||
|
||||
std::array<ck::index_t, NumDimSpatial> filter_strides{2, 2};
|
||||
std::array<ck::index_t, NumDimSpatial> filter_dilations{1, 1};
|
||||
std::array<ck::index_t, NumDimSpatial> input_left_pads{1, 1};
|
||||
std::array<ck::index_t, NumDimSpatial> input_right_pads{1, 1};
|
||||
|
||||
// matrix
|
||||
std::array<ck::index_t, 2> in_mtx_lengths{N * Ho * Wo, Y * X * C};
|
||||
std::array<ck::index_t, 2> in_mtx_strides{0, 1};
|
||||
|
||||
std::partial_sum(rbegin(in_mtx_lengths),
|
||||
std::prev(rend(in_mtx_lengths)),
|
||||
std::next(rbegin(in_mtx_strides)),
|
||||
std::multiplies<>{});
|
||||
|
||||
// host verify
|
||||
Tensor<DataType> in_host(in_lengths, in_strides);
|
||||
Tensor<DataType> in_mtx_host_ref(in_mtx_lengths, in_mtx_strides);
|
||||
Tensor<DataType> in_mtx_host_dev(in_mtx_lengths, in_mtx_strides);
|
||||
|
||||
std::cout << " image tensor element size: " << in_host.GetElementSize() << std::endl;
|
||||
std::cout << "matrix tensor element size: " << in_mtx_host_ref.GetElementSize() << std::endl;
|
||||
|
||||
std::cout << " image tensor element space size: " << in_host.GetElementSpaceSize() << std::endl;
|
||||
std::cout << "matrix tensor element sapce size: " << in_mtx_host_ref.GetElementSpaceSize()
|
||||
<< std::endl;
|
||||
|
||||
ck::utils::FillUniformDistributionIntegerValue<DataType>{-5.f, 5.f}(in_host);
|
||||
|
||||
reference_im2col(in_mtx_host_ref,
|
||||
in_host,
|
||||
N,
|
||||
K,
|
||||
C,
|
||||
Y,
|
||||
X,
|
||||
Hi,
|
||||
Wi,
|
||||
Ho,
|
||||
Wo,
|
||||
filter_strides[0],
|
||||
filter_strides[1],
|
||||
filter_dilations[0],
|
||||
filter_dilations[1],
|
||||
input_left_pads[0],
|
||||
input_left_pads[1],
|
||||
input_right_pads[0],
|
||||
input_right_pads[1]);
|
||||
|
||||
DeviceMem in_buf(sizeof(DataType) * in_host.GetElementSpaceSize());
|
||||
DeviceMem in_mtx_buf(sizeof(DataType) * in_mtx_host_ref.GetElementSpaceSize());
|
||||
|
||||
in_buf.ToDevice(in_host.mData.data());
|
||||
|
||||
constexpr ck::index_t kBlockSize = 256;
|
||||
|
||||
constexpr ck::index_t kGemmMPerBlock = 256;
|
||||
constexpr ck::index_t kGemmKPerBlock = 128;
|
||||
|
||||
ck::index_t kGridSize = (N * Ho * Wo) / kGemmMPerBlock;
|
||||
|
||||
float ave_time =
|
||||
launch_kernel(StreamConfig{nullptr, true},
|
||||
Im2Col<2, DataType, kBlockSize, kGemmMPerBlock, kGemmKPerBlock>{},
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0,
|
||||
in_lengths,
|
||||
in_strides,
|
||||
wei_lengths,
|
||||
wei_strides,
|
||||
out_lengths,
|
||||
out_strides,
|
||||
filter_strides,
|
||||
filter_dilations,
|
||||
input_left_pads,
|
||||
input_right_pads,
|
||||
in_mtx_lengths,
|
||||
in_mtx_strides,
|
||||
static_cast<DataType*>(in_buf.GetDeviceBuffer()),
|
||||
static_cast<DataType*>(in_mtx_buf.GetDeviceBuffer()));
|
||||
|
||||
std::size_t num_btype = sizeof(DataType) * in_host.GetElementSize() +
|
||||
sizeof(DataType) * in_mtx_host_ref.GetElementSize();
|
||||
|
||||
float gb_per_sec = num_btype / 1.E6 / ave_time;
|
||||
|
||||
std::cout << "Perf: " << ave_time << " ms, " << gb_per_sec << " GB/s" << std::endl;
|
||||
|
||||
in_mtx_buf.FromDevice(in_mtx_host_dev.mData.data());
|
||||
|
||||
return !ck::utils::check_err(in_mtx_host_dev, in_mtx_host_ref);
|
||||
}
|
||||
105
example/91_tile_program/reduce.cpp
Normal file
105
example/91_tile_program/reduce.cpp
Normal file
@@ -0,0 +1,105 @@
|
||||
#include <cstring>
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/cluster_descriptor.hpp"
|
||||
#include "ck/tensor/tensor_view.hpp"
|
||||
#include "ck/host_utility/device_prop.hpp"
|
||||
#include "ck/host_utility/kernel_launch.hpp"
|
||||
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/fill.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
|
||||
#include "reduce.hpp"
|
||||
|
||||
template <typename ADataType, typename AccDataType, typename BDataType>
|
||||
void reference_reduce(const Tensor<ADataType>& a_m_n, Tensor<BDataType>& b_m)
|
||||
{
|
||||
auto f = [&](auto m) {
|
||||
const int N = a_m_n.mDesc.GetLengths()[1];
|
||||
|
||||
AccDataType v_acc = 0;
|
||||
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
const ADataType v_a = a_m_n(m, n);
|
||||
|
||||
v_acc += v_a;
|
||||
}
|
||||
|
||||
b_m(m) = ck::type_convert<BDataType>(v_acc);
|
||||
};
|
||||
|
||||
make_ParallelTensorFunctor(f, b_m.mDesc.GetLengths()[0])(std::thread::hardware_concurrency());
|
||||
}
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
using ADataType = ck::half_t;
|
||||
using AccDataType = float;
|
||||
using BDataType = ck::half_t;
|
||||
|
||||
ck::index_t M = 3328;
|
||||
ck::index_t N = 4096;
|
||||
|
||||
if(argc == 3)
|
||||
{
|
||||
M = std::stoi(argv[1]);
|
||||
N = std::stoi(argv[2]);
|
||||
}
|
||||
|
||||
std::array<ck::index_t, 2> a_lengths{M, N};
|
||||
std::array<ck::index_t, 2> a_strides{N, 1};
|
||||
|
||||
std::array<ck::index_t, 1> b_lengths{M};
|
||||
std::array<ck::index_t, 1> b_strides{1};
|
||||
|
||||
// host verify
|
||||
Tensor<ADataType> a_host(a_lengths, a_strides);
|
||||
Tensor<BDataType> b_host_ref(b_lengths, b_strides);
|
||||
Tensor<BDataType> b_host_dev(b_lengths, b_strides);
|
||||
|
||||
ck::utils::FillUniformDistributionIntegerValue<ADataType>{-5.f, 5.f}(a_host);
|
||||
|
||||
// reference
|
||||
reference_reduce<ADataType, AccDataType, BDataType>(a_host, b_host_ref);
|
||||
|
||||
DeviceMem a_buf(sizeof(ADataType) * a_host.GetElementSpaceSize());
|
||||
DeviceMem b_buf(sizeof(BDataType) * b_host_ref.GetElementSpaceSize());
|
||||
|
||||
a_buf.ToDevice(a_host.mData.data());
|
||||
|
||||
constexpr ck::index_t kMPerBlock = 128;
|
||||
constexpr ck::index_t kNPerBlock = 128;
|
||||
|
||||
constexpr ck::index_t kBlockSize = 256;
|
||||
ck::index_t kGridSize = (M / kMPerBlock);
|
||||
|
||||
std::cout << "grid size " << kGridSize << std::endl;
|
||||
|
||||
const auto kernel =
|
||||
Reduce<ADataType, AccDataType, BDataType, kBlockSize, kMPerBlock, kNPerBlock>{};
|
||||
|
||||
float ave_time = launch_kernel(StreamConfig{nullptr, true},
|
||||
kernel,
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0,
|
||||
static_cast<ADataType*>(a_buf.GetDeviceBuffer()),
|
||||
static_cast<BDataType*>(b_buf.GetDeviceBuffer()),
|
||||
M,
|
||||
N);
|
||||
|
||||
b_buf.FromDevice(b_host_dev.mData.data());
|
||||
|
||||
std::size_t num_btype = sizeof(ADataType) * M * N + sizeof(BDataType) * M;
|
||||
|
||||
float gb_per_sec = num_btype / 1.E6 / ave_time;
|
||||
|
||||
std::cout << "Perf: " << ave_time << " ms, " << gb_per_sec << " GB/s" << std::endl;
|
||||
|
||||
return !ck::utils::check_err(b_host_dev, b_host_ref);
|
||||
}
|
||||
141
example/91_tile_program/reduce.hpp
Normal file
141
example/91_tile_program/reduce.hpp
Normal file
@@ -0,0 +1,141 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/tensor_adaptor.hpp"
|
||||
|
||||
#include "ck/tile_program/tile/tile_distribution.hpp"
|
||||
#include "ck/tile_program/tile/tile_window.hpp"
|
||||
#include "ck/tile_program/tile/load_tile.hpp"
|
||||
#include "ck/tile_program/tile/store_tile.hpp"
|
||||
#include "ck/tile_program/tile/tile_elementwise.hpp"
|
||||
#include "ck/tile_program/block_tile/block_reduce.hpp"
|
||||
|
||||
template <typename ADataType,
|
||||
typename AccDataType,
|
||||
typename BDataType,
|
||||
ck::index_t kBlockSize,
|
||||
ck::index_t kMPerBlock,
|
||||
ck::index_t kNPerBlock>
|
||||
struct Reduce
|
||||
{
|
||||
#if 0
|
||||
__device__ static constexpr auto MakeABlockTileDistribution()
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
|
||||
// 2x2 wave
|
||||
return make_static_tile_distribution(
|
||||
StaticTileDistributionEncoding<Sequence<>,
|
||||
Tuple<Sequence<2, 2, 4, 2, 4>, Sequence<2, 2, 32>>,
|
||||
Tuple<Sequence<1, 2>, Sequence<1, 2>>,
|
||||
Tuple<Sequence<1, 1>, Sequence<3, 2>>,
|
||||
Sequence<1, 2, 1, 1>,
|
||||
Sequence<0, 0, 2, 4>>{});
|
||||
}
|
||||
#elif 0
|
||||
__device__ static constexpr auto MakeABlockTileDistribution()
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
|
||||
// 2x2 wave
|
||||
return make_static_tile_distribution(
|
||||
StaticTileDistributionEncoding<Sequence<>,
|
||||
Tuple<Sequence<2, 2, 32>, Sequence<2, 2, 4, 2, 4>>,
|
||||
Tuple<Sequence<2, 1>, Sequence<2, 1>>,
|
||||
Tuple<Sequence<1, 1>, Sequence<3, 2>>,
|
||||
Sequence<2, 1, 2, 2>,
|
||||
Sequence<0, 0, 2, 4>>{});
|
||||
}
|
||||
#elif 1
|
||||
__device__ static constexpr auto MakeABlockTileDistribution()
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
|
||||
// 4x1 wave
|
||||
return make_static_tile_distribution(
|
||||
StaticTileDistributionEncoding<Sequence<>,
|
||||
Tuple<Sequence<1, 4, 4, 2, 4>, Sequence<4, 1, 32>>,
|
||||
Tuple<Sequence<1, 2>, Sequence<1, 2>>,
|
||||
Tuple<Sequence<1, 1>, Sequence<3, 2>>,
|
||||
Sequence<1, 2, 1, 1>,
|
||||
Sequence<0, 0, 2, 4>>{});
|
||||
}
|
||||
#endif
|
||||
|
||||
__device__ void
|
||||
operator()(const ADataType* p_a, BDataType* p_b, ck::index_t M, ck::index_t N) const
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
using namespace ck::tile_program::block;
|
||||
|
||||
const auto a_m_n = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_a, make_tuple(M, N), make_tuple(N, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
const auto iM = get_block_id() * kMPerBlock;
|
||||
|
||||
// A window
|
||||
auto a_block_window =
|
||||
make_tile_window(a_m_n,
|
||||
make_tuple(Number<kMPerBlock>{}, Number<kNPerBlock>{}),
|
||||
{iM, 0},
|
||||
MakeABlockTileDistribution());
|
||||
|
||||
const auto f_reduce = [](const auto& v0, const auto& v1) { return v0 + v1; };
|
||||
|
||||
const ADataType reduce_init_value = 0;
|
||||
|
||||
constexpr auto reduce_dims = Sequence<1>{};
|
||||
|
||||
// Acc tile
|
||||
// FIXME: support cross warp reduction
|
||||
auto acc_block_tensor = decltype(block_tile_reduce<AccDataType>(
|
||||
load_tile(a_block_window), reduce_dims, f_reduce, reduce_init_value)){};
|
||||
|
||||
// init Acc tile
|
||||
tile_elementwise_inout(
|
||||
[&](auto& acc) { acc = type_convert<AccDataType>(reduce_init_value); },
|
||||
acc_block_tensor);
|
||||
|
||||
// loop
|
||||
index_t iN = 0;
|
||||
|
||||
do
|
||||
{
|
||||
const auto a_block_tensor = load_tile(a_block_window);
|
||||
|
||||
// FIXME: support cross warp reduction
|
||||
block_tile_reduce(acc_block_tensor, a_block_tensor, reduce_dims, f_reduce);
|
||||
|
||||
move_tile_window(a_block_window, {0, kNPerBlock});
|
||||
|
||||
iN += kNPerBlock;
|
||||
|
||||
} while(iN < N);
|
||||
|
||||
// FIXME: support cross warp reduction
|
||||
block_tile_reduce_sync(acc_block_tensor, f_reduce);
|
||||
|
||||
// convert acc_block_tensor to b_block_tensor
|
||||
const auto b_block_tensor = tile_elementwise_in(
|
||||
[](const auto& acc) { return type_convert<BDataType>(acc); }, acc_block_tensor);
|
||||
|
||||
// B
|
||||
const auto b_m = make_naive_tensor_view_packed<AddressSpaceEnum::Global>(
|
||||
p_b, make_tuple(M), Number<32>{});
|
||||
|
||||
// B window
|
||||
auto b_block_window = make_tile_window(b_m, make_tuple(Number<kMPerBlock>{}), {iM});
|
||||
|
||||
// store B tile
|
||||
store_tile(b_block_window, b_block_tensor);
|
||||
}
|
||||
};
|
||||
36
example/91_tile_program/reference_batched_gemm.hpp
Normal file
36
example/91_tile_program/reference_batched_gemm.hpp
Normal file
@@ -0,0 +1,36 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
|
||||
template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>
|
||||
void reference_batched_gemm(const Tensor<ADataType>& a_b_m_k,
|
||||
const Tensor<BDataType>& b_b_n_k,
|
||||
Tensor<CDataType>& c_b_m_n)
|
||||
{
|
||||
const int N = b_b_n_k.mDesc.GetLengths()[1];
|
||||
const int K = b_b_n_k.mDesc.GetLengths()[2];
|
||||
|
||||
auto f = [&](auto batch, auto m) {
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
AccDataType v_acc = 0;
|
||||
|
||||
for(int k = 0; k < K; ++k)
|
||||
{
|
||||
ADataType v_a = a_b_m_k(batch, m, k);
|
||||
BDataType v_b = b_b_n_k(batch, n, k);
|
||||
|
||||
v_acc += ck::type_convert<AccDataType>(v_a) * ck::type_convert<AccDataType>(v_b);
|
||||
}
|
||||
|
||||
c_b_m_n(batch, m, n) = ck::type_convert<CDataType>(v_acc);
|
||||
}
|
||||
};
|
||||
|
||||
make_ParallelTensorFunctor(f, c_b_m_n.mDesc.GetLengths()[0], c_b_m_n.mDesc.GetLengths()[1])(
|
||||
std::thread::hardware_concurrency());
|
||||
}
|
||||
46
example/91_tile_program/reference_batched_softmax.hpp
Normal file
46
example/91_tile_program/reference_batched_softmax.hpp
Normal file
@@ -0,0 +1,46 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
|
||||
template <typename ADataType, typename AccDataType, typename BDataType>
|
||||
void reference_batched_softmax(const Tensor<ADataType>& a_b_m_n, Tensor<BDataType>& b_b_m_n)
|
||||
{
|
||||
const int N = a_b_m_n.mDesc.GetLengths()[2];
|
||||
|
||||
auto f = [&](auto batch, auto m) {
|
||||
AccDataType v_max = ck::NumericLimits<ADataType>::Lowest();
|
||||
|
||||
// max
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
const ADataType v_a = a_b_m_n(batch, m, n);
|
||||
|
||||
v_max = v_max < v_a ? v_a : v_max;
|
||||
}
|
||||
|
||||
AccDataType v_exp_sum = 0;
|
||||
|
||||
// sum
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
const ADataType v_a = a_b_m_n(batch, m, n);
|
||||
|
||||
v_exp_sum += ck::math::exp(v_a - v_max);
|
||||
}
|
||||
|
||||
// elementwise
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
const ADataType v_a = a_b_m_n(batch, m, n);
|
||||
|
||||
b_b_m_n(batch, m, n) = ck::math::exp(v_a - v_max) / v_exp_sum;
|
||||
}
|
||||
};
|
||||
|
||||
make_ParallelTensorFunctor(f, b_b_m_n.mDesc.GetLengths()[0], b_b_m_n.mDesc.GetLengths()[1])(
|
||||
std::thread::hardware_concurrency());
|
||||
}
|
||||
35
example/91_tile_program/reference_gemm.hpp
Normal file
35
example/91_tile_program/reference_gemm.hpp
Normal file
@@ -0,0 +1,35 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
|
||||
template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>
|
||||
void reference_gemm(const Tensor<ADataType>& a_m_k,
|
||||
const Tensor<BDataType>& b_n_k,
|
||||
Tensor<CDataType>& c_m_n)
|
||||
{
|
||||
const int N = b_n_k.mDesc.GetLengths()[0];
|
||||
const int K = b_n_k.mDesc.GetLengths()[1];
|
||||
|
||||
auto f = [&](auto m) {
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
AccDataType v_acc = 0;
|
||||
|
||||
for(int k = 0; k < K; ++k)
|
||||
{
|
||||
ADataType v_a = a_m_k(m, k);
|
||||
BDataType v_b = b_n_k(n, k);
|
||||
|
||||
v_acc += ck::type_convert<AccDataType>(v_a) * ck::type_convert<AccDataType>(v_b);
|
||||
}
|
||||
|
||||
c_m_n(m, n) = ck::type_convert<CDataType>(v_acc);
|
||||
}
|
||||
};
|
||||
|
||||
make_ParallelTensorFunctor(f, c_m_n.mDesc.GetLengths()[0])(std::thread::hardware_concurrency());
|
||||
}
|
||||
45
example/91_tile_program/reference_softmax.hpp
Normal file
45
example/91_tile_program/reference_softmax.hpp
Normal file
@@ -0,0 +1,45 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
|
||||
template <typename ADataType, typename AccDataType, typename BDataType>
|
||||
void reference_softmax(const Tensor<ADataType>& a_m_n, Tensor<BDataType>& b_m_n)
|
||||
{
|
||||
auto f = [&](auto m) {
|
||||
const int N = a_m_n.mDesc.GetLengths()[1];
|
||||
|
||||
AccDataType v_max = ck::NumericLimits<ADataType>::Lowest();
|
||||
|
||||
// max
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
const ADataType v_a = a_m_n(m, n);
|
||||
|
||||
v_max = v_max < v_a ? v_a : v_max;
|
||||
}
|
||||
|
||||
AccDataType v_exp_sum = 0;
|
||||
|
||||
// sum
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
const ADataType v_a = a_m_n(m, n);
|
||||
|
||||
v_exp_sum += ck::math::exp(v_a - v_max);
|
||||
}
|
||||
|
||||
// elementwise
|
||||
for(int n = 0; n < N; ++n)
|
||||
{
|
||||
const ADataType v_a = a_m_n(m, n);
|
||||
|
||||
b_m_n(m, n) = ck::math::exp(v_a - v_max) / v_exp_sum;
|
||||
}
|
||||
};
|
||||
|
||||
make_ParallelTensorFunctor(f, b_m_n.mDesc.GetLengths()[0])(std::thread::hardware_concurrency());
|
||||
}
|
||||
85
example/91_tile_program/softmax.cpp
Normal file
85
example/91_tile_program/softmax.cpp
Normal file
@@ -0,0 +1,85 @@
|
||||
#include <cstring>
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/cluster_descriptor.hpp"
|
||||
#include "ck/tensor/tensor_view.hpp"
|
||||
#include "ck/host_utility/device_prop.hpp"
|
||||
#include "ck/host_utility/kernel_launch.hpp"
|
||||
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/fill.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
|
||||
#include "reference_softmax.hpp"
|
||||
#include "softmax.hpp"
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
using ADataType = ck::half_t;
|
||||
using AccDataType = float;
|
||||
using BDataType = ck::half_t;
|
||||
|
||||
ck::index_t M = 13312;
|
||||
ck::index_t N = 4096;
|
||||
|
||||
if(argc == 3)
|
||||
{
|
||||
M = std::stoi(argv[1]);
|
||||
N = std::stoi(argv[2]);
|
||||
}
|
||||
|
||||
std::array<ck::index_t, 2> a_lengths{M, N};
|
||||
std::array<ck::index_t, 2> a_strides{N, 1};
|
||||
|
||||
std::array<ck::index_t, 2> b_lengths{M, N};
|
||||
std::array<ck::index_t, 2> b_strides{N, 1};
|
||||
|
||||
// host verify
|
||||
Tensor<ADataType> a_host(a_lengths, a_strides);
|
||||
Tensor<BDataType> b_host_ref(b_lengths, b_strides);
|
||||
Tensor<BDataType> b_host_dev(b_lengths, b_strides);
|
||||
|
||||
ck::utils::FillUniformDistributionIntegerValue<ADataType>{-5.f, 5.f}(a_host);
|
||||
|
||||
// reference
|
||||
reference_softmax<ADataType, AccDataType, BDataType>(a_host, b_host_ref);
|
||||
|
||||
DeviceMem a_buf(sizeof(ADataType) * a_host.GetElementSpaceSize());
|
||||
DeviceMem b_buf(sizeof(BDataType) * b_host_ref.GetElementSpaceSize());
|
||||
|
||||
a_buf.ToDevice(a_host.mData.data());
|
||||
|
||||
constexpr ck::index_t kMPerBlock = 128;
|
||||
constexpr ck::index_t kNPerBlock = 128;
|
||||
|
||||
constexpr ck::index_t kBlockSize = 256;
|
||||
ck::index_t kGridSize = (M / kMPerBlock);
|
||||
|
||||
std::cout << "grid size " << kGridSize << std::endl;
|
||||
|
||||
const auto kernel =
|
||||
Softmax<ADataType, AccDataType, BDataType, kBlockSize, kMPerBlock, kNPerBlock>{};
|
||||
|
||||
float ave_time = launch_kernel(StreamConfig{nullptr, true},
|
||||
kernel,
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0,
|
||||
static_cast<ADataType*>(a_buf.GetDeviceBuffer()),
|
||||
static_cast<BDataType*>(b_buf.GetDeviceBuffer()),
|
||||
M,
|
||||
N);
|
||||
|
||||
b_buf.FromDevice(b_host_dev.mData.data());
|
||||
|
||||
std::size_t num_btype = sizeof(ADataType) * M * N + sizeof(BDataType) * M * N;
|
||||
|
||||
float gb_per_sec = num_btype / 1.E6 / ave_time;
|
||||
|
||||
std::cout << "Perf: " << ave_time << " ms, " << gb_per_sec << " GB/s" << std::endl;
|
||||
|
||||
return !ck::utils::check_err(b_host_dev, b_host_ref);
|
||||
}
|
||||
273
example/91_tile_program/softmax.hpp
Normal file
273
example/91_tile_program/softmax.hpp
Normal file
@@ -0,0 +1,273 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor.hpp"
|
||||
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
|
||||
#include "ck/tensor_description/tensor_adaptor.hpp"
|
||||
|
||||
#include "ck/tile_program/tile/tile_distribution.hpp"
|
||||
#include "ck/tile_program/tile/tile_window.hpp"
|
||||
#include "ck/tile_program/tile/load_tile.hpp"
|
||||
#include "ck/tile_program/tile/store_tile.hpp"
|
||||
#include "ck/tile_program/tile/tile_elementwise.hpp"
|
||||
#include "ck/tile_program/block_tile/block_reduce.hpp"
|
||||
|
||||
template <typename ADataType,
|
||||
typename AccDataType,
|
||||
typename BDataType,
|
||||
ck::index_t kBlockSize,
|
||||
ck::index_t kMPerBlock,
|
||||
ck::index_t kNPerBlock>
|
||||
struct Softmax
|
||||
{
|
||||
__device__ static constexpr auto MakeABlockTileDistribution()
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
|
||||
// 4x1 wave
|
||||
return make_static_tile_distribution(
|
||||
StaticTileDistributionEncoding<Sequence<>,
|
||||
Tuple<Sequence<1, 4, 4, 2, 4>, Sequence<4, 1, 32>>,
|
||||
Tuple<Sequence<1, 2>, Sequence<1, 2>>,
|
||||
Tuple<Sequence<1, 1>, Sequence<3, 2>>,
|
||||
Sequence<1, 2, 1, 1>,
|
||||
Sequence<0, 0, 2, 4>>{});
|
||||
}
|
||||
|
||||
__device__ void
|
||||
MultiPassSoftmax(const ADataType* p_a, BDataType* p_b, ck::index_t M, ck::index_t N) const
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
using namespace ck::tile_program::block;
|
||||
|
||||
constexpr auto I0 = Number<0>{};
|
||||
constexpr auto I1 = Number<1>{};
|
||||
|
||||
// A DRAM tensor view
|
||||
const auto a_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_a, make_tuple(M, N), make_tuple(N, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
const auto iM = get_block_id() * kMPerBlock;
|
||||
|
||||
// A DRAM window
|
||||
auto a_dram_window =
|
||||
make_tile_window(a_dram,
|
||||
make_tuple(Number<kMPerBlock>{}, Number<kNPerBlock>{}),
|
||||
{iM, 0},
|
||||
MakeABlockTileDistribution());
|
||||
|
||||
// m = rowmax(A)
|
||||
const auto f_max = [](auto e0, auto e1) { return max(e0, e1); };
|
||||
|
||||
auto m = decltype(block_tile_reduce<AccDataType>(
|
||||
load_tile(a_dram_window), Sequence<1>{}, f_max, ADataType{})){};
|
||||
|
||||
tile_elementwise_inout(
|
||||
[&](auto& e) { e = type_convert<AccDataType>(NumericLimits<ADataType>::Lowest()); }, m);
|
||||
|
||||
index_t iN = 0;
|
||||
|
||||
do
|
||||
{
|
||||
// load A tile from DRAM
|
||||
const auto a = load_tile(a_dram_window);
|
||||
|
||||
// m = rowmax(A)
|
||||
block_tile_reduce(m, a, Sequence<1>{}, f_max);
|
||||
|
||||
move_tile_window(a_dram_window, {0, kNPerBlock});
|
||||
|
||||
iN += kNPerBlock;
|
||||
|
||||
} while(iN < N);
|
||||
|
||||
// cross lane reduce: max
|
||||
block_tile_reduce_sync(m, f_max);
|
||||
|
||||
// reset window location
|
||||
iN = 0;
|
||||
move_tile_window(a_dram_window, {0, -N});
|
||||
|
||||
// l = rowsum(exp(A - m))
|
||||
auto l = make_static_distributed_tensor<AccDataType>(m.GetTileDistribution());
|
||||
|
||||
tile_elementwise_inout([&](auto& e) { e = 0; }, l);
|
||||
|
||||
do
|
||||
{
|
||||
// load A tile from DRAM
|
||||
const auto a = load_tile(a_dram_window);
|
||||
|
||||
constexpr auto a_spans = decltype(a)::GetDistributedSpans();
|
||||
|
||||
sweep_tile_span(a_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
sweep_tile_span(a_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
// l = rowsum(exp(A - m))
|
||||
l(i_idx) += math::exp(a[i_j_idx] - m[i_idx]);
|
||||
});
|
||||
});
|
||||
|
||||
move_tile_window(a_dram_window, {0, kNPerBlock});
|
||||
|
||||
iN += kNPerBlock;
|
||||
|
||||
} while(iN < N);
|
||||
|
||||
// cross lane reduce: sum
|
||||
block_tile_reduce_sync(l, [](auto e0, auto e1) { return e0 + e1; });
|
||||
|
||||
// reset window location
|
||||
iN = 0;
|
||||
move_tile_window(a_dram_window, {0, -N});
|
||||
|
||||
// B DRAM tensor view
|
||||
const auto b_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_b, make_tuple(M, N), make_tuple(N, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
// B DRAM window
|
||||
auto b_dram_window = make_tile_window(
|
||||
b_dram, make_tuple(Number<kMPerBlock>{}, Number<kNPerBlock>{}), {iM, 0});
|
||||
|
||||
// B = exp(A - m) / l
|
||||
do
|
||||
{
|
||||
// load A tile from DRAM
|
||||
const auto a = load_tile(a_dram_window);
|
||||
|
||||
constexpr auto a_spans = decltype(a)::GetDistributedSpans();
|
||||
|
||||
auto b = make_static_distributed_tensor<BDataType>(a.GetTileDistribution());
|
||||
|
||||
sweep_tile_span(a_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
sweep_tile_span(a_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
// B = exp(A - m) / l
|
||||
b(i_j_idx) =
|
||||
type_convert<BDataType>(math::exp(a[i_j_idx] - m[i_idx]) / l[i_idx]);
|
||||
});
|
||||
});
|
||||
|
||||
// store B tile
|
||||
store_tile(b_dram_window, b);
|
||||
|
||||
move_tile_window(a_dram_window, {0, kNPerBlock});
|
||||
move_tile_window(b_dram_window, {0, kNPerBlock});
|
||||
|
||||
iN += kNPerBlock;
|
||||
|
||||
} while(iN < N);
|
||||
}
|
||||
|
||||
__device__ void
|
||||
SinglePassSoftmax(const ADataType* p_a, BDataType* p_b, ck::index_t M, ck::index_t N) const
|
||||
{
|
||||
using namespace ck;
|
||||
using namespace ck::tile_program;
|
||||
using namespace ck::tile_program::block;
|
||||
|
||||
constexpr auto I0 = Number<0>{};
|
||||
constexpr auto I1 = Number<1>{};
|
||||
|
||||
// A DRAM tensor view
|
||||
const auto a_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_a, make_tuple(M, N), make_tuple(N, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
const auto iM = get_block_id() * kMPerBlock;
|
||||
|
||||
// A DRAM window
|
||||
auto a_dram_window =
|
||||
make_tile_window(a_dram,
|
||||
make_tuple(Number<kMPerBlock>{}, Number<kNPerBlock>{}),
|
||||
{iM, 0},
|
||||
MakeABlockTileDistribution());
|
||||
|
||||
// f_max
|
||||
const auto f_max = [](auto e0, auto e1) { return max(e0, e1); };
|
||||
|
||||
// m = rowmax(A)
|
||||
auto m = decltype(block_tile_reduce<AccDataType>(
|
||||
load_tile(a_dram_window), Sequence<1>{}, f_max, ADataType{})){};
|
||||
|
||||
tile_elementwise_inout(
|
||||
[&](auto& e) { e = type_convert<AccDataType>(NumericLimits<ADataType>::Lowest()); }, m);
|
||||
|
||||
// l = rowsum(exp(A - m))
|
||||
auto l = make_static_distributed_tensor<AccDataType>(m.GetTileDistribution());
|
||||
|
||||
tile_elementwise_inout([&](auto& e) { e = 0; }, l);
|
||||
|
||||
// load A tile from DRAM
|
||||
const auto a = load_tile(a_dram_window);
|
||||
|
||||
constexpr auto a_spans = decltype(a)::GetDistributedSpans();
|
||||
|
||||
// m = rowmax(A)
|
||||
block_tile_reduce(m, a, Sequence<1>{}, f_max);
|
||||
|
||||
// cross lane reduce: max
|
||||
block_tile_reduce_sync(m, f_max);
|
||||
|
||||
// l = rowsum(exp(A - m))
|
||||
sweep_tile_span(a_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
sweep_tile_span(a_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
l(i_idx) += math::exp(a[i_j_idx] - m[i_idx]);
|
||||
});
|
||||
});
|
||||
|
||||
// cross lane reduce: sum
|
||||
block_tile_reduce_sync(l, [](auto e0, auto e1) { return e0 + e1; });
|
||||
|
||||
auto b = make_static_distributed_tensor<BDataType>(a.GetTileDistribution());
|
||||
|
||||
// B = exp(A - m) / l
|
||||
sweep_tile_span(a_spans[I0], [&](auto idx0) {
|
||||
constexpr auto i_idx = make_tuple(idx0);
|
||||
|
||||
sweep_tile_span(a_spans[I1], [&](auto idx1) {
|
||||
constexpr auto i_j_idx = make_tuple(idx0, idx1);
|
||||
|
||||
b(i_j_idx) = type_convert<BDataType>(math::exp(a[i_j_idx] - m[i_idx]) / l[i_idx]);
|
||||
});
|
||||
});
|
||||
|
||||
// B DRAM tensor view
|
||||
const auto b_dram = make_naive_tensor_view<AddressSpaceEnum::Global>(
|
||||
p_b, make_tuple(M, N), make_tuple(N, 1), Number<32>{}, Number<1>{});
|
||||
|
||||
// B DRAM window
|
||||
auto b_dram_window = make_tile_window(
|
||||
b_dram, make_tuple(Number<kMPerBlock>{}, Number<kNPerBlock>{}), {iM, 0});
|
||||
|
||||
// store B tile
|
||||
store_tile(b_dram_window, b);
|
||||
}
|
||||
|
||||
__device__ void
|
||||
operator()(const ADataType* p_a, BDataType* p_b, ck::index_t M, ck::index_t N) const
|
||||
{
|
||||
if(N > kNPerBlock)
|
||||
{
|
||||
MultiPassSoftmax(p_a, p_b, M, N);
|
||||
}
|
||||
else
|
||||
{
|
||||
SinglePassSoftmax(p_a, p_b, M, N);
|
||||
}
|
||||
}
|
||||
};
|
||||
Reference in New Issue
Block a user