mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-04 13:41:24 +00:00
37a8c1f756
* Redesign the DPP8 GEMM kernel to use warp-wise component * Review: Improve error messages * Review: Remove unnecessary empty lines * Review: Fix M, N per thread names * Review: Rename mfma_input_type to dpp_input_type * Review: Fix tensor adaptor; remove unnecessary element * Review: Remove calls to dpp_gemm's MakeCDescriptor * Review: Add blockwise doc, change function names to include dimension names * Review: Remove duplicated code; Move Block2CtileMap alias to the top of the file * Review: Add __restrict__ keywords * Review: Use MatrixPadder for padding A, B, C matrices * Review: Remove hardcoded datatypes * Review: Change names from FloatX to XDataType * Review: Introduce AK0 and BK0 instead of a single K0 * Review: Remove construction of dpp_datatypes object * Review: Rename DppInstrRunner to DppLanegroupGemm
69 lines
2.3 KiB
C++
69 lines
2.3 KiB
C++
// 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/utility/math.hpp"
|
|
#include "ck/utility/inner_product_dpp8.hpp"
|
|
|
|
namespace ck {
|
|
|
|
namespace dpp8 {
|
|
|
|
template <class ABDataType>
|
|
struct dpp_datatypes;
|
|
|
|
template <>
|
|
struct dpp_datatypes<half_t>
|
|
{
|
|
// Dot product of `half2_t` and `half2_t` to get `float`. Reducing 2 elements from K in a
|
|
// single instruction.
|
|
using a_dtype = half_t;
|
|
using b_dtype = half_t;
|
|
using c_dtype = float;
|
|
static constexpr index_t k_per_instr = 2;
|
|
};
|
|
|
|
template <index_t MPerThread,
|
|
index_t NPerThread,
|
|
index_t KPerThread,
|
|
class BaseInputType,
|
|
class AVecDataType,
|
|
class BVecDataType,
|
|
class CVecDataType,
|
|
bool ShareA>
|
|
struct DppLanegroupGemm
|
|
{
|
|
using datatypes_conf = dpp_datatypes<BaseInputType>;
|
|
using ADataType = typename datatypes_conf::a_dtype;
|
|
using BDataType = typename datatypes_conf::b_dtype;
|
|
using CDataType = typename datatypes_conf::c_dtype;
|
|
|
|
__device__ void Run(const AVecDataType& a_vec, const BVecDataType& b_vec, CVecDataType& c_vec)
|
|
{
|
|
constexpr index_t num_c_elems_per_thread = ShareA ? MPerThread : NPerThread;
|
|
|
|
const vector_type<ADataType, KPerThread> a_vector{a_vec};
|
|
const vector_type<BDataType, KPerThread> b_vector{b_vec};
|
|
|
|
static_for<0, num_c_elems_per_thread, 1>{}([&](auto c_idx) {
|
|
float c = c_vec.template AsType<CDataType>()(c_idx);
|
|
// Next `c_idx` implies that we need to pull data from the next lane.
|
|
constexpr index_t source_lane = c_idx;
|
|
static_for<0, KPerThread / datatypes_conf::k_per_instr, 1>{}([&](auto k_chunk) {
|
|
const auto a_k_vec = a_vector.template AsType<AVecDataType>()[k_chunk];
|
|
const auto b_k_vec = b_vector.template AsType<BVecDataType>()[k_chunk];
|
|
ck::dpp8::
|
|
inner_product_dpp<AVecDataType, BVecDataType, CDataType, source_lane, ShareA>(
|
|
a_k_vec, b_k_vec, c);
|
|
});
|
|
c_vec.template AsType<CDataType>()(c_idx) = c;
|
|
});
|
|
}
|
|
};
|
|
|
|
} // namespace dpp8
|
|
|
|
} // namespace ck
|