28 lines
828 B
Plaintext
28 lines
828 B
Plaintext
#pragma once
|
|
|
|
#include <deep_gemm/common/types.hpp>
|
|
#include <deep_gemm/common/utils.cuh>
|
|
|
|
namespace deep_gemm {
|
|
|
|
struct EpilogueIdentity {
|
|
template <uint32_t STORE_BLOCK_N>
|
|
__device__ __forceinline__ static uint32_t apply_index_n(const uint32_t &n_idx) {
|
|
return n_idx;
|
|
}
|
|
};
|
|
|
|
template <uint32_t kLeft, uint32_t kMid, uint32_t kRight>
|
|
struct EpilogueHeadSplits: EpilogueIdentity {
|
|
template <uint32_t STORE_BLOCK_N>
|
|
__device__ __forceinline__ static uint32_t apply_index_n(const uint32_t &n_idx) {
|
|
DG_STATIC_ASSERT(kLeft % STORE_BLOCK_N == 0 and kMid % STORE_BLOCK_N == 0
|
|
and kRight % STORE_BLOCK_N == 0, "Invalid head splits config");
|
|
return n_idx + (n_idx + kRight) / (kLeft + kRight) * kMid;
|
|
}
|
|
};
|
|
|
|
#pragma clang diagnostic pop
|
|
|
|
} // namespace deep_gemm
|