Add more GPU architectures support (#112)
* Add more GPU architectures support * Update layout.py * Optimize performance, Add SM90 support, Add 1D2D SM100 support * Add fmtlib submodule at commit 553ec11 --------- Co-authored-by: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
This commit is contained in:
Ray Wang
@@ -0,0 +1,298 @@
|
||||
#pragma once
|
||||
|
||||
#include "../../utils/math.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
struct MulticastConfig {
|
||||
int num_multicast;
|
||||
bool is_multicast_on_a;
|
||||
|
||||
MulticastConfig(const int& num_multicast, const bool& is_multicast_on_a):
|
||||
num_multicast(num_multicast), is_multicast_on_a(is_multicast_on_a) {
|
||||
DG_HOST_ASSERT(1 <= num_multicast and num_multicast <= 2);
|
||||
}
|
||||
};
|
||||
|
||||
struct SharedMemoryConfig {
|
||||
int smem_size;
|
||||
int swizzle_a_mode;
|
||||
int swizzle_b_mode;
|
||||
int swizzle_cd_mode;
|
||||
};
|
||||
|
||||
struct ThreadConfig {
|
||||
int num_threads;
|
||||
|
||||
// SM90
|
||||
int num_tma_threads;
|
||||
int num_math_threads;
|
||||
|
||||
// SM100
|
||||
int num_non_epilogue_threads;
|
||||
int num_epilogue_threads;
|
||||
|
||||
static ThreadConfig sm90(const int& num_tma_threads,
|
||||
const int& num_math_threads) {
|
||||
auto config = ThreadConfig();
|
||||
config.num_threads = num_tma_threads + num_math_threads;
|
||||
config.num_tma_threads = num_tma_threads;
|
||||
config.num_math_threads = num_math_threads;
|
||||
return config;
|
||||
}
|
||||
|
||||
static ThreadConfig sm100(const int& num_non_epilogue_threads,
|
||||
const int& num_epilogue_threads) {
|
||||
auto config = ThreadConfig();
|
||||
config.num_threads = num_non_epilogue_threads + num_epilogue_threads;
|
||||
config.num_non_epilogue_threads = num_non_epilogue_threads;
|
||||
config.num_epilogue_threads = num_epilogue_threads;
|
||||
return config;
|
||||
}
|
||||
};
|
||||
|
||||
struct GemmConfig {
|
||||
// Templated configs
|
||||
GemmType gemm_type;
|
||||
KernelType kernel_type;
|
||||
at::ScalarType ab_dtype, cd_dtype;
|
||||
cute::UMMA::Major major_a;
|
||||
cute::UMMA::Major major_b;
|
||||
bool with_accumulation;
|
||||
int block_m, block_n, block_k;
|
||||
int num_stages, num_last_stages;
|
||||
|
||||
// Runtime configs
|
||||
int num_sms;
|
||||
|
||||
// Structured configs
|
||||
MulticastConfig multicast_config;
|
||||
SharedMemoryConfig smem_config;
|
||||
ThreadConfig thread_config;
|
||||
};
|
||||
|
||||
static bool is_multicast_legal(const int& shape_dim, const int& block_dim,
|
||||
const int& num_multicast, const int& num_sms,
|
||||
const bool& require_divisible) {
|
||||
const bool& divisible = ceil_div(shape_dim, block_dim) % num_multicast == 0 or not require_divisible;
|
||||
return divisible and num_sms % num_multicast == 0;
|
||||
}
|
||||
|
||||
static int get_swizzle_mode(const int& block_size, const int& elem_size) {
|
||||
// `> 0` means interleaving
|
||||
// 16B actually means non-swizzling (but interleaving)
|
||||
for (const int& mode: {128, 64, 32, 16}) {
|
||||
if ((block_size * elem_size) % mode == 0)
|
||||
return mode;
|
||||
}
|
||||
DG_HOST_UNREACHABLE("Unreachable");
|
||||
}
|
||||
|
||||
template <typename ArchSpec>
|
||||
static SharedMemoryConfig get_smem_config(const KernelType& kernel_type,
|
||||
const int& m, const int& n, const int& k,
|
||||
const int& block_m, const int& block_n, const int& block_k,
|
||||
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
||||
const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype,
|
||||
const int& num_stages, const MulticastConfig& multicast_config) {
|
||||
const int& ab_elem_size = static_cast<int>(c10::elementSize(ab_dtype));
|
||||
const int& cd_elem_size = static_cast<int>(c10::elementSize(cd_dtype));
|
||||
|
||||
const int& load_block_m = ArchSpec::get_ab_load_block_m(multicast_config, block_m);
|
||||
const int& load_block_n = ArchSpec::get_ab_load_block_n(multicast_config, block_n);
|
||||
const int& swizzle_a_mode = get_swizzle_mode(major_a == cute::UMMA::Major::K ? block_k : load_block_m, ab_elem_size);
|
||||
const int& swizzle_b_mode = get_swizzle_mode(major_b == cute::UMMA::Major::K ? block_k : load_block_n, ab_elem_size);
|
||||
const int& swizzle_cd_mode = get_swizzle_mode(block_n, cd_elem_size);
|
||||
|
||||
// Different archs have different epilogue pipelines
|
||||
const int& smem_cd = ArchSpec::get_smem_cd_size(kernel_type, block_m, block_n, swizzle_cd_mode, cd_dtype);
|
||||
|
||||
// A/B shared memory
|
||||
const int& smem_a_per_stage = load_block_m * block_k * ab_elem_size;
|
||||
const int& smem_b_per_stage = load_block_n * block_k * ab_elem_size;
|
||||
|
||||
// SF shared memory
|
||||
const auto& [smem_sfa_per_stage, smem_sfb_per_stage] =
|
||||
ArchSpec::get_sf_smem_size_per_stage(kernel_type, block_m, block_n, block_k, ab_dtype, cd_dtype);
|
||||
const int& smem_extra_sfb = ArchSpec::get_extra_sfb_smem_size(m, n, k, block_m, block_n, block_k);
|
||||
|
||||
// M-barriers and tensor memory pointers
|
||||
const int& smem_barrier = ArchSpec::get_barrier_smem_size(num_stages);
|
||||
const int& smem_tmem_ptr = ArchSpec::get_tmem_ptr_smem_size();
|
||||
|
||||
// Sum them up
|
||||
int smem_size = 0;
|
||||
smem_size += smem_cd;
|
||||
smem_size += num_stages * smem_a_per_stage;
|
||||
smem_size += num_stages * smem_b_per_stage;
|
||||
smem_size += num_stages * smem_sfa_per_stage;
|
||||
smem_size += num_stages * smem_sfb_per_stage;
|
||||
smem_size += smem_extra_sfb;
|
||||
smem_size += smem_barrier;
|
||||
smem_size += smem_tmem_ptr;
|
||||
|
||||
return SharedMemoryConfig {
|
||||
.smem_size = smem_size,
|
||||
.swizzle_a_mode = swizzle_a_mode,
|
||||
.swizzle_b_mode = swizzle_b_mode,
|
||||
.swizzle_cd_mode = swizzle_cd_mode,
|
||||
};
|
||||
}
|
||||
|
||||
template <typename ArchSpec>
|
||||
static GemmConfig get_best_config(const GemmType& gemm_type, const KernelType& kernel_type,
|
||||
const int& m, const int& n, const int& k, const int& num_groups,
|
||||
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
||||
const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype,
|
||||
const bool& with_accumulation, const int& num_sms) {
|
||||
DG_HOST_ASSERT(ab_dtype == torch::kFloat8_e4m3fn);
|
||||
DG_HOST_ASSERT(cd_dtype == torch::kBFloat16 or cd_dtype == torch::kFloat);
|
||||
|
||||
// Select M/N block sizes
|
||||
// TODO: support `% 16 == 8` block size on SM90
|
||||
const auto& block_ms = gemm_type == GemmType::MGroupedContiguous ?
|
||||
std::vector{get_mk_alignment_for_contiguous_layout()} : std::vector{64, 128, 256};
|
||||
std::vector<int> block_ns;
|
||||
for (int i = 16; i <= 256; i += 16)
|
||||
block_ns.push_back(i);
|
||||
|
||||
// K block size is selected in a fixed manner
|
||||
const auto& block_k = 128 / static_cast<int>(c10::elementSize(ab_dtype));
|
||||
|
||||
// Some util functions
|
||||
const auto& get_num_blocks = [=](const int& block_m, const int& block_n) {
|
||||
return ceil_div(m, block_m) * ceil_div(n, block_n) * num_groups;
|
||||
};
|
||||
const auto& get_num_waves = [=](const int& block_m, const int& block_n) {
|
||||
return ceil_div(get_num_blocks(block_m, block_n), num_sms);
|
||||
};
|
||||
const auto& get_last_wave_util = [=](const int& block_m, const int& block_n) {
|
||||
const auto& num_last_blocks = get_num_blocks(block_m, block_n) % num_sms;
|
||||
return num_last_blocks == 0 ? num_sms : num_last_blocks;
|
||||
};
|
||||
|
||||
// Decide block sizes by waves
|
||||
int best_block_m = 0, best_block_n = 0;
|
||||
int best_num_waves = 0, best_last_util = 0;
|
||||
for (const auto& block_m: block_ms) {
|
||||
for (const auto& block_n: block_ns) {
|
||||
const int& num_waves = get_num_waves(block_m, block_n);
|
||||
const auto& last_util = get_last_wave_util(block_m, block_n);
|
||||
if (not ArchSpec::is_block_size_legal(kernel_type, major_a, major_b, ab_dtype, cd_dtype, block_m, block_n))
|
||||
continue;
|
||||
|
||||
bool success = false;
|
||||
if (best_block_m == 0 or best_block_n == 0 or num_waves < best_num_waves) {
|
||||
success = true;
|
||||
} else if (num_waves == best_num_waves) {
|
||||
// Check last wave utilization
|
||||
success = last_util > best_last_util;
|
||||
if (last_util == best_last_util) {
|
||||
// Case 1: same `block_m`, smaller `block_n` (wasted)
|
||||
success |= block_m == best_block_m and block_n < best_block_n;
|
||||
// Case 2: same `block_n`, smaller `block_m` (wasted)
|
||||
success |= block_n == best_block_n and block_m < best_block_m;
|
||||
// Case 3: different for both `block_m` and `block_n`, larger `block_n` is better
|
||||
success |= block_m != best_block_m and block_n > best_block_n;
|
||||
}
|
||||
}
|
||||
|
||||
// Replace with the new config if successful
|
||||
if (success) {
|
||||
best_block_m = block_m, best_block_n = block_n;
|
||||
best_num_waves = num_waves, best_last_util = last_util;
|
||||
}
|
||||
}
|
||||
}
|
||||
DG_HOST_ASSERT(best_block_m > 0 and best_block_n > 0);
|
||||
|
||||
// Decide the number of TMA multicasts and whether broadcast on A
|
||||
MulticastConfig best_multicast_config = {1, true};
|
||||
const auto& [is_legal_on_a, is_legal_on_b] = ArchSpec::get_multicast_legality(
|
||||
gemm_type, m, n, best_block_m, best_block_n, num_sms);
|
||||
const bool is_legal[2] = {is_legal_on_a, is_legal_on_b};
|
||||
bool order[2] = {false, true};
|
||||
if (best_block_m > best_block_n)
|
||||
std::swap(order[0], order[1]);
|
||||
for (const bool& is_multicast_on_a: order) {
|
||||
if (m >= 512 and is_legal[static_cast<int>(is_multicast_on_a)]) {
|
||||
best_multicast_config = {2, is_multicast_on_a};
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Always pick the largest number of stage
|
||||
constexpr int smem_capacity = ArchSpec::smem_capacity;
|
||||
int best_num_stages = 0;
|
||||
SharedMemoryConfig best_smem_config;
|
||||
for (int num_stages = std::min(12, ceil_div(k, block_k)); num_stages > 0; -- num_stages) {
|
||||
if (not ArchSpec::is_num_stages_legal(ab_dtype, cd_dtype, num_stages, best_block_m, best_block_n, block_k))
|
||||
continue;
|
||||
|
||||
best_smem_config = get_smem_config<ArchSpec>(kernel_type,
|
||||
m, n, k,
|
||||
best_block_m, best_block_n, block_k,
|
||||
major_a, major_b,
|
||||
ab_dtype, cd_dtype,
|
||||
num_stages, best_multicast_config);
|
||||
if (best_smem_config.smem_size <= smem_capacity) {
|
||||
best_num_stages = num_stages;
|
||||
break;
|
||||
}
|
||||
}
|
||||
DG_HOST_ASSERT(best_num_stages != 0);
|
||||
|
||||
// Recompute the minimal number of SMs required
|
||||
// NOTES: less L2 cache usage and less GPU frequency drop
|
||||
int num_min_sms = num_sms;
|
||||
if (ArchSpec::should_minimize_num_sms()) {
|
||||
num_min_sms = ceil_div(ceil_div(m, best_block_m) * ceil_div(n, best_block_n) * num_groups, best_num_waves);
|
||||
num_min_sms = align(num_min_sms, best_multicast_config.num_multicast);
|
||||
DG_HOST_ASSERT(num_min_sms <= num_sms);
|
||||
}
|
||||
|
||||
const auto& config = GemmConfig {
|
||||
.gemm_type = gemm_type,
|
||||
.kernel_type = kernel_type,
|
||||
.ab_dtype = ab_dtype,
|
||||
.cd_dtype = cd_dtype,
|
||||
.major_a = major_a,
|
||||
.major_b = major_b,
|
||||
.with_accumulation = with_accumulation,
|
||||
.block_m = best_block_m,
|
||||
.block_n = best_block_n,
|
||||
.block_k = block_k,
|
||||
.num_stages = best_num_stages,
|
||||
.num_last_stages = ceil_div(k, block_k) % best_num_stages,
|
||||
.num_sms = num_min_sms,
|
||||
.multicast_config = best_multicast_config,
|
||||
// ReSharper disable once CppLocalVariableMightNotBeInitialized
|
||||
.smem_config = best_smem_config,
|
||||
.thread_config = ArchSpec::get_thread_config(kernel_type, best_block_m, best_block_n)
|
||||
};
|
||||
|
||||
// Print configs for the first time
|
||||
if (get_env<int>("DG_JIT_DEBUG") or get_env<int>("DG_PRINT_CONFIGS")) {
|
||||
auto key = std::make_tuple(gemm_type, kernel_type, m, n, k, num_groups, major_a, major_b,
|
||||
ab_dtype, cd_dtype, with_accumulation, num_sms);
|
||||
static std::set<decltype(key)> printed;
|
||||
if (not printed.contains(key)) {
|
||||
printf("Gemm type: %d, kernel type: %d, M: %d, N: %d, K: %d, groups: %d, "
|
||||
"A major: %d, B major: %d, AB dtype: %s, CD dtype: %s, accumulation: %d, "
|
||||
"SM limit: %d -> block M: %d, block N: %d, block K: %d, stages: %d, last stages: %d, "
|
||||
"SMs: %d, multicast: %d, multicast on A: %d, shared memory: %d bytes, swizzle A: %d, "
|
||||
"swizzle B: %d, swizzle CD: %d, threads: %d\n",
|
||||
static_cast<int>(gemm_type), static_cast<int>(kernel_type), m, n, k, num_groups,
|
||||
static_cast<int>(major_a), static_cast<int>(major_b), c10::toString(ab_dtype), c10::toString(cd_dtype),
|
||||
static_cast<int>(with_accumulation), num_sms, best_block_m, best_block_n, block_k,
|
||||
best_num_stages, config.num_last_stages, num_min_sms, best_multicast_config.num_multicast,
|
||||
static_cast<int>(best_multicast_config.is_multicast_on_a),
|
||||
best_smem_config.smem_size, best_smem_config.swizzle_a_mode, best_smem_config.swizzle_b_mode,
|
||||
best_smem_config.swizzle_cd_mode, config.thread_config.num_threads);
|
||||
printed.insert(key);
|
||||
}
|
||||
}
|
||||
return config;
|
||||
}
|
||||
|
||||
} // namespace deep_gemm
|
||||
@@ -0,0 +1,144 @@
|
||||
#pragma once
|
||||
|
||||
#include <cute/arch/mma_sm100_desc.hpp>
|
||||
// Reuse some types in the JIT modules
|
||||
#include <deep_gemm/common/types.hpp>
|
||||
|
||||
#include "common.hpp"
|
||||
#include "../../utils/exception.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
struct SM100ArchSpec {
|
||||
static constexpr int smem_capacity = 232448;
|
||||
|
||||
static int get_ab_load_block_m(const MulticastConfig& config, const int& block_m) {
|
||||
return block_m / (config.is_multicast_on_a ? config.num_multicast : 1);
|
||||
}
|
||||
|
||||
static int get_ab_load_block_n(const MulticastConfig& config, const int& block_n) {
|
||||
return block_n / (config.is_multicast_on_a ? 1 : config.num_multicast);
|
||||
}
|
||||
|
||||
static int get_cd_store_block_m(const int& block_m) {
|
||||
constexpr int layout_ad_m = 128;
|
||||
return std::min(block_m, layout_ad_m);
|
||||
}
|
||||
|
||||
static int get_cd_store_block_n(const int& block_n) {
|
||||
return block_n;
|
||||
}
|
||||
|
||||
static std::pair<int, int> get_sf_uttcp_aligned_block_sizes(
|
||||
const int& block_m, const int& block_n, const at::ScalarType& ab_dtype) {
|
||||
constexpr int num_utccp_aligned_elems = 128;
|
||||
DG_HOST_ASSERT(block_m % num_utccp_aligned_elems == 0);
|
||||
switch (ab_dtype) {
|
||||
case torch::kBFloat16: return {0, 0};
|
||||
case torch::kFloat8_e4m3fn: return {align(block_m, num_utccp_aligned_elems), align(block_n, num_utccp_aligned_elems)};
|
||||
default: DG_HOST_UNREACHABLE("Unknown dtype");
|
||||
}
|
||||
}
|
||||
|
||||
static bool is_block_size_legal(const KernelType& kernel_type,
|
||||
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
||||
const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype,
|
||||
const int& block_m, const int& block_n) {
|
||||
// Layout A/D does not support `block_m == 64` and `block_n % 16 != 0`
|
||||
if (block_m == 64 or block_n % 16 != 0)
|
||||
return false;
|
||||
|
||||
// Performance is lower with 1D1D and `block_m == 256`
|
||||
if (kernel_type == KernelType::Kernel1D1D and major_b == cute::UMMA::Major::K and block_m != 128)
|
||||
return false;
|
||||
|
||||
// 1D2D kernels' maximum block N is 128
|
||||
// 1D2D kernels require more friendly block Ns
|
||||
if (kernel_type == KernelType::Kernel1D2D and (block_n > 128 or 128 % block_n != 0))
|
||||
return false;
|
||||
|
||||
// Check tensor memory validity
|
||||
int sf_block_m = 0, sf_block_n = 0;
|
||||
if (kernel_type == KernelType::Kernel1D1D) {
|
||||
const auto& [sf_block_m_, sf_block_n_] = get_sf_uttcp_aligned_block_sizes(block_m, block_n, ab_dtype);
|
||||
sf_block_m = sf_block_m_, sf_block_n = sf_block_n_;
|
||||
}
|
||||
if (((2 * block_n) + (sf_block_m / 32) + (sf_block_n / 32)) > 512)
|
||||
return false;
|
||||
|
||||
// NOTES: when B is MN-major, we restrict `block_n` to multiples of 64,
|
||||
// since TMA performance degrades when `swizzle_b <= 32B` (i.e., when `block_ns % 64 != 0`), even with 3D TMA
|
||||
return major_b == cute::UMMA::Major::K or block_n % 64 == 0;
|
||||
}
|
||||
|
||||
static bool is_num_stages_legal(const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype,
|
||||
const int& num_stages,
|
||||
const int& block_m, const int& block_n, const int& block_k) {
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool should_minimize_num_sms() {
|
||||
return false;
|
||||
}
|
||||
|
||||
static std::pair<bool, bool> get_multicast_legality(const GemmType& gemm_type,
|
||||
const int& m, const int& n, const int& block_m, const int& block_n,
|
||||
const int& num_sms) {
|
||||
// TODO: support other layouts
|
||||
return {
|
||||
is_multicast_legal(m, block_m, 2, num_sms, true) and (gemm_type == GemmType::Normal or gemm_type == GemmType::KGroupedContiguous),
|
||||
false,
|
||||
};
|
||||
}
|
||||
|
||||
static ThreadConfig get_thread_config(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n) {
|
||||
return ThreadConfig::sm100(128, kernel_type == KernelType::Kernel1D1D ? 128 : block_m);
|
||||
}
|
||||
|
||||
static int get_smem_cd_size(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n,
|
||||
const int& swizzle_cd_mode,
|
||||
const at::ScalarType& cd_dtype) {
|
||||
constexpr static int layout_ad_m = 128;
|
||||
return (kernel_type == KernelType::Kernel1D1D ? std::min(block_m, layout_ad_m) : block_m) * swizzle_cd_mode * 2;
|
||||
}
|
||||
|
||||
static std::pair<int, int> get_sf_smem_size_per_stage(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n, const int& block_k,
|
||||
const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype) {
|
||||
if (ab_dtype == torch::kBFloat16)
|
||||
return {0, 0};
|
||||
|
||||
int smem_sfa_per_stage = 0;
|
||||
int smem_sfb_per_stage = 0;
|
||||
if (kernel_type == KernelType::Kernel1D1D) {
|
||||
const auto [sf_block_m, sf_block_n] = get_sf_uttcp_aligned_block_sizes(block_m, block_n, ab_dtype);
|
||||
smem_sfa_per_stage = sf_block_m * 4;
|
||||
smem_sfb_per_stage = sf_block_n * 4;
|
||||
} else {
|
||||
smem_sfa_per_stage = block_m * 4;
|
||||
smem_sfb_per_stage = 0;
|
||||
}
|
||||
return {smem_sfa_per_stage, smem_sfb_per_stage};
|
||||
}
|
||||
|
||||
static int get_extra_sfb_smem_size(const int& m, const int& n, const int& k,
|
||||
const int& block_m, const int& block_n, const int& block_k) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
static int get_barrier_smem_size(const int& num_stages) {
|
||||
// TODO: remove SF barriers for BF16 GEMMs
|
||||
// TMA full/empty barriers, with-SF full barriers, tensor memory full/empty barriers
|
||||
// NOTES: 1D2D kernel will not use the with-SF full barriers
|
||||
// NOTES: some shapes may only have 1 epilogue stage, but we still allocate space for 2 stages
|
||||
return num_stages * 8 * 3 + 2 * 8 * 2;
|
||||
}
|
||||
|
||||
static int get_tmem_ptr_smem_size() {
|
||||
return 4;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace deep_gemm
|
||||
@@ -0,0 +1,115 @@
|
||||
#pragma once
|
||||
|
||||
#include <cute/arch/mma_sm100_desc.hpp>
|
||||
// Reuse some types in the JIT modules
|
||||
#include <deep_gemm/common/types.hpp>
|
||||
|
||||
#include "common.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
struct SM90ArchSpec {
|
||||
static constexpr int smem_capacity = 232448;
|
||||
|
||||
static int get_ab_load_block_m(const MulticastConfig& multicast_config, const int& block_m) {
|
||||
return block_m;
|
||||
}
|
||||
|
||||
static int get_ab_load_block_n(const MulticastConfig& multicast_config, const int& block_n) {
|
||||
return block_n;
|
||||
}
|
||||
|
||||
static int get_cd_store_block_m(const int& block_m) {
|
||||
return block_m;
|
||||
}
|
||||
|
||||
static int get_cd_store_block_n(const int& block_n) {
|
||||
return block_n;
|
||||
}
|
||||
|
||||
static bool is_block_size_legal(const KernelType& kernel_type,
|
||||
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
||||
const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype,
|
||||
const int& block_m, const int& block_n) {
|
||||
// FP32 output does not support `block_m == 256`
|
||||
if (cd_dtype == at::kFloat and block_m == 256)
|
||||
return false;
|
||||
|
||||
// Must be some fixed block N selections
|
||||
if (block_n > 128 and kernel_type == KernelType::Kernel1D1D and (block_n != 136 or block_n != 152))
|
||||
return false;
|
||||
if (block_n > 128 and kernel_type == KernelType::Kernel1D2D and (block_n != 144 or block_n != 160))
|
||||
return false;
|
||||
|
||||
// Avoid bank conflicts for FP32 output
|
||||
if (cd_dtype == torch::kFloat and block_n % 16 == 0)
|
||||
return false;
|
||||
|
||||
// The block sizes cannot be too large (for enough registers), so at least one dim less than 128
|
||||
return block_m <= 128 or block_n <= 128;
|
||||
}
|
||||
|
||||
static bool is_num_stages_legal(const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype,
|
||||
const int& num_stages,
|
||||
const int& block_m, const int& block_n, const int& block_k) {
|
||||
// Unrolling both stages and `num_former_iters` will cause large code size
|
||||
if (ab_dtype == torch::kFloat8_e4m3fn and block_k % block_n != 0 and block_k / std::gcd(block_n, block_k) <= 4)
|
||||
return num_stages <= 4;
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool should_minimize_num_sms() {
|
||||
return true;
|
||||
}
|
||||
|
||||
static std::pair<bool, bool> get_multicast_legality(const GemmType& gemm_type,
|
||||
const int& m, const int& n, const int& block_m, const int& block_n,
|
||||
const int& num_sms) {
|
||||
return {
|
||||
is_multicast_legal(n, block_n, 2, num_sms, gemm_type == GemmType::MGroupedMasked),
|
||||
is_multicast_legal(m, block_m, 2, num_sms, false) and gemm_type != GemmType::MGroupedMasked,
|
||||
};
|
||||
}
|
||||
|
||||
static ThreadConfig get_thread_config(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n) {
|
||||
return ThreadConfig::sm90(128, (block_m == 64 ? 1 : 2) * 128);
|
||||
}
|
||||
|
||||
static int get_smem_cd_size(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n,
|
||||
const int& swizzle_cd_mode, const at::ScalarType& cd_dtype) {
|
||||
return block_m * block_n * static_cast<int>(c10::elementSize(cd_dtype));
|
||||
}
|
||||
|
||||
static std::pair<int, int> get_sf_smem_size_per_stage(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n, const int& block_k,
|
||||
const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype) {
|
||||
if (ab_dtype == torch::kBFloat16)
|
||||
return {0, 0};
|
||||
|
||||
int smem_sfa_per_stage = block_m * static_cast<int>(sizeof(float));
|
||||
int smem_sfb_per_stage = 0;
|
||||
// TODO: figure out here
|
||||
if (kernel_type == KernelType::Kernel1D1D)
|
||||
smem_sfb_per_stage = align(block_n * 4, block_k);
|
||||
return {smem_sfa_per_stage, smem_sfb_per_stage};
|
||||
}
|
||||
|
||||
static int get_extra_sfb_smem_size(const int& m, const int& n, const int& k,
|
||||
const int& block_m, const int& block_n, const int& block_k) {
|
||||
const auto& use_uniform_sfb = block_k % block_n == 0 ? 1 : 2;
|
||||
return align<int>(ceil_div(k, block_k) * static_cast<int>(sizeof(float)) * use_uniform_sfb, 8);
|
||||
}
|
||||
|
||||
static int get_barrier_smem_size(const int& num_stages) {
|
||||
// For 1D1D kernels, there is an extra barrier for accumulation
|
||||
return (num_stages + 1) * 8 * 2;
|
||||
}
|
||||
|
||||
static int get_tmem_ptr_smem_size() {
|
||||
return 0;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace deep_gemm
|
||||
Reference in New Issue
Block a user