diff --git a/CHANGELOG.md b/CHANGELOG.md index 38669385f3..f21795012d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ Documentation for Composable Kernel available at [https://rocm.docs.amd.com/proj ## Composable Kernel 1.2.0 for ROCm 7.0.0 ### Added +* Added support for B Tensor type pk_int4_t in the CK TILE weight preshuffle GEMM. * Added support for B Tensor Preshuffle in CK TILE Grouped GEMM. * Added a basic copy kernel example and supporting documentation for new CK Tile developers. * Added support for bf16, f32, and f16 for 2D and 3D NGCHW grouped convolution backward data @@ -15,6 +16,7 @@ Documentation for Composable Kernel available at [https://rocm.docs.amd.com/proj * Added support for GKCYX layout for grouped convolution backward data (NGCHW/GKCYX/NGKHW). * Added support for Stream-K version of mixed fp8/bf16 GEMM * Added support for Multiple D GEMM +* Added support for Multiple ABD GEMM * Added GEMM pipeline for microscaling (MX) FP8/FP6/FP4 data types * Added support for FP16 2:4 structured sparsity to universal GEMM. * Added support for Split K for grouped convolution backward data. @@ -29,6 +31,7 @@ Documentation for Composable Kernel available at [https://rocm.docs.amd.com/proj * Added benchmarking support for tile engine GEMM Multi D. * Added block scaling support in CK_TILE GEMM, allowing flexible use of quantization matrices from either A or B operands. * Added the row-wise column-wise quantization for CK_TILE GEMM & CK_TILE Grouped GEMM. +* Added tensor-wise quantization for CK_TILE GEMM ### Optimized diff --git a/example/ck_tile/01_fmha/README.md b/example/ck_tile/01_fmha/README.md index cb6cd44f64..7f55d7412f 100644 --- a/example/ck_tile/01_fmha/README.md +++ b/example/ck_tile/01_fmha/README.md @@ -131,4 +131,4 @@ TBD ## FP8 experimental support As described in [this blog](https://blog.hippoml.com/8bit-hippoattention-up-to-3x-faster-compared-to-flashattentionv2-8f9def90b482), we have an experimental support for fp8 fmha kernels, you can evaluate the performance by setting the arg `-prec=fp8` to the `tile_example_fmha_fwd`, on a gfx942 machine and ROCm 6.0+. -Currently we only support `-vlayout=c`( `hdim*seqlen` for V matrix) and `-squant=1`(static quantization) with `hdim=128` for fp8 now. Full feature support will come later. +Currently we only support `-vlayout=r`( `seqlen*hdim` for V matrix) for fp8 and fp8bf16 now. Full feature support will come later. diff --git a/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py b/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py index 42a9d5148a..802c9e51d7 100644 --- a/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py +++ b/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py @@ -7,7 +7,8 @@ FWD_DTYPE_MAP = { "bf16" : "FmhaFwdBf16", "fp8" : "FmhaFwdFp8", "fp8fp16": "FmhaFwdFp8Fp16", - "fp8bf16": "FmhaFwdFp8Bf16" + "fp8bf16": "FmhaFwdFp8Bf16", + "fp8fp32": "FmhaFwdFp8Fp32" } BWD_DTYPE_MAP = { diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py index d9452206e7..cfb96b7d53 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py @@ -163,7 +163,7 @@ float fmha_fwd(fmha_fwd_traits t, fmha_fwd_args a, const ck_tile::stream_config& [[maybe_unused]] auto get_num_blocks = [&](unsigned kM0) {{ return get_num_thread_blocks(a.batch, a.nhead_q, a.max_seqlen_q, kM0); }}; - + const bool has_load_tr = ck_tile::is_load_tr_supported(); {F_dispatch} @@ -248,11 +248,11 @@ class FmhaFwdApiTrait: if self.spad == 't' : return f'true /*a.seqlen_q % {self.bm0} != 0*/' # TODO: order of get_pipelines() matters! (ugly) else : return f'a.seqlen_q % {self.bm0} == 0' else: assert False - + @property def seqtune(self) -> str: if self.bm0 == 128: return 'true/*fall back to largest tile*/' # group mode only generate spad/skpad == true - else: + else: return f'a.seqlen_q <= {self.bm0}' @property @@ -351,7 +351,7 @@ class FmhaFwdPipeline: if self.F_squant == 't' : n += '_squant' else: n += '_nsquant' - + if self.F_trload == 't' : n += '_trload' else: n += '_ntrload' @@ -378,7 +378,7 @@ class FmhaFwdApiPool: "t": "has_load_tr", "f": "true" } - + per_tr_load =str() for tr_load in ["t", "f"]: per_dtypes=str() @@ -550,12 +550,16 @@ class KernelComponentFactory: (192,192) : [FmhaFwdTileSize(128, 128, 32, 192, 32, 192, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, 1)], (256,256) : [FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)], } - elif dtype == 'fp8' or dtype == 'bf8': + elif dtype == 'fp8' or dtype == 'fp8bf16': return { (64,64 ) : [FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 2, 1, 1, 2, 1, 1, 32, 32, 32, 32, 32, 32, -1)], (128,128) : [FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1)], (256,256) : [FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1)], } + elif dtype == 'fp8fp32': + return { + (128,128) : [FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1)], + } else: return None @@ -567,9 +571,9 @@ class KernelComponentFactory: # TODO: the order of List matters! the later in this list will be also be checked later # TODO: currently for qr pipeline, let 't' padding to appear later!! # TODO: how to design this more generic? - squant = 't' if dtype == 'fp8' else 'f' pipelines = [] if dtype in ['fp16', 'bf16']: + squant = 'f' for logits, mask, bias, lse, dropout, skip in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"]): if hdim == 256 and hdim_v == 256: pipelines.append(FmhaFwdPipeline('qr', 'row', 'f', 'f', 'f', 'f', logits, bias, lse, dropout, squant, mask, skip, 'f')) @@ -589,11 +593,12 @@ class KernelComponentFactory: pipelines.append(FmhaFwdPipeline('qr_async_trload', 'row', 'f', 'f', 't', 't', logits, bias, lse, dropout, squant, mask, skip, 't')) if receipt == 1 and bias != "bias": pipelines.append(FmhaFwdPipeline('qr', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, skip, 'f')) # TODO: cover arbitraty hdim - elif dtype in ['fp8', 'bf8']: + elif dtype in ['fp8', 'fp8bf16', 'fp8fp32']: # no need lse/dropout kernels - for logits, mask, bias in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys()): - pipelines.append(FmhaFwdPipeline('qr', 'col', 'f', 'f', 'f', 'f', logits, bias, 'f', 'f', squant, mask, 'f', 'f')) - elif dtype in ['fp8fp16', 'fp8bf16']: + for logits, squant, mask, bias in itertools.product(["f"], ["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys()): + pipelines.append(FmhaFwdPipeline('qr', 'row', 'f', 'f', 'f', 'f', logits, bias, 'f', 'f', squant, mask, 'f', 'f')) + pipelines.append(FmhaFwdPipeline('qr', 'row', 't', 't', 'f', 'f', logits, bias, 'f', 'f', squant, mask, 'f', 'f')) + elif dtype in ['fp8fp16', 'bf8']: # TODO None else: @@ -674,25 +679,34 @@ def get_fwd_blobs(kernel_filter : Optional[str], receipt, optdim_list, mask_impl continue # Aiter(mha_fwd) integration elif receipt == 100: - cond = dtype in ['fp16', 'bf16'] + cond = dtype in ['fp16', 'bf16', 'fp8bf16'] cond &= mode == 'batch' cond &= pipeline.F_vlayout == 'row' - cond &= pipeline.F_squant == 'f' + if dtype == 'fp8bf16': + cond &= hdim == 128 if not cond: continue # Aiter(mha_varlen_fwd) integration elif receipt == 200: - cond = dtype in ['fp16', 'bf16'] + cond = dtype in ['fp16', 'bf16', 'fp8bf16'] cond &= mode == 'group' cond &= pipeline.F_vlayout == 'row' - cond &= pipeline.F_squant == 'f' + if dtype == 'fp8bf16': + cond &= hdim == 128 if not cond: continue # aiter::mha_fwd C++ api integration elif receipt == 600: - cond = dtype in ['fp16', 'bf16'] + cond = dtype in ['fp16', 'bf16', 'fp8bf16'] cond &= pipeline.F_vlayout == 'row' - cond &= pipeline.F_squant == 'f' + if dtype == 'fp8bf16': + cond &= hdim == 128 + if not cond: + continue + elif receipt == 888: + cond = dtype in ['fp8', 'fp8bf16', 'fp8fp32'] + cond &= pipeline.F_vlayout == 'row' + cond &= hdim == 128 if not cond: continue diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py index 3b48b3d005..cee1505486 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py @@ -645,7 +645,6 @@ def get_fmha_fwd_tile_dict_from_dtype(dtype : str) -> Optional[dict]: return { '64' : FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 2, 1, 1, 2, 1, 1, 32, 32, 32, 32, 32, 32, -1), '128' : FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1), - '256' : FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 32, 32, 32, 32, -1), } else: return None diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py b/example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py index 7b93e9654c..df6b422981 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py @@ -465,14 +465,14 @@ def get_fwd_blobs(kernel_filter : Optional[str], receipt, optdim_list, mask_impl squant = 't' if dtype == 'fp8' else 'f' pipelines = [] if dtype in ['fp16', 'bf16']: - for logits, mask, bias, pagedkv, skip in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"]): - pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'col', 't', 'f', 'f', 'f', logits, bias, 'f', pagedkv, squant, mask, skip)) - pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'col', 't', 't', 'f', 'f', logits, bias, 'f', pagedkv, squant, mask, skip)) + for logits, mask, bias, pagedkv, skip in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t"], ["f"]): pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'row', 't', 'f', 'f', 'f', logits, bias, 'f', pagedkv, squant, mask, skip)) pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'row', 't', 't', 'f', 'f', logits, bias, 'f', pagedkv, squant, mask, skip)) elif dtype in ['fp8', 'bf8']: - # TODO - None + # no need lse/dropout kernels + for logits, mask, bias in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys()): + pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'row', 'f', 'f', 'f', 'f', logits, bias, 'f', 't', squant, mask, 'f')) + pipelines.append(FmhaFwdPipeline('qr_pagedkv', 'row', 't', 't', 'f', 'f', logits, bias, 'f', 't', squant, mask, 'f')) elif dtype in ['fp8fp16', 'fp8bf16']: # TODO None diff --git a/example/ck_tile/01_fmha/example_fmha_fwd.cpp b/example/ck_tile/01_fmha/example_fmha_fwd.cpp index c3bbb7a558..91cb9f55be 100644 --- a/example/ck_tile/01_fmha/example_fmha_fwd.cpp +++ b/example/ck_tile/01_fmha/example_fmha_fwd.cpp @@ -44,21 +44,15 @@ auto create_args(int argc, char* argv[]) .insert("scale_s", "0", "scale factor of S. 0 means equal to 1/sqrt(hdim).\n" - "note when squant=1, this value will be modified by range_q/k") + "note when squant=1, this value will be modified") .insert("logits_soft_cap", "0", "attention logits soft capping value.") - .insert("range_q", "16", "per-tensor quantization range of q. used if squant=1.") - .insert("range_k", "16", "per-tensor quantization range of k. used if squant=1.") - .insert("range_v", "16", "per-tensor quantization range of v. used if squant=1.") - .insert("range_p", "1", "per-tensor quantization range of p [e^(s-m)]. used if squant=1.") - .insert("range_o", "16", "per-tensor quantization range of o (p*v). used if squant=1.") .insert("squant", "auto", "if using static quantization fusion or not. auto: fp8 will default use squant, " "other will not\n" "0: no static quant(not implemented) 1: apply scale_p and scale_o with respect to " "P and O.\n" - "calculate scale_s, scale_p, scale_o according to range_q, range_k, range_v, " - "range_p, range_o") + "calculate scale_s, scale_p, scale_o auto") .insert("iperm", "1", "permute input\n" @@ -89,7 +83,7 @@ auto create_args(int argc, char* argv[]) "uf", "init method:\n ui or 0 - uniform random int\n ni - normalized random int" "\n uf or 1 - uniform random float\n nf - normalized random float" - "\n tf or 2 - trig float\n uf:q or ufq or 3 - fp8 quantization") + "\n tf or 2 - trig float\n") .insert("seed", "11939", "random seed used for initializing input tensors. 0 for " @@ -148,11 +142,6 @@ auto run(const ck_tile::ArgParser& arg_parser) uint64_t drop_offset = arg_parser.get_uint64("drop_offset"); bool drop_prefs = arg_parser.get_bool("drop_prefs"); std::string mask_str = arg_parser.get_str("mask"); - float range_q = arg_parser.get_float("range_q"); - float range_k = arg_parser.get_float("range_k"); - float range_v = arg_parser.get_float("range_v"); - float range_p = arg_parser.get_float("range_p"); - float range_o = arg_parser.get_float("range_o"); bool is_rotary_interleaved = arg_parser.get_bool("rotary_interleaved"); ck_tile::index_t num_splits = arg_parser.get_int("num_splits"); std::string init_method = arg_parser.get_str("init"); @@ -201,11 +190,6 @@ auto run(const ck_tile::ArgParser& arg_parser) drop_offset, drop_prefs, mask_str, - range_q, - range_k, - range_v, - range_p, - range_o, squant, is_rotary_interleaved, num_splits, @@ -237,6 +221,14 @@ int main(int argc, char* argv[]) { return run(arg_parser) == fwd_result::success ? 0 : -2; } + else if(data_type == "fp8bf16") + { + return run(arg_parser) == fwd_result::success ? 0 : -2; + } + else if(data_type == "fp8fp32") + { + return run(arg_parser) == fwd_result::success ? 0 : -2; + } std::cerr << "Unsupported precision: " << data_type << std::endl; return -1; } diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index df1e9e5699..c41e48e6aa 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -41,6 +41,10 @@ struct FmhaFwdFp8Bf16 { }; +struct FmhaFwdFp8Fp32 +{ +}; + template struct FmhaFwdTypeConfig; @@ -108,6 +112,38 @@ struct FmhaFwdTypeConfig using ODataType = ck_tile::bf8_t; }; +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::fp8_t; + using KDataType = ck_tile::fp8_t; + using VDataType = ck_tile::fp8_t; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = ck_tile::bf16_t; +}; + +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::fp8_t; + using KDataType = ck_tile::fp8_t; + using VDataType = ck_tile::fp8_t; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = float; +}; + struct FmhaMasks { using NoMask = ck_tile::GenericAttentionMask; diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp index 397245ab32..43f484fe14 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp @@ -50,20 +50,30 @@ auto get_elimit(std::string /*init_method*/) } template <> -auto get_elimit(std::string init_method) +auto get_elimit(std::string /*init_method*/) { - if(init_method == "ui" || init_method == "ni") - { - unsigned max_rounding_point_distance = 0; - double atol = 2e-3; - return ck_tile::make_tuple(max_rounding_point_distance, atol); - } - else - { - unsigned max_rounding_point_distance = 1; - double atol = 0.0625; - return ck_tile::make_tuple(max_rounding_point_distance, atol); - } + using TypeConfig = FmhaFwdTypeConfig; + using ODataType = typename TypeConfig::ODataType; + float o_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + double rtol = 0; + double atol = 16 * (o_dtype_max > 240 ? 2 : 1); + return ck_tile::make_tuple(rtol, atol); +} + +template <> +auto get_elimit(std::string /*init_method*/) +{ + double rtol = 1e-2; + double atol = 1.8e-1; + return ck_tile::make_tuple(rtol, atol); +} + +template <> +auto get_elimit(std::string /*init_method*/) +{ + double rtol = 1e-2; + double atol = 1.8e-1; + return ck_tile::make_tuple(rtol, atol); } int num_splits_heuristic(int batch_nhead_mblocks, int num_SMs, int max_splits) @@ -157,11 +167,6 @@ fwd_result fmha_fwd_run(mode_enum mode, uint64_t drop_offset, bool drop_prefs, std::string mask_str, - float range_q, - float range_k, - float range_v, - float range_p, - float range_o, bool squant, bool is_rotary_interleaved, ck_tile::index_t num_splits, @@ -180,6 +185,10 @@ fwd_result fmha_fwd_run(mode_enum mode, return "fp8"; else if constexpr(std::is_same_v) return "bf8"; + else if constexpr(std::is_same_v) + return "fp8bf16"; + else if constexpr(std::is_same_v) + return "fp8fp32"; else static_assert(false); }(); @@ -367,22 +376,6 @@ fwd_result fmha_fwd_run(mode_enum mode, using OaccDataType = typename TypeConfig::OaccDataType; using ODataType = typename TypeConfig::ODataType; - float q_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); - float k_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); - float v_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); - float p_dtype_max = v_dtype_max; // assume p and v is the same type - float o_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); - - float scale_p = 1.f; - float scale_o = 1.f; - - if(squant) - { - scale_s = scale_s * (range_q / q_dtype_max) * (range_k / k_dtype_max); - scale_p = p_dtype_max / range_p; - scale_o = (o_dtype_max / range_o) * (range_p / p_dtype_max) * (range_v / v_dtype_max); - } - // accumulation numbers for performance evaluation std::size_t flop = 0, num_byte = 0; auto max_seqlen_q = @@ -528,7 +521,7 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::HostTensor cache_batch_idx_host(use_cache_batch_idx ? std::array{batch} : std::array{1}); - + float max_o = 5.0; if(init_method == "ui" || init_method == "0") { ck_tile::FillUniformDistributionIntegerValue{-3.f, 3.f, next_seed()}(q_host); @@ -576,32 +569,6 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::FillTrigValue{}(vnew_host); ck_tile::FillTrigValue{}(bias_host); } - else if(init_method == "ufq" || init_method == "uf:q" || init_method == "3") - { - // suitable for fp8 quantization - if(!squant) - { - std::cerr << "init method " << init_method << " can not be used without quantization" - << std::endl; - return fwd_result::invalid_args; - } - ck_tile::FillUniformDistribution{0.f, q_dtype_max, next_seed()}(q_host); - ck_tile::FillUniformDistribution{0.f, k_dtype_max, next_seed()}(k_host); - ck_tile::FillUniformDistribution{0.f, k_dtype_max, next_seed()}(knew_host); - ck_tile::FillUniformDistribution{0.f, v_dtype_max, next_seed()}(v_host); - ck_tile::FillUniformDistribution{0.f, v_dtype_max, next_seed()}(vnew_host); - - // bias_fp8 = qscale_bias * bias_fp32 - float qscale_bias = (q_dtype_max / range_q) * (k_dtype_max / range_k); - // Assume bias is in [0.f, 1.f] in original fp32 - ck_tile::FillUniformDistribution{0.f, qscale_bias, next_seed()}(bias_host); - } - else - { - std::cerr << "Unknown value for init argument: " << init_method << std::endl; - return fwd_result::invalid_args; - } - if(bias.type == bias_enum::alibi) { auto slopes = ck_tile::get_alibi_slopes(nhead); @@ -625,8 +592,8 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::DeviceMem q_buf(q_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem k_buf(k_host.get_element_space_size_in_bytes()); - ck_tile::DeviceMem knew_buf(knew_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem v_buf(v_host.get_element_space_size_in_bytes()); + ck_tile::DeviceMem knew_buf(knew_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem vnew_buf(vnew_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem bias_buf(bias_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem lse_acc_buf(lse_acc_host.get_element_space_size_in_bytes()); @@ -650,10 +617,79 @@ fwd_result fmha_fwd_run(mode_enum mode, ck_tile::DeviceMem block_table_buf(block_table_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem cache_batch_idx_buf(cache_batch_idx_host.get_element_space_size_in_bytes()); + float scale_p = 1.f; + float scale_o = 1.f; + if(squant) + { + float q_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + float k_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + float v_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + float p_dtype_max = v_dtype_max; // assume p and v is the same type + // Q tensor + { + float max_value = ck_tile::type_convert(ck_tile::numeric::min()); + q_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + if(val > max_value) + max_value = val; + }); + + float scale = q_dtype_max / max_value; + + q_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + self(idx) = ck_tile::type_convert(val * scale); + }); + scale_s = scale_s / scale; + } + + // K tensor + { + float max_value = ck_tile::type_convert(ck_tile::numeric::min()); + k_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + if(val > max_value) + max_value = val; + }); + float scale = k_dtype_max / max_value; + k_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + self(idx) = ck_tile::type_convert(val * scale); + }); + scale_s = scale_s / scale; + } + + // V tensor + { + float max_value = ck_tile::type_convert(ck_tile::numeric::min()); + v_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + if(val > max_value) + max_value = val; + }); + + float scale = k_dtype_max / max_value; + v_host.ForEach([&](auto& self, auto idx) { + float val = ck_tile::type_convert(self(idx)); + self(idx) = ck_tile::type_convert(val * scale); + }); + + scale_o = (1.0 / p_dtype_max) / scale; + } + + scale_p = p_dtype_max; + + if constexpr(std::is_same_v) + { + float o_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + scale_o = scale_o * o_dtype_max / max_o; + } + } + q_buf.ToDevice(q_host.data()); k_buf.ToDevice(k_host.data()); - knew_buf.ToDevice(knew_host.data()); v_buf.ToDevice(v_host.data()); + knew_buf.ToDevice(knew_host.data()); vnew_buf.ToDevice(vnew_host.data()); bias_buf.ToDevice(bias_host.data()); seqstart_q.ToDevice(seqstart_q_host.data()); @@ -1103,7 +1139,9 @@ fwd_result fmha_fwd_run(mode_enum mode, lse_buf.FromDevice(lse_host.data()); randval_buf.FromDevice(randval_host.data()); - constexpr bool supports_squant = std::is_same_v; + constexpr bool supports_squant = std::is_same_v || + std::is_same_v || + std::is_same_v; auto p_compute_element_func = [&]() { if constexpr(supports_squant) @@ -1113,9 +1151,11 @@ fwd_result fmha_fwd_run(mode_enum mode, }(); auto oacc_element_func = [&]() { - if constexpr(supports_squant) + if constexpr(std::is_same_v && supports_squant) return ck_tile::composes(ck_tile::saturates{}, ck_tile::scales{scale_o}); + else if constexpr(supports_squant) + return ck_tile::scales{scale_o}; else return ck_tile::identity{}; }(); diff --git a/example/ck_tile/01_fmha/script/smoke_test_fwd.sh b/example/ck_tile/01_fmha/script/smoke_test_fwd.sh index c087a1fb3e..afd0c728c6 100755 --- a/example/ck_tile/01_fmha/script/smoke_test_fwd.sh +++ b/example/ck_tile/01_fmha/script/smoke_test_fwd.sh @@ -94,7 +94,30 @@ run_fp8_tests() { for b in 1 2 ; do for hdim in 64 128 256 ; do - run_exe -prec=fp8 -init=3 -b=$b -h=1 -d=128 -s=128 -bias=$bias -iperm=$perm -operm=$perm -vlayout=c -squant=1 -kname=$KNAME $COMMON_ARGS + $EXE -prec=fp8 -init=0 -b=$b -h=1 -d=128 -s=128 -bias=$bias -iperm=$perm -operm=$perm -vlayout=r -squant=1 -kname=$KNAME $COMMON_ARGS + + done ; done ; done ; done +} + +run_fp8bf16_tests() { + for perm in 0 1 ; do + for bias in "n" "e" "a" ; do + for b in 1 2 ; do + for hdim in 64 128 256 ; do + + $EXE -prec=fp8bf16 -init=0 -b=$b -h=1 -d=128 -s=128 -bias=$bias -iperm=$perm -operm=$perm -vlayout=r -squant=1 -kname=$KNAME $COMMON_ARGS + + done ; done ; done ; done +} + +run_fp8fp32_tests() { + for perm in 0 1 ; do + for bias in "n" "e" "a" ; do + for b in 1 2 ; do + for hdim in 64 128 256 ; do + + $EXE -prec=fp8fp32 -init=0 -b=$b -h=1 -d=128 -s=128 -bias=$bias -iperm=$perm -operm=$perm -vlayout=r -squant=1 -kname=$KNAME $COMMON_ARGS + done ; done ; done ; done } @@ -117,7 +140,9 @@ run_fp16_appendkv_tests() { set -x run_fp16_bf16_tests -# run_fp8_tests +run_fp8_tests +run_fp8bf16_tests +run_fp8fp32_tests if [ $TEST_APPENDKV -eq 1 ] ; then run_fp16_appendkv_tests diff --git a/example/ck_tile/17_grouped_gemm/quant_grouped_gemm.cpp b/example/ck_tile/17_grouped_gemm/quant_grouped_gemm.cpp index 83542e76f1..409bb173a1 100644 --- a/example/ck_tile/17_grouped_gemm/quant_grouped_gemm.cpp +++ b/example/ck_tile/17_grouped_gemm/quant_grouped_gemm.cpp @@ -13,7 +13,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/epilogue.hpp" #include "ck_tile/ops/gemm.hpp" -#include "ck_tile/ops/gemm_group_quant.hpp" +#include "ck_tile/ops/gemm_quant.hpp" #include "ck_tile/host.hpp" #include "quant_grouped_gemm.hpp" @@ -65,15 +65,15 @@ float grouped_gemm_tileloop(const ck_tile::stream_config& s, constexpr auto memory_operation = memory_operation_.value; constexpr bool transpose_c = false; - using QuantGemmProblem = ck_tile::GemmRowColQuantPipelineProblem; + using QuantGemmProblem = ck_tile::GemmRowColTensorQuantPipelineProblem; using GemmPipeline = typename PipelineTypeTraits< GemmConfig::Pipeline>::template GemmPipeline; diff --git a/example/ck_tile/20_grouped_convolution/CMakeLists.txt b/example/ck_tile/20_grouped_convolution/CMakeLists.txt index 5cb1d2650e..10332137e2 100644 --- a/example/ck_tile/20_grouped_convolution/CMakeLists.txt +++ b/example/ck_tile/20_grouped_convolution/CMakeLists.txt @@ -7,5 +7,8 @@ target_compile_options(tile_example_grouped_conv_fwd PRIVATE ${EXAMPLE_GEMM_COMP add_executable(tile_example_grouped_conv_bwd_weight EXCLUDE_FROM_ALL grouped_convolution_backward_weight.cpp) target_compile_options(tile_example_grouped_conv_bwd_weight PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) +add_executable(tile_example_grouped_conv_bwd_weight_two_stage EXCLUDE_FROM_ALL grouped_convolution_backward_weight_two_stage.cpp) +target_compile_options(tile_example_grouped_conv_bwd_weight_two_stage PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) + add_executable(tile_example_grouped_conv_bwd_data EXCLUDE_FROM_ALL grouped_convolution_backward_data.cpp) target_compile_options(tile_example_grouped_conv_bwd_data PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_data.cpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_data.cpp index c456e7a2e8..4f9362beb2 100644 --- a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_data.cpp +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_data.cpp @@ -41,8 +41,8 @@ float grouped_conv_bwd_data(const ck_tile::GroupedConvBwdDataHostArgs& args, constexpr ck_tile::index_t N_Warp_Tile = GemmWarpConfig::N_Warp_Tile; constexpr ck_tile::index_t K_Warp_Tile = GemmWarpConfig::K_Warp_Tile; - constexpr ck_tile::index_t VectorSizeA = 8; - constexpr ck_tile::index_t VectorSizeB = 8; + constexpr ck_tile::index_t VectorSizeA = 1; + constexpr ck_tile::index_t VectorSizeB = 1; constexpr ck_tile::index_t VectorSizeC = 8; // Implicit GEMM Traits @@ -51,20 +51,29 @@ float grouped_conv_bwd_data(const ck_tile::GroupedConvBwdDataHostArgs& args, ck_tile::sequence, ck_tile::sequence>; - constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default; - using TilePartitioner = ck_tile::GemmTile1DPartitioner; - using GroupedConvTraitsType = - ck_tile::GroupedConvTraits; - using CodegenPipelineProblem = - ck_tile::GemmPipelineProblem; + constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default; + using TilePartitioner = ck_tile::GemmTile1DPartitioner; + using GroupedConvTraitsType = ck_tile::GroupedConvTraits; + using CodegenPipelineProblem = ck_tile::GemmPipelineProblem< + InDataType, + WeiDataType, + AccDataType, + CodegenShape, + typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdData, + ck_tile::element_wise::PassThrough, + ck_tile::element_wise::PassThrough, + InDataType, + true, + GroupedConvTraitsType::VectorSizeA, + GroupedConvTraitsType::VectorSizeB>; using CodegenPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1; const auto Run = [&](const auto memory_operation_) { @@ -90,7 +99,7 @@ float grouped_conv_bwd_data(const ck_tile::GroupedConvBwdDataHostArgs& args, memory_operation, 1, true, - VectorSizeC>>; + GroupedConvTraitsType::VectorSizeC>>; using Kernel = ck_tile::GroupedConvolutionBackwardDataKernel, - typename DsLayout = ck_tile::tuple<>, - typename CDEElementWise = ck_tile::element_wise::PassThrough> -float grouped_conv_bwd_weight(const ck_tile::GroupedConvBwdWeightHostArgs& args, - const ck_tile::stream_config& s) -{ - constexpr int kBlockPerCu = 1; - - constexpr ck_tile::index_t M_Tile = 64; - constexpr ck_tile::index_t N_Tile = 64; - constexpr ck_tile::index_t K_Tile = 64; - - constexpr ck_tile::index_t M_Warp = 2; - constexpr ck_tile::index_t N_Warp = 2; - constexpr ck_tile::index_t K_Warp = 1; - - constexpr ck_tile::index_t M_Warp_Tile = GemmWarpConfig::M_Warp_Tile; - constexpr ck_tile::index_t N_Warp_Tile = GemmWarpConfig::N_Warp_Tile; - constexpr ck_tile::index_t K_Warp_Tile = GemmWarpConfig::K_Warp_Tile; - - constexpr ck_tile::index_t VectorSizeA = 8; - constexpr ck_tile::index_t VectorSizeB = 8; - constexpr ck_tile::index_t VectorSizeC = 8; - - // Implicit GEMM Traits - using CodegenShape = - ck_tile::TileGemmShape, - ck_tile::sequence, - ck_tile::sequence>; - - constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default; - using TilePartitioner = ck_tile::GemmTile1DPartitioner; - using GroupedConvTraitsType = - ck_tile::GroupedConvTraits; - using CodegenPipelineProblem = - ck_tile::GemmPipelineProblem; - using CodegenPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1; - - const auto Run = [&](const auto memory_operation_) { - constexpr auto memory_operation = memory_operation_.value; - - using ConvEpilogue = ck_tile::CShuffleEpilogue< - ck_tile::CShuffleEpilogueProblem>; - - using Kernel = ck_tile::GroupedConvolutionBackwardWeightKernel; - auto kargs = Kernel::MakeKernelArgs(args); - - const dim3 grids = Kernel::GridSize(kargs); - const dim3 blocks = Kernel::BlockSize(); - - if(!Kernel::IsSupportedArgument(kargs)) - { - throw std::runtime_error("Wrong! Arguments not supported! Skipping conv!\n"); - } - - if(s.log_level_ > 0) - { - std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n' - << "shape: " << CodegenShape::GetName() << '\n' - << "problem: " << CodegenPipelineProblem::GetName() << '\n' - << "pipeline: " << CodegenPipeline::GetName() << '\n' - << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}" - << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}" - << '\n' - << "Vector size A: " << CodegenPipeline::GetVectorSizeA() - << ", Vector size B: " << CodegenPipeline::GetVectorSizeB() - << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl; - } - - float ave_time = ck_tile::launch_kernel_time_mask( - s, - Kernel::Preprocess(kargs, s), - ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); - - return ave_time; - }; - - if(args.k_batch == 1) - { - return Run(ck_tile::integral_constant{}); - } - else - { - return Run(ck_tile::integral_constant{}); - } -} - +#include "grouped_convolution_backward_weight_invoker.hpp" #include "run_grouped_convolution_bwd_weight_example.inc" -template -int run_grouped_conv_bwd_weight_example_prec_type( - std::string in_layout, std::string wei_layout, std::string out_layout, int argc, char* argv[]) -{ - using NWGC = ck_tile::tensor_layout::convolution::NWGC; - using NHWGC = ck_tile::tensor_layout::convolution::NHWGC; - using NDHWGC = ck_tile::tensor_layout::convolution::NDHWGC; - - using GKXC = ck_tile::tensor_layout::convolution::GKXC; - using GKYXC = ck_tile::tensor_layout::convolution::GKYXC; - using GKZYXC = ck_tile::tensor_layout::convolution::GKZYXC; - - using NWGK = ck_tile::tensor_layout::convolution::NWGK; - using NHWGK = ck_tile::tensor_layout::convolution::NHWGK; - using NDHWGK = ck_tile::tensor_layout::convolution::NDHWGK; - - if(in_layout == "NWGC" && wei_layout == "GKXC" && out_layout == "NWGK") - { - return run_grouped_conv_bwd_weight_example_with_layouts{}, - GemmWarpConfig, - InPrecType, - WeiPrecType, - OutPrecType>( - argc, argv, NWGC{}, GKXC{}, NWGK{}); - } - else if(in_layout == "NHWGC" && wei_layout == "GKYXC" && out_layout == "NHWGK") - { - return run_grouped_conv_bwd_weight_example_with_layouts{}, - GemmWarpConfig, - InPrecType, - WeiPrecType, - OutPrecType>( - argc, argv, NHWGC{}, GKYXC{}, NHWGK{}); - } - else if(in_layout == "NDHWGC" && wei_layout == "GKZYXC" && out_layout == "NDHWGK") - { - return run_grouped_conv_bwd_weight_example_with_layouts{}, - GemmWarpConfig, - InPrecType, - WeiPrecType, - OutPrecType>( - argc, argv, NDHWGC{}, GKZYXC{}, NDHWGK{}); - } - else - { - throw std::runtime_error("Unsupported memory layout!"); - } -} - template -int run_grouped_conv_bwd_weight_example(int argc, char* argv[]) +int run_grouped_conv_bwd_weight_example(ck_tile::ArgParser& arg_parser) { - auto [result, arg_parser] = create_args(argc, argv); - if(!result) - return -1; + using Invoker = GroupedConvolutionBackwardWeightInvoker; std::string data_type = arg_parser.get_str("prec"); std::string in_layout = arg_parser.get_str("in_layout"); @@ -208,13 +26,17 @@ int run_grouped_conv_bwd_weight_example(int argc, char* argv[]) if(data_type == "fp16") { - return run_grouped_conv_bwd_weight_example_prec_type( - in_layout, wei_layout, out_layout, argc, argv); + return run_grouped_conv_bwd_weight_example_prec_type( + in_layout, wei_layout, out_layout, arg_parser); } else if(data_type == "bf16") { - return run_grouped_conv_bwd_weight_example_prec_type( - in_layout, wei_layout, out_layout, argc, argv); + return run_grouped_conv_bwd_weight_example_prec_type( + in_layout, wei_layout, out_layout, arg_parser); } else { @@ -224,9 +46,22 @@ int run_grouped_conv_bwd_weight_example(int argc, char* argv[]) int main(int argc, char* argv[]) { + + auto [result, arg_parser] = create_args(argc, argv); + if(!result) + return -1; + + try + { #if CK_TILE_USE_WMMA - return !run_grouped_conv_bwd_weight_example(argc, argv); + return !run_grouped_conv_bwd_weight_example(arg_parser); #else - return !run_grouped_conv_bwd_weight_example(argc, argv); + return !run_grouped_conv_bwd_weight_example(arg_parser); #endif + } + catch(const std::runtime_error& e) + { + std::cerr << "Runtime error: " << e.what() << '\n'; + return EXIT_FAILURE; + } } diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_invoker.hpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_invoker.hpp new file mode 100644 index 0000000000..f46707d1d2 --- /dev/null +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_invoker.hpp @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +#pragma once + +#include "grouped_convolution_utils.hpp" + +struct GroupedConvolutionBackwardWeightInvoker +{ + template , + typename DsLayout = ck_tile::tuple<>, + typename CDEElementWise = ck_tile::element_wise::PassThrough> + static float grouped_conv_bwd_weight(const ck_tile::GroupedConvBwdWeightHostArgs& args, + const ck_tile::stream_config& s) + { + constexpr int kBlockPerCu = 1; + + constexpr ck_tile::index_t M_Tile = 64; + constexpr ck_tile::index_t N_Tile = 64; + constexpr ck_tile::index_t K_Tile = 64; + + constexpr ck_tile::index_t M_Warp = 2; + constexpr ck_tile::index_t N_Warp = 2; + constexpr ck_tile::index_t K_Warp = 1; + + constexpr ck_tile::index_t M_Warp_Tile = GemmWarpConfig::M_Warp_Tile; + constexpr ck_tile::index_t N_Warp_Tile = GemmWarpConfig::N_Warp_Tile; + constexpr ck_tile::index_t K_Warp_Tile = GemmWarpConfig::K_Warp_Tile; + + constexpr ck_tile::index_t VectorSizeA = 1; + constexpr ck_tile::index_t VectorSizeB = 1; + constexpr ck_tile::index_t VectorSizeC = 8; + + // Implicit GEMM Traits + using CodegenShape = + ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + + constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default; + using TilePartitioner = ck_tile::GemmTile1DPartitioner; + using GroupedConvTraitsType = ck_tile::GroupedConvTraits; + using CodegenPipelineProblem = ck_tile::GemmPipelineProblem< + InDataType, + WeiDataType, + AccDataType, + CodegenShape, + typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdWeight, + ck_tile::element_wise::PassThrough, + ck_tile::element_wise::PassThrough, + InDataType, + true, + GroupedConvTraitsType::VectorSizeA, + GroupedConvTraitsType::VectorSizeB>; + using CodegenPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1; + + const auto Run = [&](const auto memory_operation_) { + constexpr auto memory_operation = memory_operation_.value; + + using ConvEpilogue = ck_tile::CShuffleEpilogue>; + + using Kernel = ck_tile::GroupedConvolutionBackwardWeightKernel; + auto kargs = Kernel::MakeKernelArgs(args); + + const dim3 grids = Kernel::GridSize(kargs); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Wrong! Arguments not supported! Skipping conv!\n"); + } + + if(s.log_level_ > 0) + { + std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n' + << "shape: " << CodegenShape::GetName() << '\n' + << "problem: " << CodegenPipelineProblem::GetName() << '\n' + << "pipeline: " << CodegenPipeline::GetName() << '\n' + << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}" + << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z + << "}" << '\n' + << "Vector size A: " << CodegenPipeline::GetVectorSizeA() + << ", Vector size B: " << CodegenPipeline::GetVectorSizeB() + << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl; + } + + float ave_time = ck_tile::launch_kernel_time_mask( + s, + Kernel::Preprocess(kargs, s), + ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); + + return ave_time; + }; + + if(args.k_batch == 1) + { + return Run(ck_tile::integral_constant{}); + } + else + { + return Run(ck_tile::integral_constant{}); + } + } +}; diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage.cpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage.cpp new file mode 100644 index 0000000000..82068156e5 --- /dev/null +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage.cpp @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include + +#include +#include +#include +#include +#include + +#include "ck_tile/host.hpp" +#include "grouped_convolution_utils.hpp" +#include "grouped_convolution_backward_weight_two_stage_invoker.hpp" +#include "run_grouped_convolution_bwd_weight_example.inc" + +template +int run_grouped_conv_bwd_weight_example(ck_tile::ArgParser& arg_parser) +{ + using Invoker = GroupedConvolutionBackwardWeightTwoStageInvoker; + + std::string data_type = arg_parser.get_str("prec"); + std::string in_layout = arg_parser.get_str("in_layout"); + std::string wei_layout = arg_parser.get_str("wei_layout"); + std::string out_layout = arg_parser.get_str("out_layout"); + + if(data_type == "fp16") + { + return run_grouped_conv_bwd_weight_example_prec_type( + in_layout, wei_layout, out_layout, arg_parser); + } + else if(data_type == "bf16") + { + return run_grouped_conv_bwd_weight_example_prec_type( + in_layout, wei_layout, out_layout, arg_parser); + } + else + { + throw std::runtime_error("Unsupported data type for this operation!"); + } +} + +int main(int argc, char* argv[]) +{ + + auto [result, arg_parser] = create_args(argc, argv); + if(!result) + return -1; + + try + { +#if CK_TILE_USE_WMMA + return !run_grouped_conv_bwd_weight_example(arg_parser); +#else + return !run_grouped_conv_bwd_weight_example(arg_parser); +#endif + } + catch(const std::runtime_error& e) + { + std::cerr << "Runtime error: " << e.what() << '\n'; + return EXIT_FAILURE; + } +} diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage_invoker.hpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage_invoker.hpp new file mode 100644 index 0000000000..a8f321e7f9 --- /dev/null +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage_invoker.hpp @@ -0,0 +1,215 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +#pragma once + +#include "grouped_convolution_utils.hpp" + +struct GroupedConvolutionBackwardWeightTwoStageInvoker +{ + template , + typename DsLayout = ck_tile::tuple<>, + typename CDEElementWise = ck_tile::element_wise::PassThrough> + static float grouped_conv_bwd_weight(const ck_tile::GroupedConvBwdWeightHostArgs& args, + const ck_tile::stream_config& s) + { + using WorkspaceDataType = float; + + constexpr int kBlockPerCu = 1; + + constexpr ck_tile::index_t M_Tile = 64; + constexpr ck_tile::index_t N_Tile = 64; + constexpr ck_tile::index_t K_Tile = 64; + + constexpr ck_tile::index_t M_Warp = 2; + constexpr ck_tile::index_t N_Warp = 2; + constexpr ck_tile::index_t K_Warp = 1; + + constexpr ck_tile::index_t M_Warp_Tile = GemmWarpConfig::M_Warp_Tile; + constexpr ck_tile::index_t N_Warp_Tile = GemmWarpConfig::N_Warp_Tile; + constexpr ck_tile::index_t K_Warp_Tile = GemmWarpConfig::K_Warp_Tile; + + constexpr ck_tile::index_t VectorSizeA = 1; + constexpr ck_tile::index_t VectorSizeB = 1; + constexpr ck_tile::index_t VectorSizeC = 1; + + // Implicit GEMM Traits + using CodegenShape = + ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + + constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default; + using TilePartitioner = ck_tile::GemmTile1DPartitioner; + using GroupedConvTraitsType = ck_tile::GroupedConvTraits; + using CodegenPipelineProblem = ck_tile::GemmPipelineProblem< + OutDataType, // A: Out + InDataType, // B: In + AccDataType, + CodegenShape, + typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsBwdWeight, + ck_tile::element_wise::PassThrough, + ck_tile::element_wise::PassThrough, + InDataType, + true, + GroupedConvTraitsType::VectorSizeA, + GroupedConvTraitsType::VectorSizeB>; + using CodegenPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1; + + const auto Run = [&](const auto memory_operation_) { + constexpr auto memory_operation = memory_operation_.value; + + using ConvEpilogue = ck_tile::CShuffleEpilogue>; + + using Kernel = ck_tile::GroupedConvolutionBackwardWeightKernel; + + const ck_tile::index_t spatial_lengths_accum = + std::accumulate(args.filter_spatial_lengths_.begin(), + args.filter_spatial_lengths_.end(), + 1, + std::multiplies()); + ck_tile::DeviceMem ws_m_n_dev_buf(args.G_ * args.K_ * args.C_ * spatial_lengths_accum * + sizeof(WorkspaceDataType)); + ck_tile::GroupedConvBwdWeightHostArgs ws_args = + ck_tile::GroupedConvBwdWeightHostArgs(args); + auto c_ptr = ws_args.wei_ptr; + ws_args.wei_ptr = ws_m_n_dev_buf.GetDeviceBuffer(); + auto kargs = Kernel::MakeKernelArgs(ws_args); + + const dim3 grids = Kernel::GridSize(kargs); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Wrong! Arguments not supported! Skipping conv!\n"); + } + + using XElementwiseOperation = ck_tile::element_wise::UnaryConvert; + using BlockTile = ck_tile::sequence<2048>; + using BlockWarps = ck_tile::sequence<8>; + using WarpTile = ck_tile::sequence<64>; + + using ElementwiseShape = + ck_tile::ElementWiseShape; + using Problem = ck_tile::ElementWisePipelineProblem; + using ElementwiseKernel = + ck_tile::ElementWiseKernel; + + ck_tile::index_t total_elements = 1; + std::vector shape = { + static_cast(args.G_ * args.K_), + static_cast(args.C_ * spatial_lengths_accum)}; + + for(auto d : shape) + total_elements *= d; + + const ck_tile::index_t kBlockSize = ElementwiseKernel::BlockSize(); + + constexpr ck_tile::index_t elements_per_block = BlockTile::at(ck_tile::number<0>{}); + ck_tile::index_t kGridSize = + (total_elements + elements_per_block - 1) / elements_per_block; + + auto input_tensors = + ck_tile::make_tuple(static_cast(ws_args.wei_ptr)); + auto input_size = ck_tile::make_tuple(shape[0], shape[1]); + + // Check if the kernel configuration is supported + if(!ElementwiseKernel::IsSupportedArgument(input_size)) + { + throw std::runtime_error( + "Wrong! Elementwise arguments not supported! Skipping gemm!\n"); + } + + if(s.log_level_ > 0) + { + std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n' + << "shape: " << CodegenShape::GetName() << '\n' + << "problem: " << CodegenPipelineProblem::GetName() << '\n' + << "pipeline: " << CodegenPipeline::GetName() << '\n' + << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}" + << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z + << "}" << '\n' + << "Vector size A: " << CodegenPipeline::GetVectorSizeA() + << ", Vector size B: " << CodegenPipeline::GetVectorSizeB() + << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl; + } + + auto preprocess = [&]() { + if(args.k_batch > 1) + ck_tile::hip_check_error( + hipMemsetAsync(ws_args.wei_ptr, + 0, + shape[0] * shape[1] * sizeof(WorkspaceDataType), + s.stream_id_)); + }; + + return ck_tile::launch_kernel_time_mask( + s, + preprocess, + ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs), + ck_tile::make_kernel(ElementwiseKernel{}, + kGridSize, + kBlockSize, + 0, + input_size, + ck_tile::make_tuple(shape[1], 1), // Input Stride + ck_tile::make_tuple(shape[1], 1), // Output Stride + input_tensors, + static_cast(c_ptr))); + }; + + if(args.k_batch == 1) + { + return Run(ck_tile::integral_constant{}); + } + else + { + return Run(ck_tile::integral_constant{}); + } + } +}; diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_forward.cpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_forward.cpp index d74fa1b292..cebfa90579 100644 --- a/example/ck_tile/20_grouped_convolution/grouped_convolution_forward.cpp +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_forward.cpp @@ -50,20 +50,29 @@ float grouped_conv_fwd(const ck_tile::GroupedConvFwdHostArgs& args, const ck_til ck_tile::sequence, ck_tile::sequence>; - constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default; - using TilePartitioner = ck_tile::GemmTile1DPartitioner; - using GroupedConvTraitsType = - ck_tile::GroupedConvTraits; - using CodegenPipelineProblem = - ck_tile::GemmPipelineProblem; + constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default; + using TilePartitioner = ck_tile::GemmTile1DPartitioner; + using GroupedConvTraitsType = ck_tile::GroupedConvTraits; + using CodegenPipelineProblem = ck_tile::GemmPipelineProblem< + InDataType, + WeiDataType, + AccDataType, + CodegenShape, + typename GroupedConvTraitsType::GroupedConvImplicitGemmTraitsFwd, + ck_tile::element_wise::PassThrough, + ck_tile::element_wise::PassThrough, + InDataType, + true, + GroupedConvTraitsType::VectorSizeA, + GroupedConvTraitsType::VectorSizeB>; using CodegenPipeline = ck_tile::GemmPipelineAGmemBGmemCRegV1; const auto Run = [&](const auto memory_operation_) { @@ -89,7 +98,7 @@ float grouped_conv_fwd(const ck_tile::GroupedConvFwdHostArgs& args, const ck_til memory_operation, 1, true, - VectorSizeC>>; + GroupedConvTraitsType::VectorSizeC>>; using Kernel = ck_tile::GroupedConvolutionForwardKernel( + float ave_time = Invoker::template grouped_conv_bwd_weight( args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); - std::size_t flop = args.GetFlops(); - std::size_t num_byte = args.GetByte(); - float tflops = static_cast(flop) / 1.E9 / ave_time; - float gb_per_sec = num_byte / 1.E6 / ave_time; - - std::cout << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " - << std::endl; - return ave_time; } template -int run_grouped_conv_bwd_weight_example_with_layouts( - int argc, char* argv[], const InLayout, const WeiLayout, const OutLayout) +int run_grouped_conv_bwd_weight_example_with_layouts(ck_tile::ArgParser& arg_parser, + const InLayout, + const WeiLayout, + const OutLayout) { - auto [result, arg_parser] = create_args(argc, argv); - if(!result) - return -1; - using AccDataType = float; std::vector filter_spatial_lengths; @@ -138,17 +130,27 @@ int run_grouped_conv_bwd_weight_example_with_layouts( std::cout << "weight: " << weight.mDesc << std::endl; std::cout << "output: " << output.mDesc << std::endl; - invoke_grouped_conv_bwd_weight(args, n_warmup, n_repeat); + float ave_time = invoke_grouped_conv_bwd_weight(args, n_warmup, n_repeat); weight_dev_buf.FromDevice(weight.data()); + + std::size_t flop = args.GetFlops(); + std::size_t num_byte = args.GetByte(); + float tflops = static_cast(flop) / 1.E9 / ave_time; + float gb_per_sec = num_byte / 1.E6 / ave_time; + + std::cout << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " + << std::endl; + bool pass = true; if(arg_parser.get_int("v") == 1) @@ -189,3 +191,61 @@ int run_grouped_conv_bwd_weight_example_with_layouts( return pass; } + +template +int run_grouped_conv_bwd_weight_example_prec_type(std::string in_layout, + std::string wei_layout, + std::string out_layout, + ck_tile::ArgParser& arg_parser) +{ + using NWGC = ck_tile::tensor_layout::convolution::NWGC; + using NHWGC = ck_tile::tensor_layout::convolution::NHWGC; + using NDHWGC = ck_tile::tensor_layout::convolution::NDHWGC; + + using GKXC = ck_tile::tensor_layout::convolution::GKXC; + using GKYXC = ck_tile::tensor_layout::convolution::GKYXC; + using GKZYXC = ck_tile::tensor_layout::convolution::GKZYXC; + + using NWGK = ck_tile::tensor_layout::convolution::NWGK; + using NHWGK = ck_tile::tensor_layout::convolution::NHWGK; + using NDHWGK = ck_tile::tensor_layout::convolution::NDHWGK; + + if(in_layout == "NWGC" && wei_layout == "GKXC" && out_layout == "NWGK") + { + return run_grouped_conv_bwd_weight_example_with_layouts{}, + GemmWarpConfig, + Invoker, + InPrecType, + WeiPrecType, + OutPrecType>( + arg_parser, NWGC{}, GKXC{}, NWGK{}); + } + else if(in_layout == "NHWGC" && wei_layout == "GKYXC" && out_layout == "NHWGK") + { + return run_grouped_conv_bwd_weight_example_with_layouts{}, + GemmWarpConfig, + Invoker, + InPrecType, + WeiPrecType, + OutPrecType>( + arg_parser, NHWGC{}, GKYXC{}, NHWGK{}); + } + else if(in_layout == "NDHWGC" && wei_layout == "GKZYXC" && out_layout == "NDHWGK") + { + return run_grouped_conv_bwd_weight_example_with_layouts{}, + GemmWarpConfig, + Invoker, + InPrecType, + WeiPrecType, + OutPrecType>( + arg_parser, NDHWGC{}, GKZYXC{}, NDHWGK{}); + } + else + { + throw std::runtime_error("Unsupported memory layout!"); + } +} diff --git a/example/ck_tile/22_gemm_multi_abd/CMakeLists.txt b/example/ck_tile/22_gemm_multi_abd/CMakeLists.txt new file mode 100644 index 0000000000..f382e0cf45 --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/CMakeLists.txt @@ -0,0 +1 @@ +add_executable(tile_example_gemm_multi_abd_fp16 EXCLUDE_FROM_ALL gemm_multi_abd_fp16.cpp) diff --git a/example/ck_tile/22_gemm_multi_abd/README.md b/example/ck_tile/22_gemm_multi_abd/README.md new file mode 100644 index 0000000000..c272df3fb5 --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/README.md @@ -0,0 +1,35 @@ +#Multiple ABD GEMM + +This folder contains example for Multiple ABD GEMM using ck_tile tile-programming implementation. + +## build +``` +#in the root of ck_tile +mkdir build && cd build +#you can replace < arch> with the appropriate architecture(for example gfx90a or gfx942) or \ + leave it blank +sh ../script/cmake-ck-dev.sh ../ +#The basic pipeline method on the gemm calculation +make tile_example_gemm_multi_abd_fp16 -j +``` +This will result in an executable `build/bin/tile_example_gemm_multi_abd_fp16` + +## example +``` +args: + -m M dimensions - (Default: 3840) + -n N dimensions - (Default: 4096) + -k K dimensions - (Default: 4096) +-as_layout Tensor A layout (default:R) +-bs_layout Tensor B layout (default:C) +-ds_layout Tensor D layout (default:R) +-e_layout Tensor E layout (default:R) +-stride_as Tensor A strides - (Default: 0) +-stride_bs Tensor B strides - (Default: 0) +-stride_e Tensor C strides - (Default: 0) +-stride_ds Tensor D strides - (Default: 0) +-validate 0. No validation, 1. Validation on GPU. (Default: 1) + -warmup Number of iterations before benchmark the kernel. (Default: 10) + -repeat Number of iterations to benchmark the kernel. (Default: 100) + -kbatch kbatch for SplitK. (Default: 1) +``` \ No newline at end of file diff --git a/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.cpp b/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.cpp new file mode 100644 index 0000000000..6d955c3a09 --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.cpp @@ -0,0 +1,184 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include + +#include +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/epilogue.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/host.hpp" +#include "gemm_multi_abd_fp16.hpp" +#include "utils.hpp" + +template +auto gemm_multi_abd(const gemm_multi_abd_kargs& args, const ck_tile::stream_config& s) -> float +{ + constexpr ck_tile::index_t M_Tile = GemmConfig::M_Tile; + constexpr ck_tile::index_t N_Tile = GemmConfig::N_Tile; + constexpr ck_tile::index_t K_Tile = GemmConfig::K_Tile; + + constexpr ck_tile::index_t M_Warp = GemmConfig::M_Warp; + constexpr ck_tile::index_t N_Warp = GemmConfig::N_Warp; + constexpr ck_tile::index_t K_Warp = GemmConfig::K_Warp; + + constexpr ck_tile::index_t M_Warp_Tile = GemmConfig::M_Warp_Tile; + constexpr ck_tile::index_t N_Warp_Tile = GemmConfig::N_Warp_Tile; + constexpr ck_tile::index_t K_Warp_Tile = GemmConfig::K_Warp_Tile; + + constexpr bool DoubleSmemBuffer = GemmConfig::DoubleSmemBuffer; + constexpr bool kPadM = false; + constexpr bool kPadN = false; + constexpr bool kPadK = false; + + constexpr bool TransposeC = false; + + constexpr int kBlockPerCu = 1; + constexpr ck_tile::index_t TileParitionerGroupNum = 8; + constexpr ck_tile::index_t TileParitionerM01 = 4; + + using GemmShape = + ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + + using TilePartitioner = ck_tile:: + GemmSpatiallyLocalTilePartitioner; + + using Traits = ck_tile::TileGemmTraits; + + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits; + using GemmPipelineProblem = + ck_tile::GemmPipelineProblem; + + using BaseGemmPipeline = typename PipelineTypeTraits< + GemmConfig::Pipeline>::template UniversalGemmPipeline; + + const ck_tile::index_t k_grain = args.k_batch * K_Tile; + const ck_tile::index_t K_split = (args.K + k_grain - 1) / k_grain * K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + float ave_time{0}; + + const auto Run = + [&](const auto has_hot_loop_, const auto tail_number_, const auto memory_operation_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + constexpr auto scheduler = GemmConfig::Scheduler; + constexpr auto memory_operation = memory_operation_.value; + + using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem; + + using GemmPipeline = typename PipelineTypeTraits< + GemmConfig::Pipeline>::template GemmPipeline; + + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem>; + + using Kernel = ck_tile::GemmKernelMultiABD; + auto kargs = Kernel::MakeKernelArgs(args); + + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n"); + } + + if(s.log_level_ > 0) + { + std::cout << "Launching kernel with args:" << " grid: {" << grids.x << ", " + << grids.y << ", " << grids.z << "}" << ", blocks: {" << blocks.x << ", " + << blocks.y << ", " << blocks.z << "}" << std::endl; + } + + ave_time = ck_tile::launch_kernel( + s, ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); + return ave_time; + }; + + const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) { + if(args.k_batch == 1) + { + Run(has_hot_loop_, + tail_number_, + ck_tile::integral_constant{}); + } + else + { + Run(has_hot_loop_, + tail_number_, + ck_tile::integral_constant{}); + } + }; + + BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num); + + return ave_time; +} + +#include "run_gemm_multi_abd_fp16_example.inc" + +int main(int argc, char* argv[]) +{ +#if CK_TILE_USE_WMMA + return !run_multiple_abd_gemm_example(argc, argv); +#else + return !run_multiple_abd_gemm_example(argc, argv); +#endif +} diff --git a/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.hpp b/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.hpp new file mode 100644 index 0000000000..35bc232eca --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/gemm_multi_abd_fp16.hpp @@ -0,0 +1,186 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" + +#define CK_TILE_PIPELINE_COMPUTE_V3 1 +#define CK_TILE_PIPELINE_MEMORY 2 +#define CK_TILE_PIPELINE_COMPUTE_V4 3 + +#ifndef CK_TILE_PIPELINE_DEFAULT +#define CK_TILE_PIPELINE_DEFAULT CK_TILE_PIPELINE_COMPUTE_V3 +#endif + +using A0DataType = ck_tile::half_t; +using A1DataType = ck_tile::half_t; + +using B0DataType = ck_tile::half_t; +using B1DataType = ck_tile::half_t; + +using D0DataType = ck_tile::half_t; +using D1DataType = ck_tile::half_t; + +using EDataType = ck_tile::half_t; + +using AsDataType = ck_tile::tuple; +using BsDataType = ck_tile::tuple; +using DsDataType = ck_tile::tuple; + +using AccDataType = float; + +struct GemmConfigMemory +{ + // Memory friendly for Interwave scheduler + static constexpr ck_tile::index_t M_Tile = 128; + static constexpr ck_tile::index_t N_Tile = 32; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 4; + static constexpr ck_tile::index_t N_Warp = 1; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 32; + static constexpr ck_tile::index_t N_Warp_Tile = 32; + static constexpr ck_tile::index_t K_Warp_Tile = 8; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_MEMORY; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Interwave; +}; + +struct GemmConfigV3 +{ + // Compute friendly for Intrawave scheduler + static constexpr ck_tile::index_t M_Tile = 256; + static constexpr ck_tile::index_t N_Tile = 256; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 2; + static constexpr ck_tile::index_t N_Warp = 2; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 32; + static constexpr ck_tile::index_t N_Warp_Tile = 32; + static constexpr ck_tile::index_t K_Warp_Tile = 16; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; +}; + +struct GemmConfigV4 +{ + // Compute friendly for Intrawave scheduler + // Using the ping pong reader in the lds level + static constexpr ck_tile::index_t M_Tile = 256; + static constexpr ck_tile::index_t N_Tile = 256; + static constexpr ck_tile::index_t K_Tile = 32; + + static constexpr ck_tile::index_t M_Warp = 2; + static constexpr ck_tile::index_t N_Warp = 2; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 32; + static constexpr ck_tile::index_t N_Warp_Tile = 32; + static constexpr ck_tile::index_t K_Warp_Tile = 16; + + static constexpr bool DoubleSmemBuffer = true; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V4; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; +}; + +struct GemmConfigV3_Wmma +{ + // Compute friendly for Intrawave scheduler + static constexpr ck_tile::index_t M_Tile = 128; + static constexpr ck_tile::index_t N_Tile = 128; + static constexpr ck_tile::index_t K_Tile = 64; + + static constexpr ck_tile::index_t M_Warp = 2; + static constexpr ck_tile::index_t N_Warp = 2; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = 16; + + static constexpr bool DoubleSmemBuffer = false; + static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; +}; + +template +struct PipelineTypeTraits; + +template <> +struct PipelineTypeTraits +{ + template + using GemmPipeline = ck_tile::GemmPipelineAgBgCrMem; + template + using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrMem; +}; + +template <> +struct PipelineTypeTraits +{ + template + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3; + template + using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; +}; + +template <> +struct PipelineTypeTraits +{ + template + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV4; + template + using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV4; +}; + +auto create_args(int argc, char* argv[]) +{ + ck_tile::ArgParser arg_parser; + arg_parser.insert("m", "3840", "m dimension") + .insert("n", "4096", "n dimension") + .insert("k", "4096", "k dimension") + .insert("as_layout", "R", "As tensor data layout - Row by default") + .insert("bs_layout", "C", "Bs tensor data layout - Col by default") + .insert("ds_layout", "R", "Ds tensor data layout - Row by default") + .insert("e_layout", "R", "E tensor data layout - Row by default") + .insert("stride_as", "0", "Tensor A stride") + .insert("stride_bs", "0", "Tensor B stride") + .insert("stride_ds", "0", "Tensor Ds stride") + .insert("stride_e", "0", "Tensor E stride") + .insert("v", "1", "0. No validation, 1. Validation on GPU") + .insert("warmup", "50", "number of iterations before benchmark the kernel") + .insert("repeat", "100", "number of iterations to benchmark the kernel") + .insert("kbatch", "1", "kbatch for SplitK"); + + bool result = arg_parser.parse(argc, argv); + return std::make_tuple(result, arg_parser); +} +using gemm_multi_abd_kargs = + ck_tile::GemmMultiABDHostArgs; + +template +float gemm_multi_abd(const gemm_multi_abd_kargs& kargs, const ck_tile::stream_config& s); diff --git a/example/ck_tile/22_gemm_multi_abd/run_gemm_multi_abd_fp16_example.inc b/example/ck_tile/22_gemm_multi_abd/run_gemm_multi_abd_fp16_example.inc new file mode 100644 index 0000000000..881961c9db --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/run_gemm_multi_abd_fp16_example.inc @@ -0,0 +1,311 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once +#include + +template +float invoke_gemm_multi_abd(const std::array& as_m_k_dev_buf, + const std::array& bs_k_n_dev_buf, + const std::array& ds_m_n_dev_buf, + void* e_m_n_dev_buf, + ck_tile::index_t M, + ck_tile::index_t N, + ck_tile::index_t K, + const std::array& StrideAs, + const std::array& StrideBs, + const std::array& StrideDs, + ck_tile::index_t StrideE, + int n_warmup, + int n_repeat, + int k_batch) +{ + gemm_multi_abd_kargs gemm_descs({as_m_k_dev_buf, + bs_k_n_dev_buf, + ds_m_n_dev_buf, + e_m_n_dev_buf, + k_batch, + M, + N, + K, + StrideAs, + StrideBs, + StrideDs, + StrideE}); + + float ave_time = gemm_multi_abd( + gemm_descs, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); + + std::string op_name{"Gemm Multiple-ABD"}; + + std::size_t flop = 0, num_btype = 0; + + flop += std::size_t(2) * M * N * K; + + num_btype += + sizeof(A0DataType) * M * K + sizeof(B0DataType) * K * N + sizeof(EDataType) * M * N; + + float tflops = static_cast(flop) / 1.E9 / ave_time; + float gb_per_sec = num_btype / 1.E6 / ave_time; + + std::cout << "Run Gemm Multiple-ABD kernel with:\n"; + std::cout << "M =" << M << " N =" << N << " K =" << K << "\n"; + std::cout << "StrideA = " << StrideAs[0] << " StrideB = " << StrideBs[0] + << " StrideE = " << StrideE << "\n"; + std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " + << "\n"; + + return ave_time; +} + +template +int run_gemm_multi_abd_example_with_layouts(int argc, + char* argv[], + const A0Layout a0_layout = A0Layout{}, + const A1Layout a1_layout = A1Layout{}, + const B0Layout b0_layout = B0Layout{}, + const B1Layout b1_layout = B1Layout{}, + const D0Layout d0_layout = D0Layout{}, + const D1Layout d1_layout = D1Layout{}, + const ELayout e_layout = ELayout{}) +{ + auto [result, arg_parser] = create_args(argc, argv); + if(!result) + { + return -1; + } + using AElementWiseFn = ck_tile::element_wise::AddScale; + using BElementWiseFn = ck_tile::element_wise::AddScale; + using CDEElementWiseFn = ck_tile::element_wise::MultiDMultiply; + using AsLayout = ck_tile::tuple; + using BsLayout = ck_tile::tuple; + using DsLayout = ck_tile::tuple; + + ck_tile::index_t M = arg_parser.get_int("m"); + ck_tile::index_t N = arg_parser.get_int("n"); + ck_tile::index_t K = arg_parser.get_int("k"); + + ck_tile::index_t StrideA = arg_parser.get_int("stride_as"); + ck_tile::index_t StrideB = arg_parser.get_int("stride_bs"); + ck_tile::index_t StrideD = arg_parser.get_int("stride_ds"); + ck_tile::index_t StrideE = arg_parser.get_int("stride_e"); + + ck_tile::index_t StrideA0 = StrideA; + ck_tile::index_t StrideA1 = StrideA; + + ck_tile::index_t StrideB0 = StrideB; + ck_tile::index_t StrideB1 = StrideB; + + ck_tile::index_t StrideD0 = StrideD; + ck_tile::index_t StrideD1 = StrideD; + + const int n_warmup = arg_parser.get_int("warmup"); + const int n_repeat = arg_parser.get_int("repeat"); + const int k_batch = arg_parser.get_int("kbatch"); + + StrideA0 = get_default_stride(M, N, StrideA0, is_row_major(a1_layout)); + StrideA1 = get_default_stride(M, N, StrideA1, is_row_major(a1_layout)); + + StrideB0 = get_default_stride(K, N, StrideB0, is_row_major(b0_layout)); + StrideB1 = get_default_stride(K, N, StrideB1, is_row_major(b1_layout)); + + StrideD0 = get_default_stride(M, N, StrideD0, is_row_major(d0_layout)); + StrideD1 = get_default_stride(M, N, StrideD1, is_row_major(d1_layout)); + + StrideE = get_default_stride(M, N, StrideE, is_row_major(e_layout)); + + ck_tile::HostTensor a0_m_k_tesnor( + host_tensor_descriptor(M, K, StrideA0, is_row_major(a0_layout))); + ck_tile::HostTensor a1_m_k_tesnor( + host_tensor_descriptor(M, K, StrideA1, is_row_major(a1_layout))); + + ck_tile::HostTensor b0_k_n_tensors( + host_tensor_descriptor(K, N, StrideB0, is_row_major(b0_layout))); + ck_tile::HostTensor b1_k_n_tensors( + host_tensor_descriptor(K, N, StrideB1, is_row_major(b1_layout))); + + ck_tile::HostTensor d0_m_n_tensors( + host_tensor_descriptor(M, N, StrideD0, is_row_major(d0_layout))); + ck_tile::HostTensor d1_m_n_tensors( + host_tensor_descriptor(M, N, StrideD1, is_row_major(d1_layout))); + + ck_tile::HostTensor e_m_n_device_result( + host_tensor_descriptor(M, N, StrideE, is_row_major(e_layout))); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(a0_m_k_tesnor); + ck_tile::FillUniformDistribution{-1.f, 1.f}(a1_m_k_tesnor); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(b0_k_n_tensors); + ck_tile::FillUniformDistribution{-1.f, 1.f}(b1_k_n_tensors); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(d0_m_n_tensors); + ck_tile::FillUniformDistribution{-1.f, 1.f}(d1_m_n_tensors); + + ck_tile::DeviceMem a0_m_k_dev_buf(a0_m_k_tesnor.get_element_space_size_in_bytes()); + ck_tile::DeviceMem a1_m_k_dev_buf(a1_m_k_tesnor.get_element_space_size_in_bytes()); + + ck_tile::DeviceMem b0_k_n_dev_buf(b0_k_n_tensors.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b1_k_n_dev_buf(b1_k_n_tensors.get_element_space_size_in_bytes()); + + ck_tile::DeviceMem d0_m_n_dev_buf(d0_m_n_tensors.get_element_space_size_in_bytes()); + ck_tile::DeviceMem d1_m_n_dev_buf(d1_m_n_tensors.get_element_space_size_in_bytes()); + + ck_tile::DeviceMem e_m_n_dev_buf(e_m_n_device_result.get_element_space_size_in_bytes()); + + a0_m_k_dev_buf.ToDevice(a0_m_k_tesnor.mData.data()); + a1_m_k_dev_buf.ToDevice(a1_m_k_tesnor.mData.data()); + + b0_k_n_dev_buf.ToDevice(b0_k_n_tensors.mData.data()); + b1_k_n_dev_buf.ToDevice(b1_k_n_tensors.mData.data()); + + d0_m_n_dev_buf.ToDevice(d0_m_n_tensors.mData.data()); + d1_m_n_dev_buf.ToDevice(d1_m_n_tensors.mData.data()); + + e_m_n_dev_buf.SetZero(); + e_m_n_device_result.SetZero(); + + std::array as_ptr_buf = {a0_m_k_dev_buf.GetDeviceBuffer(), + a1_m_k_dev_buf.GetDeviceBuffer()}; + + std::array bs_ptr_buf = {b0_k_n_dev_buf.GetDeviceBuffer(), + b1_k_n_dev_buf.GetDeviceBuffer()}; + + std::array ds_ptr_buf = {d0_m_n_dev_buf.GetDeviceBuffer(), + d1_m_n_dev_buf.GetDeviceBuffer()}; + + std::array strideAs = {StrideA0, StrideA1}; + std::array strideBs = {StrideB0, StrideB1}; + std::array strideDs = {StrideD0, StrideD1}; + + invoke_gemm_multi_abd(as_ptr_buf, + bs_ptr_buf, + ds_ptr_buf, + e_m_n_dev_buf.GetDeviceBuffer(), + M, + N, + K, + strideAs, + strideBs, + strideDs, + StrideE, + n_warmup, + n_repeat, + k_batch); + + e_m_n_dev_buf.FromDevice(e_m_n_device_result.data()); + + ck_tile::HostTensor a_m_k_host_ref_element_result( + host_tensor_descriptor(M, K, StrideA0, is_row_major(a0_layout))); + ck_tile::HostTensor b_k_n_host_ref_element_result( + host_tensor_descriptor(K, N, StrideB0, is_row_major(b0_layout))); + ck_tile::HostTensor e_m_n_host_ref( + host_tensor_descriptor(M, N, StrideE, is_row_major(e_layout))); + a_m_k_host_ref_element_result.SetZero(); + b_k_n_host_ref_element_result.SetZero(); + e_m_n_host_ref.SetZero(); + + ck_tile::reference_gemm_multiple_abd({a0_m_k_tesnor, a1_m_k_tesnor}, + {b0_k_n_tensors, b1_k_n_tensors}, + {d0_m_n_tensors, d1_m_n_tensors}, + a_m_k_host_ref_element_result, + b_k_n_host_ref_element_result, + e_m_n_host_ref); + + bool pass{true}; + if(arg_parser.get_int("v")) + { + const float max_accumulated_value = + *std::max_element(e_m_n_host_ref.mData.begin(), e_m_n_host_ref.mData.end()); + + const auto rtol_atol = calculate_rtol_atol(K, 1, max_accumulated_value); + + pass &= ck_tile::check_err(e_m_n_device_result, + e_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + std::cout << "Relative error threshold: " << rtol_atol.at(ck_tile::number<0>{}) + << std::endl; + std::cout << "Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + std::cout << "The CPU veification result is: " << (pass ? "correct" : "fail") << std::endl; + } + return pass; +} + +template +int run_multiple_abd_gemm_example(int argc, char* argv[]) +{ + auto [result, arg_parser] = create_args(argc, argv); + if(!result) + { + return -1; + } + + const std::string as_layout = arg_parser.get_str("as_layout"); + const std::string bs_layout = arg_parser.get_str("bs_layout"); + + using Row = ck_tile::tensor_layout::gemm::RowMajor; + using Col = ck_tile::tensor_layout::gemm::ColumnMajor; + + if(as_layout == "R" && bs_layout == "C") + { + return run_gemm_multi_abd_example_with_layouts( + argc, argv, Row{}, Row{}, Col{}, Col{}, Row{}, Row{}, Row{}); + } + else + { + throw std::runtime_error("Unsupported data layout configuration for A,B and C tensors!"); + } +} diff --git a/example/ck_tile/22_gemm_multi_abd/utils.hpp b/example/ck_tile/22_gemm_multi_abd/utils.hpp new file mode 100644 index 0000000000..38bf8623d4 --- /dev/null +++ b/example/ck_tile/22_gemm_multi_abd/utils.hpp @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +template +static constexpr inline auto is_row_major(Layout layout_) +{ + return ck_tile::bool_constant, + ck_tile::tensor_layout::gemm::RowMajor>>{}; +} + +auto calculate_rtol_atol(const ck_tile::index_t K, + const ck_tile::index_t kbatch, + const float max_accumulated_value) +{ + using ComputeTypeAB = + std::conditional_t; + + using ComputeType = + std::conditional_t; + // Calculate thresholds + const auto rtol = ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(K, kbatch)); + + const auto atol = ck_tile::get_absolute_threshold( + max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); + + // Calculate error due to split_k accumulation + const auto rtol_split_k = + ck_tile::get_relative_threshold(kbatch); + + const auto atol_split_k = ck_tile::get_absolute_threshold( + max_accumulated_value, kbatch); + + // Use higher threshold + return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); +} diff --git a/example/ck_tile/38_block_scale_gemm/README.md b/example/ck_tile/38_block_scale_gemm/README.md index 9acc4f9bfc..9b2610813c 100644 --- a/example/ck_tile/38_block_scale_gemm/README.md +++ b/example/ck_tile/38_block_scale_gemm/README.md @@ -5,6 +5,7 @@ This folder contains examples of quant GEMMs using the ck_tile tile-programming - AQuant kernel with blocks of A matrix sharing scales: custom GEMM pipeline - BQuant kernel with blocks of B matrix sharing scales: custom GEMM pipeline - Row and Column-wise scaled: scaling implemented in Epilogue +- Tensor-wise scaled: scaling implemented in Epilogue ## build ``` @@ -14,7 +15,6 @@ mkdir build && cd build ../script/cmake-ck-dev.sh ../ # Compile the quant kernels make tile_example_gemm_quant_basic -j -make tile_example_gemm_bquant_basic -j ``` This will result in an executable `build/bin/tile_example_gemm_quant_basic` @@ -37,7 +37,7 @@ args: -warmup number of iterations before benchmark the kernel (default:10) -repeat number of iterations to benchmark the kernel (default:100) -timer gpu:gpu timer, cpu:cpu timer (default:gpu) - -quant_mode Which quant method to use (aquant, rowcol) + -quant_mode Which quant method to use (aquant, bquant, tensor, rowcol) ``` User need to select correct mapping of config for each quant mode: diff --git a/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp b/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp index 79c6cca6cb..91f799f194 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_quant_basic.cpp @@ -66,19 +66,21 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str constexpr auto tail_number_v = tail_number_.value; constexpr bool transpose_c = false; + // row-col and tensor quants use the regular pipeline, A/B quants use their own using PipelineProblem = std::conditional_t< - QuantMode == ck_tile::QuantType::RowColQuant, - ck_tile::GemmRowColQuantPipelineProblem, + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant, + ck_tile::GemmRowColTensorQuantPipelineProblem, std::conditional_t>>; using GemmPipeline = std::conditional_t< - QuantMode == ck_tile::QuantType::RowColQuant, + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant, ck_tile::GemmPipelineAgBgCrCompV3, std::conditional_t, @@ -241,10 +244,18 @@ int run_gemm_example(int argc, char* argv[]) ck_tile::QuantType::RowColQuant>( a_layout, b_layout, argc, argv); } + else if(quant_mode == "tensor") + { + return run_gemm_example_prec_type, + TypeConfig, + 128, + ck_tile::QuantType::TensorQuant>( + a_layout, b_layout, argc, argv); + } else { throw std::runtime_error( - "Unsupported quantization mode! Use 'aquant', 'bquant' or 'rowcol'"); + "Unsupported quantization mode! Use 'aquant', 'bquant', 'tensor' or 'rowcol'"); } } else if(data_type == "bf8") @@ -276,10 +287,18 @@ int run_gemm_example(int argc, char* argv[]) ck_tile::QuantType::RowColQuant>( a_layout, b_layout, argc, argv); } + else if(quant_mode == "tensor") + { + return run_gemm_example_prec_type, + TypeConfig, + 128, + ck_tile::QuantType::TensorQuant>( + a_layout, b_layout, argc, argv); + } else { throw std::runtime_error( - "Unsupported quantization mode! Use 'aquant', 'bquant' or 'rowcol'"); + "Unsupported quantization mode! Use 'aquant', 'bquant', 'tensor' or 'rowcol'"); } } else if(data_type == "i4fp8") diff --git a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp index ccf07460fa..e5313d8aaf 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp @@ -9,7 +9,7 @@ #include "ck_tile/host/kernel_launch.hpp" #include "ck_tile/ops/epilogue.hpp" #include "ck_tile/ops/gemm.hpp" -#include "ck_tile/ops/gemm_group_quant.hpp" +#include "ck_tile/ops/gemm_quant.hpp" template constexpr ck_tile::index_t get_k_warp_tile() @@ -241,7 +241,7 @@ auto create_args(int argc, char* argv[]) .insert("init", "0", "0:random, 1:linear, 2:constant(1)") .insert("flush_cache", "true", "flush cache before running the kernel, defaults to true") .insert("rotating_count", "1", "rotating count, defaults to 1") - .insert("quant_mode", "aquant", "Choose aquant (default), bquant or rowcol"); + .insert("quant_mode", "aquant", "Choose aquant (default), bquant, tensor or rowcol"); bool result = arg_parser.parse(argc, argv); return std::make_tuple(result, arg_parser); diff --git a/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc b/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc index 0f45811ff3..8e9456e973 100644 --- a/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc +++ b/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc @@ -119,11 +119,7 @@ float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, } std::cout << " Acc_Type = " << DataTypeTraits::name << " C_Type = " << DataTypeTraits::name - << " QuantMode = " - << (QuantMode == ck_tile::QuantType::AQuantGrouped - ? "AQuantGrouped" - : (QuantMode == ck_tile::QuantType::BQuantGrouped ? "BQuantGrouped" - : "RowColQuant")) + << " QuantMode = " << quant_type_to_string(QuantMode) << " PreshuffleQuant = " << (GemmConfig::PreshuffleQuant ? "true" : "false") << " : " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " << std::endl; @@ -183,10 +179,11 @@ int run_gemm_example_with_layouts(int argc, AQK = 0; // No A quantization BQK = K / QuantGroupSize; // Group quantization: BQK = K / GroupSize } - else if constexpr(QuantMode == ck_tile::QuantType::RowColQuant) + else if constexpr(QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant) { - AQK = 1; // Row quantization: tensor shape [M, 1] - BQK = N; // Column quantization: tensor shape [1, N] + AQK = 1; // Row quantization: tensor shape [M, 1] or [1] + BQK = 1; // Column quantization: tensor shape [1, N] or [1] } else { @@ -227,6 +224,11 @@ int run_gemm_example_with_layouts(int argc, stride_AQ = ck_tile::get_default_stride(M, 1, stride_AQ, is_row_major(aq_layout)); stride_BQ = ck_tile::get_default_stride(1, N, stride_BQ, is_row_major(bq_layout)); } + else if constexpr(QuantMode == ck_tile::QuantType::TensorQuant) + { + stride_AQ = 1; // Tensor quantization: tensor shape [1] + stride_BQ = 1; // Tensor quantization: tensor shape [1] + } ck_tile::HostTensor a_m_k( ck_tile::host_tensor_descriptor(M, K, stride_A, is_row_major(a_layout))); @@ -237,28 +239,30 @@ int run_gemm_example_with_layouts(int argc, // Create AQ tensor with appropriate shape std::unique_ptr> aq_tensor_ptr = nullptr; - if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped) + if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped || + QuantMode == ck_tile::QuantType::RowColQuant) { aq_tensor_ptr = std::make_unique>( ck_tile::host_tensor_descriptor(M, AQK, stride_AQ, is_row_major(aq_layout))); } - else if(QuantMode == ck_tile::QuantType::RowColQuant) + else if constexpr(QuantMode == ck_tile::QuantType::TensorQuant) { aq_tensor_ptr = std::make_unique>( - ck_tile::host_tensor_descriptor(M, AQK, stride_AQ, is_row_major(aq_layout))); + ck_tile::host_tensor_descriptor(1, 1, stride_AQ, is_row_major(aq_layout))); } - // Create BQ tensor only for RowColQuant mode + // Create BQ tensor with appropriate shape std::unique_ptr> bq_tensor_ptr = nullptr; - if constexpr(QuantMode == ck_tile::QuantType::BQuantGrouped) + if constexpr(QuantMode == ck_tile::QuantType::BQuantGrouped || + QuantMode == ck_tile::QuantType::RowColQuant) { bq_tensor_ptr = std::make_unique>( ck_tile::host_tensor_descriptor(BQK, N, stride_BQ, is_row_major(bq_layout))); } - else if constexpr(QuantMode == ck_tile::QuantType::RowColQuant) + else if constexpr(QuantMode == ck_tile::QuantType::TensorQuant) { bq_tensor_ptr = std::make_unique>( - ck_tile::host_tensor_descriptor(1, N, stride_BQ, is_row_major(bq_layout))); + ck_tile::host_tensor_descriptor(1, 1, stride_BQ, is_row_major(bq_layout))); } std::random_device rd; @@ -282,7 +286,7 @@ int run_gemm_example_with_layouts(int argc, *bq_tensor_ptr); ck_tile::FillUniformDistribution{-5.0f, 5.0f, fill_seed(gen)}(a_m_k); } - else + else if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped) { if constexpr(std::is_same_v) { @@ -296,12 +300,15 @@ int run_gemm_example_with_layouts(int argc, ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}( *aq_tensor_ptr); ck_tile::FillUniformDistribution{-5.0f, 5.0f, fill_seed(gen)}(b_k_n); - - if constexpr(QuantMode == ck_tile::QuantType::RowColQuant) - { - ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}( - *bq_tensor_ptr); - } + } + else + { + ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}(a_m_k); + ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}(b_k_n); + ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}( + *aq_tensor_ptr); + ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}( + *bq_tensor_ptr); } } else if(init_method == 1) @@ -343,7 +350,8 @@ int run_gemm_example_with_layouts(int argc, std::unique_ptr aq_dev_buf_ptr = nullptr; if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant) { aq_dev_buf_ptr = std::make_unique(aq_tensor_ptr->get_element_space_size_in_bytes()); @@ -351,14 +359,16 @@ int run_gemm_example_with_layouts(int argc, std::unique_ptr bq_dev_buf_ptr = nullptr; if constexpr(QuantMode == ck_tile::QuantType::BQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant) { bq_dev_buf_ptr = std::make_unique(bq_tensor_ptr->get_element_space_size_in_bytes()); } if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant) { if constexpr(GemmConfig::PreshuffleQuant) { @@ -398,7 +408,8 @@ int run_gemm_example_with_layouts(int argc, c_m_n_dev_result.SetZero(); if constexpr(QuantMode == ck_tile::QuantType::BQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant) { bq_dev_buf_ptr->ToDevice(bq_tensor_ptr->data()); } @@ -412,15 +423,9 @@ int run_gemm_example_with_layouts(int argc, CLayout, QuantGroupSize, QuantMode>(a_m_k_dev_buf, - (QuantMode == ck_tile::QuantType::AQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) - ? aq_dev_buf_ptr.get() - : nullptr, + aq_dev_buf_ptr.get(), b_k_n_dev_buf, - (QuantMode == ck_tile::QuantType::BQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) - ? bq_dev_buf_ptr.get() - : nullptr, + bq_dev_buf_ptr.get(), c_m_n_dev_buf, M, N, @@ -467,7 +472,7 @@ int run_gemm_example_with_layouts(int argc, QuantGroupSize, false>(a_m_k, *bq_tensor_ptr, b_k_n, c_m_n_host_ref); } - else + else if constexpr(QuantMode == ck_tile::QuantType::RowColQuant) { ck_tile::reference_gemm_rowcol_quant( a_m_k, *aq_tensor_ptr, b_k_n, *bq_tensor_ptr, c_m_n_host_ref); } + else if constexpr(QuantMode == ck_tile::QuantType::TensorQuant) + { + ck_tile::reference_gemm_tensor_quant( + a_m_k, *aq_tensor_ptr, b_k_n, *bq_tensor_ptr, c_m_n_host_ref); + } const float max_accumulated_value = *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); diff --git a/example/ck_tile/CMakeLists.txt b/example/ck_tile/CMakeLists.txt index 8fce70ba04..75d32a5eb0 100644 --- a/example/ck_tile/CMakeLists.txt +++ b/example/ck_tile/CMakeLists.txt @@ -21,6 +21,7 @@ add_subdirectory(18_flatmm) add_subdirectory(19_gemm_multi_d) add_subdirectory(20_grouped_convolution) add_subdirectory(21_elementwise) +add_subdirectory(22_gemm_multi_abd) add_subdirectory(35_batched_transpose) add_subdirectory(38_block_scale_gemm) add_subdirectory(39_copy) diff --git a/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp b/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp index 7296e4faaa..18223c78f7 100644 --- a/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp +++ b/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp @@ -11,6 +11,8 @@ namespace ck { namespace tensor_operation { namespace device { +#define DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS 1 + template ()) diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp index 934dc7ee8e..987a1e273a 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp @@ -671,6 +671,7 @@ struct DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle end(a_g_n_k_wos_lengths), begin(output_spatial_lengths_)); +#if !DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS if(split_k < 0) { ck::index_t gemmM, gemmN; @@ -683,6 +684,7 @@ struct DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle grid_size); } else +#endif { k_batch_ = split_k; } @@ -939,6 +941,12 @@ struct DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle static bool IsSupportedArgument(const Argument& arg) { +#if DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS + if(arg.k_batch_ < 0) + { + return false; + } +#endif if(!ck::is_xdl_wmma_supported()) { return false; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp index b361409e38..22fc13bae4 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp @@ -553,6 +553,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle conv_ngchw_to_nhwgc_transformer.TransposeWeiStrides(e_g_k_c_xs_lengths, e_g_k_c_xs_strides); +#if !DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS if(split_k < 0) { ck::index_t gemmM, gemmN; @@ -565,6 +566,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle grid_size); } else +#endif { k_batch_ = split_k; } @@ -934,6 +936,12 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle static bool IsSupportedArgument(const Argument& arg) { +#if DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS + if(arg.k_batch_ < 0) + { + return false; + } +#endif if(!ck::is_xdl_wmma_supported()) { return false; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp index 8bf188be2e..735eebbdf6 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp @@ -524,6 +524,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 end(a_g_n_k_wos_lengths), begin(output_spatial_lengths_)); +#if !DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS if(split_k < 0) { ck::index_t gemmM, gemmN, gemmK; @@ -549,6 +550,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 } } else +#endif { k_batch_ = split_k; } @@ -1275,6 +1277,13 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 static bool IsSupportedArgument(const Argument& arg) { +#if DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS + if(arg.k_batch_ < 0) + { + return false; + } +#endif + const index_t GemmM = arg.a_grid_desc_kbatch_k0_m_k1_.GetLength(I1); const index_t GemmN = arg.b_grid_desc_kbatch_k0_n_k1_.GetLength(I1); const index_t GemmK = arg.a_grid_desc_kbatch_k0_m_k1_.GetLength(I0) * diff --git a/include/ck/utility/amd_ck_fp8.hpp b/include/ck/utility/amd_ck_fp8.hpp index 2edbb7c789..0b73f76155 100644 --- a/include/ck/utility/amd_ck_fp8.hpp +++ b/include/ck/utility/amd_ck_fp8.hpp @@ -33,8 +33,34 @@ namespace ck { -using f8_fnuz_t = _BitInt(8); -using bf8_fnuz_t = unsigned _BitInt(8); +struct f8_fnuz_t +{ + using data_type = unsigned char; + data_type m_data; + __host__ __device__ explicit constexpr f8_fnuz_t(data_type in_data) : m_data(in_data) {} + __host__ __device__ explicit constexpr f8_fnuz_t() = default; + __host__ __device__ bool constexpr operator==(f8_fnuz_t other) const + { + return m_data == other.m_data; + } + __host__ __device__ explicit constexpr operator data_type() const { return m_data; } +}; + +struct bf8_fnuz_t +{ + using data_type = unsigned char; + data_type m_data; + __host__ __device__ explicit constexpr bf8_fnuz_t(data_type in_data) : m_data(in_data) {} + __host__ __device__ explicit constexpr bf8_fnuz_t() = default; + __host__ __device__ bool constexpr operator==(bf8_fnuz_t other) const + { + return m_data == other.m_data; + } + __host__ __device__ explicit constexpr operator data_type() const { return m_data; } +}; + +static_assert(1 == sizeof(f8_fnuz_t)); +static_assert(1 == sizeof(bf8_fnuz_t)); typedef unsigned char fp8_storage_t; diff --git a/include/ck/utility/data_type.hpp b/include/ck/utility/data_type.hpp index 48b352986e..984bb4d862 100644 --- a/include/ck/utility/data_type.hpp +++ b/include/ck/utility/data_type.hpp @@ -205,7 +205,7 @@ inline constexpr bool is_native_type() return is_same::value || is_same::value || is_same::value || is_same::value || is_same::value || is_same::value || is_same::value || is_same::value || - is_same::value || is_same::value || is_same::value; + is_same_v || is_same_v || is_same::value; } // scalar_type @@ -300,14 +300,14 @@ struct scalar_type template <> struct scalar_type { - using type = f8_fnuz_t; + using type = f8_fnuz_t::data_type; static constexpr index_t vector_size = 1; }; template <> struct scalar_type { - using type = bf8_fnuz_t; + using type = bf8_fnuz_t::data_type; static constexpr index_t vector_size = 1; }; diff --git a/include/ck/utility/dtype_vector.hpp b/include/ck/utility/dtype_vector.hpp index ae0edb35ee..27a7545a0e 100644 --- a/include/ck/utility/dtype_vector.hpp +++ b/include/ck/utility/dtype_vector.hpp @@ -1294,6 +1294,18 @@ struct nnvb_data_t_selector using type = bf8_ocp_t::data_type; }; +template <> +struct nnvb_data_t_selector +{ + using type = f8_fnuz_t::data_type; +}; + +template <> +struct nnvb_data_t_selector +{ + using type = bf8_fnuz_t::data_type; +}; + template <> struct nnvb_data_t_selector { diff --git a/include/ck/utility/f8_utils.hpp b/include/ck/utility/f8_utils.hpp index 799683ae65..748aa07f9e 100644 --- a/include/ck/utility/f8_utils.hpp +++ b/include/ck/utility/f8_utils.hpp @@ -39,7 +39,7 @@ __host__ __device__ Y run_cast_to_f8(X x, uint32_t rng) int exponent, bias; uint32_t head, mantissa, sign; // nan code is same for float and half - constexpr Y nan_code = 0x80; + constexpr uint8_t nan_code = 0x80; constexpr uint32_t nan_mask = NumericUtils::nan_mask; // convert to bitwise @@ -60,17 +60,17 @@ __host__ __device__ Y run_cast_to_f8(X x, uint32_t rng) if constexpr(negative_zero_nan) { if((x_bitwise & nan_mask) == nan_mask) - return nan_code; + return Y{nan_code}; } else { if((x_bitwise & nan_mask) == nan_mask) - return signed_inf + (mantissa != 0 ? 1 : 0); + return Y{static_cast(signed_inf + (mantissa != 0 ? 1 : 0))}; } // check if x is 0.0 if(x_bitwise == 0) - return 0; + return Y{0}; // First need to check if it is normal or denorm as there is a difference of implict 1 // Then need to adjust the exponent to align with the F8 exponent, in the meanwhile, shift @@ -178,9 +178,10 @@ In this case, the fp16 mantissa should be shift left by 1 */ // check if x is 0.0 or -0.0 if(out_exponent == 0 && mantissa == 0) - return negative_zero_nan ? 0 : (sign << (out_exp + out_mant)); + return Y{negative_zero_nan ? 0 : static_cast(sign << (out_exp + out_mant))}; mantissa &= (1 << out_mant) - 1; - return (sign << (out_exp + out_mant)) | (out_exponent << out_mant) | mantissa; + return Y{static_cast((sign << (out_exp + out_mant)) | (out_exponent << out_mant) | + mantissa)}; } template @@ -195,8 +196,8 @@ __host__ __device__ Y run_cast_from_f8(X x) constexpr int out_mant = NumericUtils::mant; // prepare the codes - constexpr X nan_code = 0x80; - using T_bitwise = typename NumericUtils::bitwise_type; + constexpr uint8_t nan_code = 0x80; + using T_bitwise = typename NumericUtils::bitwise_type; constexpr T_bitwise Inf_bitwise = NumericUtils::Inf; constexpr T_bitwise NegInf_bitwise = NumericUtils::NegInf; @@ -209,13 +210,13 @@ __host__ __device__ Y run_cast_from_f8(X x) constexpr Y Neg0 = bit_cast(Neg0_bitwise); // check if x is 0.0 - if(x == 0) + if(!static_cast(x)) return static_cast(0); // unpack the input - uint32_t sign = x >> (in_exp + in_mant); - uint32_t mantissa = x & ((1 << in_mant) - 1); - int exponent = (x & 0x7F) >> in_mant; + uint32_t sign = static_cast(x) >> (in_exp + in_mant); + uint32_t mantissa = static_cast(x) & ((1 << in_mant) - 1); + int exponent = (static_cast(x) & 0x7F) >> in_mant; constexpr int exp_low_cutoff = (1 << (out_exp - 1)) - (1 << (in_exp - 1)) + 1 - (negative_zero_nan ? 1 : 0); @@ -223,12 +224,12 @@ __host__ __device__ Y run_cast_from_f8(X x) if constexpr(negative_zero_nan) { - if(x == nan_code) + if(static_cast(x) == nan_code) return NaN; } else { - if(x == nan_code) + if(static_cast(x) == nan_code) return Neg0; if(exponent == ((1 << in_exp) - 1)) return (mantissa == 0) ? (sign ? NegInf : Inf) : NaN; diff --git a/include/ck/utility/type_convert.hpp b/include/ck/utility/type_convert.hpp index 290a6c8dd6..66d760c2b3 100644 --- a/include/ck/utility/type_convert.hpp +++ b/include/ck/utility/type_convert.hpp @@ -351,7 +351,7 @@ inline __host__ __device__ f8_fnuz_t f8_convert_sr(float x) val.fval = __builtin_amdgcn_fmed3f(val.fval, max_fp8, -max_fp8); ival = __builtin_amdgcn_cvt_sr_fp8_f32(val.fval, rng, ival, 0); // 0 pos val.i32val = ival; - return val.i8val[0]; // little endian + return f8_fnuz_t{val.i8val[0]}; // little endian #else constexpr bool negative_zero_nan = true; constexpr bool clip = true; @@ -419,7 +419,7 @@ inline __host__ __device__ bf8_fnuz_t f8_convert_sr(float x) val.fval = __builtin_amdgcn_fmed3f(val.fval, max_bf8, -max_bf8); ival = __builtin_amdgcn_cvt_sr_bf8_f32(val.fval, rng, ival, 0); // 0 pos val.i32val = ival; - return val.i8val[0]; // little endian + return bf8_fnuz_t{val.i8val[0]}; // little endian #else constexpr bool negative_zero_nan = true; constexpr bool clip = true; @@ -655,7 +655,7 @@ inline __host__ __device__ f8_fnuz_t f8_convert_rne(float x) val.fval = __builtin_amdgcn_fmed3f(val.fval, max_fp8, -max_fp8); ival = __builtin_amdgcn_cvt_pk_fp8_f32(val.fval, val.fval, ival, false); // false -> WORD0 val.i32val = ival; - return val.i8val[0]; + return f8_fnuz_t{val.i8val[0]}; #else constexpr bool negative_zero_nan = true; constexpr bool clip = true; @@ -707,7 +707,7 @@ inline __host__ __device__ bf8_fnuz_t f8_convert_rne(float x) val.fval = __builtin_amdgcn_fmed3f(val.fval, max_bf8, -max_bf8); ival = __builtin_amdgcn_cvt_pk_bf8_f32(val.fval, val.fval, ival, false); // false -> WORD0 val.i32val = ival; - return val.i8val[0]; + return bf8_fnuz_t{val.i8val[0]}; #else constexpr bool negative_zero_nan = true; constexpr bool clip = true; @@ -924,7 +924,7 @@ inline __host__ __device__ float type_convert(f8_fnuz_t x) { #if defined(__gfx94__) float fval; - uint32_t i32val = static_cast(x); + uint32_t i32val = static_cast(static_cast(x)); fval = __builtin_amdgcn_cvt_f32_fp8(i32val, 0); // asm volatile("v_cvt_f32_fp8 %0, %1 src0_sel:BYTE_0" : "=v"(fval) : "v"(i32val)); return fval; @@ -1430,7 +1430,7 @@ inline __host__ __device__ float type_convert(bf8_fnuz_t x) { #if defined(__gfx94__) float fval; - uint32_t i32val = static_cast(x); + uint32_t i32val = static_cast(static_cast(x)); fval = __builtin_amdgcn_cvt_f32_bf8(i32val, 0); // asm volatile("v_cvt_f32_bf8 %0, %1 src0_sel:BYTE_0" : "=v"(fval) : "v"(i32val)); return fval; diff --git a/include/ck_tile/core/tensor/load_tile.hpp b/include/ck_tile/core/tensor/load_tile.hpp index 8b7541bf23..a3620453b4 100644 --- a/include/ck_tile/core/tensor/load_tile.hpp +++ b/include/ck_tile/core/tensor/load_tile.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -26,6 +26,29 @@ CK_TILE_DEVICE auto load_tile(const TileWindow_& tile_window, return tile_window.load(number{}, bool_constant{}); } +/** + * @brief Load tile with elementwise function + * + * @note This function is a modification of the existing load function. + * It has been extended with two additional parameters: it takes a tuple as input + * and an elementwise function. For each A = A0, A1… AN, the elementwise function + * is additionally applied during a single read. + */ +template +CK_TILE_DEVICE auto load_tile_with_elementwise(const TileWindow_& tile_window, + ElementWise_ elementwise, + number = {}, + bool_constant = {}) +{ + // TODO: Tile windows should works with unknow number of params + // Load element_wise API works only when the input typle is a tuple-tyupe + return tile_window[number<0>{}].load( + tile_window, elementwise, number{}, bool_constant{}); +} + template +concept IsLoadableTile = requires { load_tile(std::declval()); }; + } // namespace ck_tile diff --git a/include/ck_tile/core/tensor/tensor_view.hpp b/include/ck_tile/core/tensor/tensor_view.hpp index 6fa8f898e5..fb209ba827 100644 --- a/include/ck_tile/core/tensor/tensor_view.hpp +++ b/include/ck_tile/core/tensor/tensor_view.hpp @@ -455,7 +455,7 @@ CK_TILE_HOST_DEVICE constexpr auto make_tensor_view(DataType* __restrict__ p, auto buffer_view = make_buffer_view(p, desc.get_element_space_size()); - return tensor_view{buffer_view, desc}; + return tensor_view{buffer_view, desc}; } template + CK_TILE_DEVICE auto load(const TileWindow_& tile_window, + ElementWise_ elementwise, + number = {}, + bool_constant = {}) const + { + constexpr auto tile_dstr = typename Base::TileDstr{}; + auto dst_tensor = make_static_distributed_tensor(tile_dstr); + load(dst_tensor, + tile_window, + elementwise, + number{}, + bool_constant{}); + return dst_tensor; + } + + template + CK_TILE_DEVICE auto load(DistributedTensor& dst_tensor, + const TileWindow_& tile_window, + ElementWise_ elementwise, + number = {}, + bool_constant = {}) const + { + + using Traits = typename Base::Traits; + using vector_t = typename Traits::vector_t; + using SFC_Ys = typename Traits::SFC_Ys; + + constexpr auto tile_dstr = typename Base::TileDstr{}; + constexpr auto sizeOfTuple = TileWindow_::size(); + // loop over thread tensor space [y0, y1, ...] + static_for<0, NumCoord, 1>{}([&](auto iCoord) { + /// TODO: use structure binding (to be captured later) if compiled in C++20 + auto window_adaptor_thread_coord = + tile_window[number<0>{}].pre_computed_coords_[iCoord][I0]; + auto bottom_tensor_thread_coord = + tile_window[number<0>{}].pre_computed_coords_[iCoord][I1]; + + static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) { + constexpr auto iAccess = number{}; + + // data index [y0, y1, ...] + constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); + + // read from bottom tensor + const auto idx_vec_value = generate_tuple( + [&](auto jj) { + return tile_window[number{}] + .get_bottom_tensor_view() + .template get_vectorized_elements( + bottom_tensor_thread_coord, + 0, + bool_constant{}); + }, + number{}); + + // write into distributed tensor + static_for<0, Traits::ScalarPerVector, Traits::PackedSize>{}([&](auto j) { + constexpr auto idx_ys = generate_tuple( + [&](auto jj) { + return jj == Traits::VectorDimY ? (idx_ys_start[jj] + j) + : idx_ys_start[jj]; + }, + number{}); + + constexpr index_t d = + tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys) / + Traits::PackedSize; + + ck_tile::apply( + [&](auto&&... t) { + elementwise(dst_tensor.get_thread_buffer().template at(), + t.template get_as< + typename Base::DataType>()[j / Traits::PackedSize]...); + }, + idx_vec_value); + }); + // move thread coordinate + if constexpr(iCoordAccess != (NumAccessPerCoord - 1)) + { + constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); + + constexpr auto idx_diff_ps_ys = container_concat( + generate_tuple([&](auto) { return number<0>{}; }, number{}), + idx_diff_ys); + + Base::move_window_adaptor_and_bottom_tensor_thread_coordinate( + window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys); + } + }); + }); + } + template @@ -857,6 +967,39 @@ CK_TILE_DEVICE void move_tile_window( window.move(step); } +template +CK_TILE_DEVICE void move_tile_window( + tuple>& window, + const typename tile_window_with_static_distribution::BottomTensorIndex& step) +{ + using T = tuple>; + + static constexpr auto N = T::size(); + static_for<0, N, 1>{}([&](auto Is) { window[number{}].move(step); }); +} + +template ::value>* = nullptr> +CK_TILE_DEVICE void move_tile_window(TileWindowWithStaticDistributionType& window, StepType& step) +{ + static constexpr auto N = TileWindowWithStaticDistributionType::size(); + static_for<0, N, 1>{}([&](auto Is) { window[number{}].move(step); }); +} + /** * @brief This class provides description of tile windowed view on the device memory. * diff --git a/include/ck_tile/host/reference/reference_gemm.hpp b/include/ck_tile/host/reference/reference_gemm.hpp index caa00e5994..90f68f7e2e 100644 --- a/include/ck_tile/host/reference/reference_gemm.hpp +++ b/include/ck_tile/host/reference/reference_gemm.hpp @@ -180,10 +180,6 @@ CK_TILE_HOST void reference_gemm_rowcol_quant(const HostTensor& a_m_k else v_b = fp32_val.lo; } - else if constexpr(std::is_same_v) - { - v_b = fp8_to_float_raw(b_element_op(b_k_n(k, n))); - } else { v_b = ck_tile::type_convert(b_element_op(b_k_n(k, n))); @@ -198,7 +194,57 @@ CK_TILE_HOST void reference_gemm_rowcol_quant(const HostTensor& a_m_k }; make_ParallelTensorFunctor(f_mn, M, N)(std::thread::hardware_concurrency()); - std::cout << std::endl; +} + +template +CK_TILE_HOST void reference_gemm_tensor_quant(const HostTensor& a_m_k, + const HostTensor& aq_1_1, + const HostTensor& b_k_n, + const HostTensor& bq_1_1, + HostTensor& c_m_n, + const AElementOp& a_element_op = {}, + const BElementOp& b_element_op = {}, + const ACCElementOp& acc_element_op = {}) +{ + static_assert(std::is_same_v || std::is_same_v); + static_assert(std::is_same_v || std::is_same_v); + static_assert(std::is_same_v); + static_assert(std::is_same_v || std::is_same_v); + static_assert(std::is_same_v && std::is_same_v); + const std::size_t M = a_m_k.get_length(0); + const std::size_t N = b_k_n.get_length(1); + const std::size_t K = a_m_k.get_length(1); + + auto f_mn = [&](auto m, auto n) { + // Init accumulator + AccDataType v_acc = 0; + // Get scale for A and scale for B + const AccDataType a_scale = ck_tile::type_convert(aq_1_1(0, 0)); + const AccDataType b_scale = ck_tile::type_convert(bq_1_1(0, 0)); + + // Compute the dot product + for(std::size_t k = 0; k < K; ++k) + { + AccDataType v_a = ck_tile::type_convert(a_element_op(a_m_k(m, k))); + AccDataType v_b = ck_tile::type_convert(b_element_op(b_k_n(k, n))); + + v_acc += v_a * v_b; + } + + v_acc = v_acc * a_scale * b_scale; + + c_m_n(m, n) = ck_tile::type_convert(acc_element_op(v_acc)); + }; + + make_ParallelTensorFunctor(f_mn, M, N)(std::thread::hardware_concurrency()); } template & a_m_k, make_ParallelTensorFunctor(f_mn, M, N)(std::thread::hardware_concurrency()); } +template >, + typename BDataType = remove_cvref_t>, + typename DDataType = remove_cvref_t>> +CK_TILE_HOST void +reference_gemm_multiple_abd(const std::array, AsDataType::size()>& as_m_k, + const std::array, BsDataType::size()>& bs_k_n, + const std::array, DsDataType::size()>& ds_m_n, + HostTensor& a_m_k, + HostTensor& b_k_n, + HostTensor& c_m_n, + const AElementOp& a_element_op = {}, + const BElementOp& b_element_op = {}, + const CDElementOp& acc_element_op = {}) +{ + const std::size_t M = a_m_k.get_length(0); + const std::size_t N = b_k_n.get_length(1); + const std::size_t K = a_m_k.get_length(1); + + auto as_m_k_tuple = + generate_tie([&](auto idx) -> auto& { return as_m_k[idx]; }, number{}); + + auto bs_k_n_tuple = + generate_tie([&](auto idx) -> auto& { return bs_k_n[idx]; }, number{}); + + auto ds_m_n_tuple = + generate_tie([&](auto idx) -> auto& { return ds_m_n[idx]; }, number{}); + + // Apply elementwise function to A + auto a_elementwise_fn = [&](auto i, auto j) { + ck_tile::apply([&](auto&&... t) { a_element_op(a_m_k(i, j), t(i, j)...); }, as_m_k_tuple); + }; + + make_ParallelTensorFunctor(a_elementwise_fn, M, K)(std::thread::hardware_concurrency()); + + // Apply elementwise function to B + auto b_elementwise_fn = [&](auto i, auto j) { + ck_tile::apply([&](auto&&... t) { b_element_op(b_k_n(i, j), t(i, j)...); }, bs_k_n_tuple); + }; + + make_ParallelTensorFunctor(b_elementwise_fn, K, N)(std::thread::hardware_concurrency()); + + auto f_mk_kn_mn = [&](auto m, auto n) { + AccDataType v_acc = 0; + for(std::size_t k = 0; k < K; ++k) + { + ADataType v_a = a_m_k(m, k); + BDataType v_b = b_k_n(k, n); + v_acc += + ck_tile::type_convert(v_a) * ck_tile::type_convert(v_b); + } + + CDataType v_c = 0; + + ck_tile::apply( + [&](auto&&... t) { + acc_element_op(v_c, + ck_tile::type_convert(v_acc), + ck_tile::type_convert(t(m, n))...); + }, + ds_m_n_tuple); + + c_m_n(m, n) = ck_tile::type_convert(v_c); + }; + + make_ParallelTensorFunctor(f_mk_kn_mn, M, N)(std::thread::hardware_concurrency()); +} + template +struct is_pk_int4 : std::false_type +{ +}; +template <> +struct is_pk_int4 : std::true_type +{ +}; + +template +struct InterleavedPKTypeLoader +{ + template + CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, + const WarpWindow& warp_window) + { + const element_wise::PassThroughPack8 elementwise_op{}; + + static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); + constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; + const auto in_dstr_tensors = load_tile(warp_window); + + using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); + static_for<0, thread_buffer_size, 1>{}([&](auto i) { + elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), + in_dstr_tensors.get_thread_buffer().template get_as()[i]); + }); + } +}; + +template +CK_TILE_DEVICE void load_int4_tile(WarpTile& dst, const WarpWindow& src) +{ + if constexpr(is_pk_int4>::value) + { + InterleavedPKTypeLoader::load_interleaved_pk_type(dst, src); + } + else + { + dst = load_tile(src); + } +} + +} // namespace ck_tile diff --git a/include/ck_tile/ops/elementwise/unary_element_wise_operation.hpp b/include/ck_tile/ops/elementwise/unary_element_wise_operation.hpp index 692d5ec504..221592ee10 100644 --- a/include/ck_tile/ops/elementwise/unary_element_wise_operation.hpp +++ b/include/ck_tile/ops/elementwise/unary_element_wise_operation.hpp @@ -392,6 +392,23 @@ struct PassThrough } }; +struct AddScale +{ + template + CK_TILE_HOST_DEVICE constexpr void operator()(E& a, const As&... as) const + { + // Start with the base value c + float result = ck_tile::type_convert(0.0f); + + // Add by each D parameter using fold expression + ((result += ck_tile::type_convert(as)), ...); + + a = ck_tile::type_convert(scale * result); + } + + float scale = 1.0; +}; + struct MultiDMultiply { template diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp index 628af0e0b3..6c815d804d 100644 --- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp @@ -28,8 +28,8 @@ struct GetDataType using type = typename T::DataType; // Use T::ScaleN::DataType }; -template struct CShuffleEpilogueProblem { - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; using AccDataType = remove_cvref_t; using ODataType = remove_cvref_t; using DsDataType = remove_cvref_t; @@ -83,12 +83,27 @@ template struct CShuffleEpilogue { using Problem = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; using AccDataType = remove_cvref_t; using ODataType = remove_cvref_t; using DsDataType = remove_cvref_t; using DsLayout = remove_cvref_t; + + static constexpr bool ADataTypeIsTuple = is_detected::value; + static constexpr bool BDataTypeIsTuple = is_detected::value; + + using AsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using BsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using ADataType = remove_cvref_t{}, AsDataTypeTuple>>; + using BDataType = remove_cvref_t{}, BsDataTypeTuple>>; + using ATypeToUse = std::conditional_t, BDataType, ADataType>; // Used for weight-only quantization kernel, B would be dequantized to the same data type as A @@ -289,22 +304,41 @@ struct CShuffleEpilogue CK_TILE_DEVICE void scale_tile(LdsTile& lds_tile, ScaleM& scale_m_window, ScaleN& scale_n_window) { - // Load tiles - const auto scale_m_tile = load_tile(scale_m_window); - const auto scale_n_tile = load_tile(scale_n_window); - - // Compute element-wise product in-place i.e. lds_tile = lds_tile * scale_m * scale_n - tile_elementwise_inout( - element_wise::MultiDMultiply{}, lds_tile, lds_tile, scale_m_tile, scale_n_tile); - - // Move scale windows - constexpr index_t num_access = SFC::get_num_of_access(); - if constexpr(iAccess != num_access - 1) + // Check if scales are EmptyScale first (no scaling needed) + if constexpr(std::is_same_v && std::is_same_v) { - constexpr auto step = SFC::get_forward_step(iAccess); + // No scaling needed - this is a no-op + } + // Check if scales are scalar AccDataType + else if constexpr(std::is_same_v && + std::is_same_v) + { + // Handle scalar scales + const AccDataType scale_m = scale_m_window; + const AccDataType scale_n = scale_n_window; + tile_elementwise_inout([&](auto& element) { element = element * scale_m * scale_n; }, + lds_tile); + } + // Otherwise, assume they are tile windows that can be loaded + else + { + // Load tiles + const auto scale_m_tile = load_tile(scale_m_window); + const auto scale_n_tile = load_tile(scale_n_window); - move_tile_window(scale_m_window, {step.at(number<0>{}), step.at(number<1>{})}); - move_tile_window(scale_n_window, {step.at(number<0>{}), step.at(number<1>{})}); + // Compute element-wise product in-place i.e. lds_tile = lds_tile * scale_m * scale_n + tile_elementwise_inout( + element_wise::MultiDMultiply{}, lds_tile, lds_tile, scale_m_tile, scale_n_tile); + + // Move scale windows + constexpr index_t num_access = SFC::get_num_of_access(); + if constexpr(iAccess != num_access - 1) + { + constexpr auto step = SFC::get_forward_step(iAccess); + + move_tile_window(scale_m_window, {step.at(number<0>{}), step.at(number<1>{})}); + move_tile_window(scale_n_window, {step.at(number<0>{}), step.at(number<1>{})}); + } } } @@ -437,6 +471,8 @@ struct CShuffleEpilogue // Optional scales (must share the same distribution to match per-thread indexing) constexpr bool has_scales = !std::is_same::value && !std::is_same::value; + constexpr bool has_scalar_scales = + std::is_same_v && std::is_same_v; // Tiles to hold row/col scales when present using SMType = typename GetDataType>::type; @@ -447,8 +483,11 @@ struct CShuffleEpilogue // Build windows only if scales are provided auto scale_m_window = [&]() { - if constexpr(has_scales) + if constexpr(has_scales && !has_scalar_scales) { + static_assert( + IsLoadableTile, + "ScaleM must be a loadable tile"); return make_tile_window(scale_m, dram_tile_distribution); } else @@ -457,8 +496,11 @@ struct CShuffleEpilogue } }(); auto scale_n_window = [&]() { - if constexpr(has_scales) + if constexpr(has_scales && !has_scalar_scales) { + static_assert( + IsLoadableTile, + "ScaleN must be a loadable tile"); return make_tile_window(scale_n, dram_tile_distribution); } else @@ -474,7 +516,7 @@ struct CShuffleEpilogue merge_sequences(sequence<1, NRepeat>{}, c_warp_y_lengths)); // If scales provided, load them with identical distribution - if constexpr(has_scales) + if constexpr(has_scales && IsLoadableTile && IsLoadableTile) { sm_tile = load_tile(scale_m_window); // row scales in permuted layout sn_tile = load_tile(scale_n_window); // col scales in permuted layout @@ -489,7 +531,11 @@ struct CShuffleEpilogue auto emit = [&](index_t out_idx, index_t src_row) { AccDataType v = shuffle_acc.get_thread_buffer()[base + src_row]; - if constexpr(has_scales) + if constexpr(has_scalar_scales) + { + v = static_cast(v * scale_m * scale_n); + } + else if constexpr(has_scales) { // same linear index mapping on the permuted distribution const auto s_m = static_cast(sm_tile.get_thread_buffer()[out_idx]); @@ -580,10 +626,19 @@ struct CShuffleEpilogue number{}); constexpr bool has_scales = - !std::is_same::value && !std::is_same::value; + !std::is_same_v && !std::is_same_v; + constexpr bool has_scalar_scales = + std::is_same_v && std::is_same_v; auto scale_m_window = [&]() { - if constexpr(has_scales) + if constexpr(has_scalar_scales) { + return scale_m; + } + else if constexpr(has_scales) + { + static_assert( + IsLoadableTile, + "ScaleM must be a loadable tile"); return make_tile_window(scale_m, lds_tile.get_tile_distribution()); } else @@ -592,8 +647,15 @@ struct CShuffleEpilogue } }(); auto scale_n_window = [&]() { - if constexpr(has_scales) + if constexpr(has_scalar_scales) { + return scale_n; + } + else if constexpr(has_scales) + { + static_assert( + IsLoadableTile, + "ScaleN must be a loadable tile"); return make_tile_window(scale_n, lds_tile.get_tile_distribution()); } else diff --git a/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp b/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp index 54becd3c0f..2843966cd7 100644 --- a/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp @@ -28,8 +28,8 @@ struct Default2DEpilogueProblem static constexpr index_t NumDTensor = 0; }; -template { - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; using CLayout = remove_cvref_t; using DsDataType = remove_cvref_t; using CDElementwise = remove_cvref_t; @@ -157,14 +157,28 @@ struct Default2DEpilogue template struct DefaultGemm2DEpilogue : public Default2DEpilogue { - using Problem = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using AccDataType = remove_cvref_t; - using ODataType = remove_cvref_t; + using Problem = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using AccDataType = remove_cvref_t; + using ODataType = remove_cvref_t; + static constexpr bool ADataTypeIsTuple = is_detected::value; + static constexpr bool BDataTypeIsTuple = is_detected::value; + + using AsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using BsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using ADataType = remove_cvref_t{}, AsDataTypeTuple>>; + using BDataType = remove_cvref_t{}, BsDataTypeTuple>>; // Used for weight-only quantization kernel, B would be dequantized to the same data type as A using BTypeToUse = std::conditional_t, ADataType, BDataType>; + using DsDataType = remove_cvref_t; using DsLayout = remove_cvref_t; using CDElementwise = remove_cvref_t; diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp index 6405ca50df..58fdad149a 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp @@ -1446,29 +1446,35 @@ struct FmhaFwdKernel auto o_acc_tile = [&]() { if constexpr(kDoFp8StaticQuant) { - return FmhaPipeline{}( - q_dram_window, - identity{}, // q_element_func - k_dram_window, - identity{}, // k_element_func - v_dram_window, - identity{}, // v_element_func - bias_dram_window, - identity{}, // bias_element_func - randval_dram_window, - lse_dram_window, - identity{}, // lse_element_func - identity{}, // s_acc_element_func - scales{kargs.scale_p}, // p_compute_element_func - composes(saturates{}, scales{kargs.scale_o}), // o_acc_element_func - mask, - position_encoding, - kargs.scale_s, - variant, - variant_params, - block_indices, - smem_ptr, - dropout); + auto o_acc_element_func = [&]() { + if constexpr(std::is_same_v) + return ck_tile::composes(ck_tile::saturates{}, + ck_tile::scales{kargs.scale_o}); + else + return ck_tile::scales{kargs.scale_o}; + }(); + return FmhaPipeline{}(q_dram_window, + identity{}, // q_element_func + k_dram_window, + identity{}, // k_element_func + v_dram_window, + identity{}, // v_element_func + bias_dram_window, + identity{}, // bias_element_func + randval_dram_window, + lse_dram_window, + identity{}, // lse_element_func + identity{}, // s_acc_element_func + scales{kargs.scale_p}, // p_compute_element_func + o_acc_element_func, // o_acc_element_func + mask, + position_encoding, + kargs.scale_s, + variant, + variant_params, + block_indices, + smem_ptr, + dropout); } else { diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp index b883aad155..c402eaeac4 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp @@ -559,6 +559,9 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP auto shuffled_bias_tile = make_static_distributed_tensor( Policy::template MakeShuffledBiasTileDistribution()); shuffle_tile(shuffled_bias_tile, bias_tile); + // SGrad and Bias use the same address in LDS, finish loading ds on the previous + // iteration to reuse LDS. + block_sync_lds(); store_tile(bias_lds_write_window, shuffled_bias_tile); block_sync_lds(); auto bias_s_tile = load_tile(bias_s_lds_read_window); @@ -814,6 +817,9 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP auto shuffled_bias_tile = make_static_distributed_tensor( Policy::template MakeShuffledBiasTileDistribution()); shuffle_tile(shuffled_bias_tile, bias_tile); + // SGrad and Bias use the same address in LDS, finish loading ds in the hot loop to + // reuse LDS. + block_sync_lds(); store_tile(bias_lds_write_window, shuffled_bias_tile); block_sync_lds(); auto bias_s_tile = load_tile(bias_s_lds_read_window); @@ -956,6 +962,8 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP return cast_tile(ds); } }(); + // Finish loading bias_s to reuse LDS. + block_sync_lds(); store_tile(bias_lds_write_window, dbias); block_sync_lds(); auto shuffled_dbias_tile = load_tile(dbias_lds_read_window); @@ -975,11 +983,9 @@ struct BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP gemm_3(dk_acc, dst_reg_tensor, qt_reg_tensor); - if constexpr(kHasBiasGrad) - { - // SGrad and BiasGrad use the same address in LDS. - block_sync_lds(); - } + // SGrad and Bias/BiasGrad use the same address in LDS, finish loading bias/dbias or, when + // bias is not used, loading ds in the hot loop to reuse LDS. + block_sync_lds(); store_tile(ds_lds_window, ds_gemm); block_sync_lds(); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_kr_ktr_vr.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_kr_ktr_vr.hpp index 81950bd30a..41cb4fc306 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_kr_ktr_vr.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_kr_ktr_vr.hpp @@ -698,6 +698,12 @@ struct BlockFmhaBwdDQDKDVPipelineTrLoadKRKTRVR dst_reg_tensor.get_thread_buffer() = ds_gemm.get_thread_buffer(); gemm_3(dk_acc, dst_reg_tensor, qt_reg_tensor); + if constexpr(kHasBiasGrad) + { + // SGrad and BiasGrad use the same address in LDS, finish loading dbias to reuse + // LDS. + block_sync_lds(); + } store_tile(ds_lds_window, ds_gemm); } s_waitcnt(); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp index 16d9f695df..6d90429407 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_trload_qr_qtr_dor.hpp @@ -489,7 +489,7 @@ struct BlockFmhaBwdDQDKDVPipelineTrLoadQRQTRDOR move_tile_window(k_dram_window, {kN0, 0}); async_load_tile(v_lds_write_window, v_dram_window); move_tile_window(v_dram_window, {kN0, 0}); - // __builtin_amdgcn_s_waitcnt(0); + s_waitcnt(); k_reg_tensor = load_tile(k_lds_read_window); v_reg_tensor = load_tile(v_lds_read_window); kt_reg_tensor = load_tile_transpose(kt_lds_read_window); @@ -636,7 +636,7 @@ struct BlockFmhaBwdDQDKDVPipelineTrLoadQRQTRDOR } }(); store_tile(bias_lds_write_window, dbias); - __builtin_amdgcn_s_waitcnt(3952); + s_waitcnt(); block_sync_lds(); auto shuffled_dbias_tile = load_tile(dbias_lds_read_window); auto dbias_tile = make_static_distributed_tensor( @@ -656,9 +656,15 @@ struct BlockFmhaBwdDQDKDVPipelineTrLoadQRQTRDOR dst_reg_tensor.get_thread_buffer() = ds_gemm.get_thread_buffer(); dk_acc = gemm_3(dst_reg_tensor, qt_reg_tensor); + if constexpr(kHasBiasGrad) + { + // SGrad and BiasGrad use the same address in LDS, finish loading dbias to reuse + // LDS. + block_sync_lds(); + } store_tile(ds_lds_window, ds_gemm); } - __builtin_amdgcn_s_waitcnt(3952); + s_waitcnt(); block_sync_lds(); if constexpr(is_epilogue) { diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp index 68ead7c765..ad9e2959f5 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp @@ -1941,7 +1941,7 @@ struct BlockFmhaBwdPipelineDefaultPolicy constexpr index_t smem_size_stage0_0 = smem_size_k + smem_size_kt; constexpr index_t smem_size_stage0_1 = smem_size_v; - constexpr index_t smem_size_stage1 = smem_size_qt + smem_size_q + +smem_size_dot + + constexpr index_t smem_size_stage1 = smem_size_qt + smem_size_q + smem_size_dot + smem_size_do + smem_size_lse + smem_size_d + max(smem_size_bias, smem_size_ds); diff --git a/include/ck_tile/ops/gemm.hpp b/include/ck_tile/ops/gemm.hpp index de13e305e0..6e07dbc00e 100644 --- a/include/ck_tile/ops/gemm.hpp +++ b/include/ck_tile/ops/gemm.hpp @@ -31,6 +31,7 @@ #include "ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp" #include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp" #include "ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp" +#include "ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp" #include "ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp" #include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp" #include "ck_tile/ops/gemm/kernel/streamk_gemm_kernel.hpp" diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp index e1b0792ecf..94adb42880 100644 --- a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp @@ -4,6 +4,7 @@ #pragma once #include "ck_tile/core.hpp" +#include "ck_tile/ops/common/load_interleaved_pk_type.hpp" #include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/elementwise.hpp" @@ -13,7 +14,9 @@ namespace ck_tile { // A is block window on shared memory // B is block window on shared memory // C is block distributed tensor -template +template struct BlockUniversalGemmAsBsCr { private: @@ -91,6 +94,7 @@ struct BlockUniversalGemmAsBsCr using ComputeDataType = remove_cvref_t; using CDataType = remove_cvref_t; + using Loader = remove_cvref_t>; using WarpGemm = remove_cvref_t; static constexpr index_t KIterPerWarp = Traits::KIterPerWarp; @@ -179,25 +183,6 @@ struct BlockUniversalGemmAsBsCr return b_block_dstr_encode; } - private: - template - CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, - const WarpWindow& warp_window) - { - constexpr index_t UnaryOpSize = 8; - const element_wise::PassThroughPack8 elementwise_op{}; - constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; - const auto in_dstr_tensors = load_tile(warp_window); - - static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); - - using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); - static_for<0, thread_buffer_size, 1>{}([&](auto i) { - elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), - in_dstr_tensors.get_thread_buffer().template get_as()[i]); - }); - } - template struct BlockGemmImpl { @@ -239,7 +224,7 @@ struct BlockUniversalGemmAsBsCr if constexpr(std::is_same_v) { - load_interleaved_pk_type(a_warp_tile_, a_block_window); + Loader::load_interleaved_pk_type(a_warp_tile_, a_block_window); } else { @@ -247,7 +232,7 @@ struct BlockUniversalGemmAsBsCr } if constexpr(std::is_same_v) { - load_interleaved_pk_type(b_warp_tile_, b_block_window); + Loader::load_interleaved_pk_type(b_warp_tile_, b_block_window); } else { @@ -317,7 +302,7 @@ struct BlockUniversalGemmAsBsCr { if constexpr(std::is_same_v) { - load_interleaved_pk_type(a_warp_tile_, a_block_window); + Loader::load_interleaved_pk_type(a_warp_tile_, a_block_window); } else if constexpr(ALoadTranspose) { @@ -329,7 +314,7 @@ struct BlockUniversalGemmAsBsCr } if constexpr(std::is_same_v) { - load_interleaved_pk_type(b_warp_tile_, b_block_window); + Loader::load_interleaved_pk_type(b_warp_tile_, b_block_window); } else if constexpr(BLoadTranspose) { @@ -468,7 +453,7 @@ struct BlockUniversalGemmAsBsCr if constexpr(std::is_same_v) { - load_interleaved_pk_type(a_warp_tile_, a_block_window); + Loader::load_interleaved_pk_type(a_warp_tile_, a_block_window); } else if constexpr(ALoadTranspose) { @@ -480,7 +465,7 @@ struct BlockUniversalGemmAsBsCr } if constexpr(std::is_same_v) { - load_interleaved_pk_type(b_warp_tile_, b_block_window); + Loader::load_interleaved_pk_type(b_warp_tile_, b_block_window); } else if constexpr(BLoadTranspose) { diff --git a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp index fcfbf9635f..588d903b25 100644 --- a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp @@ -90,10 +90,10 @@ struct BatchedGemmKernel !is_detected::value && !is_detected::value, "BLayout and BDataType must be scalars. Multiple parameters are not currently supported."); - /// @brief C/ELayout and C/EDataType are expected to be scalars, not a tuple. + /// @brief C/CLayout and C/EDataType are expected to be scalars, not a tuple. static_assert(!is_detected::value && !is_detected::value, - "C/ELayout and C/EDataType must be scalars."); + "C/CLayout and C/EDataType must be scalars."); struct BatchedGemmKernelArgs : ck_tile::UniversalGemmKernelArgs<> { diff --git a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp index e37b4f36d4..d632b1596c 100644 --- a/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp @@ -89,7 +89,7 @@ struct GemmKernel /// @brief Specify the layout configurations for A, B, E and D using ALayout = remove_cvref_t; using BLayout = remove_cvref_t; - using ELayout = remove_cvref_t; + using CLayout = remove_cvref_t; /// @brief Specify the data type configurations for A, B, E and D using ADataType = remove_cvref_t; @@ -106,10 +106,10 @@ struct GemmKernel !is_detected::value && !is_detected::value, "BLayout and BDataType must be scalars. Multiple parameters are not currently supported."); - /// @brief C/ELayout and C/EDataType are expected to be scalars, not a tuple. - static_assert(!is_detected::value && + /// @brief C/CLayout and C/EDataType are expected to be scalars, not a tuple. + static_assert(!is_detected::value && !is_detected::value, - "C/ELayout and C/EDataType must be scalars."); + "C/CLayout and C/EDataType must be scalars."); static constexpr index_t NumATensor = 1; static constexpr index_t NumBTensor = 1; diff --git a/include/ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp new file mode 100644 index 0000000000..3b050e03ed --- /dev/null +++ b/include/ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp @@ -0,0 +1,193 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/common.hpp" +#include "ck_tile/host/concat.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/host/stream_utils.hpp" +#include "ck_tile/core/utility/env.hpp" +#include "ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp" +#include "ck_tile/core/utility/type_traits.hpp" + +namespace ck_tile { + +/// @brief The MultiABD GEMM kernel host arguments. +/// +/// @par Overview +/// This structure is passed to @ref GemmKernelMultiABD "GemmKernelMultiABD" when creating +/// kernel arguments object. It contain all necessary information required to build proper +/// kernel argument and launch kernel on GPU. This structure defines the GEMM problem +/// configuration by stating all required information like M,N,K sizes and respective strides. +/// NumATensor describes the number of A tensors. The minimum number of tensors is 1(required). +/// NumBTensor describes the number of B tensors. The minimum number of tensors is 1(required). +/// NumDTensor describes the number of D tensors. The minimum number of tensors is 0(not +/// required). +template +struct GemmMultiABDHostArgs +{ + CK_TILE_HOST GemmMultiABDHostArgs(const std::array& as_ptr_, + const std::array& bs_ptr_, + const std::array& ds_ptr_, + void* e_ptr_, + index_t k_batch_, + index_t M_, + index_t N_, + index_t K_, + const std::array& stride_As_, + const std::array& stride_Bs_, + const std::array& stride_Ds_, + index_t stride_E_) + : as_ptr(as_ptr_), + bs_ptr(bs_ptr_), + ds_ptr(ds_ptr_), + e_ptr(e_ptr_), + M(M_), + N(N_), + K(K_), + stride_As(stride_As_), + stride_Bs(stride_Bs_), + stride_Ds(stride_Ds_), + stride_E(stride_E_), + k_batch(k_batch_) + { + } + + const std::array as_ptr; + const std::array bs_ptr; + const std::array ds_ptr; + union + { + void* e_ptr; + void* c_ptr; + }; + index_t M; + index_t N; + index_t K; + const std::array stride_As; + const std::array stride_Bs; + const std::array stride_Ds; + union + { + index_t stride_E; + index_t stride_C; + }; + + index_t k_batch; +}; + +template +struct GemmKernelMultiABD +{ + /// @brief Inject the UniversalGemmKernel base class to support execution of all necessary + /// functions. + using UniversalGemmKernel = + UniversalGemmKernel; + static constexpr index_t kBlockSize = UniversalGemmKernel::kBlockSize; + + using TilePartitioner = remove_cvref_t; + using GemmPipeline = remove_cvref_t; + using EpiloguePipeline = remove_cvref_t; + + /// @brief Specify the layout configurations for A, B, E and D + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + using DsLayout = remove_cvref_t; + + /// @brief Specify the data type configurations for A, B, E and D + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using EDataType = remove_cvref_t; + using DsDataType = remove_cvref_t; + + /// @brief ALayout and ADataType are expected to be a tuple, not a scalar. + static_assert(is_detected::value && + is_detected::value, + "ALayout and ADataType must be a tuple."); + + /// @brief BLayout and BDataType are expected to be a tuple, not a scalar. + static_assert(is_detected::value && + is_detected::value, + "BLayout and BDataType must be a tuple."); + + /// @brief CLayout and EDataType are expected to be scalars, not a tuple. + static_assert(!is_detected::value && + !is_detected::value, + "CLayout and EDataType must be a scalar."); + + /// @brief DsLayout and DsDataType are expected to be tuple, not a scalar. + static_assert(is_detected::value && + is_detected::value && + DsLayout::size() == DsDataType::size() && DsLayout::size() > 0, + "DsLayout and DsDataType must be tuples and must have the same size."); + + /// @brief The sizes of NumATensor, NumBTensor and NumDTensor is set by the user." + static constexpr index_t NumATensor = AsDataType::size(); + static constexpr index_t NumBTensor = BsDataType::size(); + static constexpr index_t NumDTensor = DsDataType::size(); + + CK_TILE_HOST static auto GetName() -> const std::string + { + return UniversalGemmKernel::GetName(); + } + + CK_TILE_HOST static constexpr auto GridSize(index_t M, index_t N, index_t KBatch) -> dim3 + { + return UniversalGemmKernel::GridSize(M, N, KBatch); + } + + CK_TILE_HOST static auto MaxOccupancyGridSize(const stream_config& s) -> dim3 + { + return UniversalGemmKernel::MaxOccupancyGridSize(s); + } + + CK_TILE_HOST static constexpr auto BlockSize() -> dim3 + { + return UniversalGemmKernel::BlockSize(); + } + + CK_TILE_HOST static constexpr auto + MakeKernelArgs(const GemmMultiABDHostArgs& hostArgs) -> + typename UniversalGemmKernel::KernelArgs + { + /// @brief Universal GEMM requires array objects and corresponding stride information for + /// matrices A, B, and D. + return UniversalGemmKernel::MakeKernelArgs( + UniversalGemmHostArgs(hostArgs.as_ptr, + hostArgs.bs_ptr, + hostArgs.ds_ptr, + hostArgs.e_ptr, + hostArgs.k_batch, + hostArgs.M, + hostArgs.N, + hostArgs.K, + hostArgs.stride_As, + hostArgs.stride_Bs, + hostArgs.stride_Ds, + hostArgs.stride_E)); + } + + CK_TILE_HOST static auto + IsSupportedArgument(const typename UniversalGemmKernel::KernelArgs& kargs) -> bool + { + // Currently MultiABD kernel doesn't support k_batch > 1 + if(kargs.k_batch > 1) + { + return false; + } + + return UniversalGemmKernel::IsSupportedArgument(kargs); + } + + CK_TILE_DEVICE auto operator()(typename UniversalGemmKernel::KernelArgs kargs) const -> void + { + UniversalGemmKernel{}.template operator()(kargs); + } +}; +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp b/include/ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp index 9d3ac8b901..b0b2905cb4 100644 --- a/include/ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp @@ -95,7 +95,7 @@ struct GemmKernelMultiD /// @brief Specify the layout configurations for A, B, E and D using ALayout = remove_cvref_t; using BLayout = remove_cvref_t; - using ELayout = remove_cvref_t; + using CLayout = remove_cvref_t; using DsLayout = remove_cvref_t; /// @brief Specify the data type configurations for A, B, E and D @@ -114,10 +114,10 @@ struct GemmKernelMultiD !is_detected::value, "BLayout and BDataType must be scalars."); - /// @brief ELayout and EDataType are expected to be scalars, not a tuple. - static_assert(!is_detected::value && + /// @brief CLayout and EDataType are expected to be scalars, not a tuple. + static_assert(!is_detected::value && !is_detected::value, - "ELayout and EDataType must be scalars."); + "CLayout and EDataType must be scalars."); /// @brief DsLayout and DsDataType are expected to be tuple, not a scalar. static_assert(is_detected::value && diff --git a/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp index e38e49f5d1..df1d6c9e4f 100644 --- a/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp @@ -120,10 +120,10 @@ struct GroupedGemmKernel !is_detected::value && !is_detected::value, "BLayout and BDataType must be scalars. Multiple parameters are not currently supported."); - /// @brief C/ELayout and C/EDataType are expected to be scalars, not a tuple. + /// @brief C/CLayout and C/EDataType are expected to be scalars, not a tuple. static_assert(!is_detected::value && !is_detected::value, - "C/ELayout and C/EDataType must be scalars."); + "C/CLayout and C/EDataType must be scalars."); using OffsetTile1DPartitioner = OffsettedTile1DPartitioner; using Kernel = GroupedGemmKernel; @@ -364,12 +364,8 @@ struct GroupedGemmKernel const TailNumber tail_num = GemmPipeline::GetBlockLoopTailNum(num_loop); // Run GEMM pipeline - const auto& c_block_tile = GemmPipeline{}.template operator()(a_block_window[Base::I0], - b_block_window[Base::I0], - num_loop, - has_hot_loop, - tail_num, - smem_ptr_0); + const auto& c_block_tile = GemmPipeline{}.template operator()( + a_block_window, b_block_window, num_loop, has_hot_loop, tail_num, smem_ptr_0); // Run Epilogue Pipeline auto& c_block_window = gemm_tile_windows.at(Base::I3); EpiloguePipeline{}.template diff --git a/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp index cfba8b6c9d..8f44108cc4 100644 --- a/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp @@ -157,23 +157,23 @@ struct UniversalGemmKernel using EpiloguePipeline = remove_cvref_t; static constexpr bool ADataTypeIsTuple = - is_detected::value; + is_detected::value; static constexpr bool BDataTypeIsTuple = - is_detected::value; + is_detected::value; static constexpr bool DDataTypeIsTuple = is_detected::value; static constexpr bool ALayoutIsTuple = - is_detected::value; + is_detected::value; static constexpr bool BLayoutIsTuple = - is_detected::value; + is_detected::value; static constexpr bool DLayoutIsTuple = is_detected::value; using AsLayout = std::conditional_t, + remove_cvref_t, remove_cvref_t>>; using BsLayout = std::conditional_t, + remove_cvref_t, remove_cvref_t>>; using DsLayout = std::conditional_t>>; using AsDataType = std::conditional_t, + remove_cvref_t, remove_cvref_t>>; using BsDataType = std::conditional_t, + remove_cvref_t, remove_cvref_t>>; using DsDataType = @@ -193,9 +193,12 @@ struct UniversalGemmKernel remove_cvref_t, remove_cvref_t>>; - using ELayout = remove_cvref_t; + using CLayout = remove_cvref_t; using EDataType = remove_cvref_t; + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; + static constexpr index_t kBlockSize = GemmPipeline::BlockSize; // Get the persistent kernel if the pipeline has it available @@ -483,7 +486,7 @@ struct UniversalGemmKernel bool DTesnorIsValid = {true}; static_for<0, NumDTensor, 1>{}([&](auto index) { using DiLayout = remove_cvref_t>; - if(std::is_same_v == false) + if(std::is_same_v == false) { DTesnorIsValid = false; } @@ -529,7 +532,7 @@ struct UniversalGemmKernel } }); - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) { if(kargs.N % TilePartitioner::NPerBlock != 0 && GemmPipeline::kPadN == false) { @@ -724,7 +727,7 @@ struct UniversalGemmKernel // TODO: enable vector write for C in ColMajor const auto& e_tensor_view = [&]() { - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) { return make_naive_tensor_view( e_ptr, @@ -818,7 +821,7 @@ struct UniversalGemmKernel // TODO vector write in for C in ColMajor const auto& e_pad_view = [&]() { const auto& e_tensor_view = views.at(I3); - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) { return pad_tensor_view(e_tensor_view, make_tuple(number{}, @@ -975,8 +978,8 @@ struct UniversalGemmKernel const auto& bs_block_window = gemm_tile_windows.at(I1); const auto& ds_block_window = gemm_tile_windows.at(I2); - const auto& c_block_tile = - GemmPipeline{}(as_block_window[I0], bs_block_window[I0], num_loop, smem_ptr_0); + const auto& c_block_tile = GemmPipeline{}.template operator()( + as_block_window, AElementWise{}, bs_block_window, BElementWise{}, num_loop, smem_ptr_0); if(UseDefaultScheduler || (get_warp_id() == 0)) { @@ -1031,8 +1034,13 @@ struct UniversalGemmKernel const auto& bs_block_window = gemm_tile_windows.at(I1); const auto& ds_block_window = gemm_tile_windows.at(I2); - const auto& c_block_tile = GemmPipeline{}( - as_block_window[I0], bs_block_window[I0], num_loop, smem_ptr_0, smem_ptr_1); + const auto& c_block_tile = GemmPipeline{}.template operator()(as_block_window, + AElementWise{}, + bs_block_window, + BElementWise{}, + num_loop, + smem_ptr_0, + smem_ptr_1); // Run Epilogue Pipeline auto& c_block_window = gemm_tile_windows.at(I3); diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp index 2bee550b3c..b5584f98df 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp @@ -11,12 +11,17 @@ namespace ck_tile { template struct GemmPipelineAgBgCrImplBase { - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; using BlockGemmShape = remove_cvref_t; + using ADataType = remove_cvref_t{}, AsDataType>>; + using ALayout = remove_cvref_t{}, AsLayout>>; + using BDataType = remove_cvref_t{}, BsDataType>>; + using BLayout = remove_cvref_t{}, BsLayout>>; + static constexpr index_t MPerBlock = BlockGemmShape::kM; static constexpr index_t NPerBlock = BlockGemmShape::kN; static constexpr index_t KPerBlock = BlockGemmShape::kK; @@ -57,6 +62,13 @@ struct GemmPipelineAgBgCrImplBase store_tile(lds_tile_window, block_tile_tmp); } + template + CK_TILE_DEVICE void LocalPrefill(DstTileWindow& lds_tile_window, + const SrcBlockTile& src_block_tile) const + { + store_tile(lds_tile_window, src_block_tile); + } + template CK_TILE_DEVICE void LocalPrefetch(DstBlockTile& dst_block_tile, const SrcTileWindow& lds_tile_window, @@ -88,23 +100,100 @@ struct GemmPipelineAgBgCrImplBase return make_tuple(std::move(a_lds_block), std::move(b_lds_block)); } + template ::value, bool>* = + nullptr> + CK_TILE_DEVICE constexpr auto CopyADramWindow(const DramBlockWindowTmp& dram_block_window_tmp, + const array& offset = {0, 0}) const + { + constexpr bool is_col_major = std::is_same_v; + + using YPerTile = std::conditional_t, number>; + using XPerTile = std::conditional_t, number>; + // A DRAM tile window for load + auto a_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(YPerTile{}, XPerTile{}), + dram_block_window_tmp[number{}].get_window_origin() + offset, + Policy::template MakeADramTileDistribution()); + }, + number{}); + return std::move(a_copy_dram_window); + } + + template ::value, bool>* = + nullptr> + CK_TILE_DEVICE constexpr auto CopyADramWindow(const DramBlockWindowTmp& dram_block_window_tmp, + const array& offset = {0, 0}) const + { + constexpr bool is_col_major = std::is_same_v; + + using YPerTile = std::conditional_t, number>; + using XPerTile = std::conditional_t, number>; + // A DRAM tile window for load + auto a_copy_dram_window = + make_tile_window(dram_block_window_tmp.get_bottom_tensor_view(), + make_tuple(YPerTile{}, XPerTile{}), + dram_block_window_tmp.get_window_origin() + offset, + Policy::template MakeADramTileDistribution()); + + return std::move(a_copy_dram_window); + } + + template ::value, bool>* = + nullptr> + CK_TILE_DEVICE constexpr auto CopyBDramWindow(const DramBlockWindowTmp& dram_block_window_tmp, + const array& offset = {0, 0}) const + { + constexpr bool is_row_major = std::is_same_v; + + using YPerTile = std::conditional_t, number>; + using XPerTile = std::conditional_t, number>; + // A DRAM tile window for load + auto a_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(YPerTile{}, XPerTile{}), + dram_block_window_tmp[number{}].get_window_origin() + offset, + Policy::template MakeBDramTileDistribution()); + }, + number{}); + return std::move(a_copy_dram_window); + } + + template ::value, bool>* = + nullptr> + CK_TILE_DEVICE constexpr auto CopyBDramWindow(const DramBlockWindowTmp& dram_block_window_tmp, + const array& offset = {0, 0}) const + { + constexpr bool is_row_major = std::is_same_v; + + using YPerTile = std::conditional_t, number>; + using XPerTile = std::conditional_t, number>; + // A DRAM tile window for load + auto a_copy_dram_window = + make_tile_window(dram_block_window_tmp.get_bottom_tensor_view(), + make_tuple(YPerTile{}, XPerTile{}), + dram_block_window_tmp.get_window_origin() + offset, + Policy::template MakeBDramTileDistribution()); + + return std::move(a_copy_dram_window); + } + template CK_TILE_DEVICE constexpr auto GetAWindows(const ADramBlockWindowTmp& a_dram_block_window_tmp, const ALdsTensorView& a_lds_block_view, const ALdsLoadTileDistr&, const array& offset = {0, 0}) const { - constexpr bool is_col_major = std::is_same_v; - - using YPerTile = std::conditional_t, number>; - using XPerTile = std::conditional_t, number>; - // A DRAM tile window for load - auto a_copy_dram_window = - make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(YPerTile{}, XPerTile{}), - a_dram_block_window_tmp.get_window_origin() + offset, - Policy::template MakeADramTileDistribution()); + auto a_copy_dram_window = CopyADramWindow(a_dram_block_window_tmp, offset); // A LDS tile window for store auto a_lds_shape = []() { @@ -138,16 +227,8 @@ struct GemmPipelineAgBgCrImplBase const BLdsLoadTileDistr&, const array& offset = {0, 0}) const { - constexpr bool is_row_major = std::is_same_v; - - using YPerTile = std::conditional_t, number>; - using XPerTile = std::conditional_t, number>; - - auto b_copy_dram_window = - make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(YPerTile{}, XPerTile{}), - b_dram_block_window_tmp.get_window_origin() + offset, - Policy::template MakeBDramTileDistribution()); + // A DRAM tile window for load + auto b_copy_dram_window = CopyBDramWindow(b_dram_block_window_tmp, offset); // TODO: Do we really need those two tile windows??? // They're exactly same... diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp index f325595769..7fadafc649 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp @@ -107,14 +107,23 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 using Base = BaseGemmPipelineAgBgCrCompV3; using PipelineImplBase = GemmPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; using BlockGemm = remove_cvref_t())>; using I0 = number<0>; @@ -386,17 +395,25 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v auto block_gemm = BlockGemm(); auto c_block_tile = block_gemm.MakeCBlockTile(); - using ABlockTileDistr = decltype(a_copy_dram_window.get_tile_distribution()); - using BBlockTileDistr = decltype(b_copy_dram_window.get_tile_distribution()); - - using ABlockTile = - decltype(make_static_distributed_tensor(ABlockTileDistr{})); - using BBlockTile = - decltype(make_static_distributed_tensor(BBlockTileDistr{})); - - ABlockTile a_block_tile; - BBlockTile b_block_tile; - using ADramTileWindowStep = typename ADramBlockWindowTmp::BottomTensorIndex; using BDramTileWindowStep = typename BDramBlockWindowTmp::BottomTensorIndex; @@ -470,45 +476,61 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 // ----------------------------------------------------------------------------------------- // Gemm pipeline start - - // prefetch - // global read 0 - Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); - // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + auto elementwise_As_res = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + // Move each A — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + auto elementwise_Bs_res = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + // Move each B — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); + // LDS write 0 if constexpr(is_a_col_major && !is_a_load_tr_v()) { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_block_tile); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_block_tile); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window, elementwise_Bs_res); } - Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + // global read 1 + + elementwise_As_res = load_tile_with_elementwise(a_copy_dram_window, a_element_func); + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + elementwise_Bs_res = load_tile_with_elementwise(b_copy_dram_window, b_element_func); + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); block_sync_lds(); - block_gemm.LocalPrefetch( - a_lds_gemm_window, b_lds_gemm_window, is_a_load_tr_v, is_b_load_tr_v); + block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); __builtin_amdgcn_sched_barrier(0); @@ -520,38 +542,42 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 { block_sync_lds(); - if constexpr(is_a_col_major && !is_a_load_tr_v()) + if constexpr(is_a_col_major) { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_block_tile); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window, elementwise_As_res); } - if constexpr(is_b_row_major && !is_b_load_tr_v()) + if constexpr(is_b_row_major) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_block_tile); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window, elementwise_Bs_res); } - Base::GlobalPrefetch(a_block_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tile, b_copy_dram_window, b_dram_tile_window_step); + elementwise_As_res = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + elementwise_Bs_res = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); block_gemm(c_block_tile, a_lds_gemm_window, b_lds_gemm_window); block_sync_lds(); - block_gemm.LocalPrefetch( - a_lds_gemm_window, b_lds_gemm_window, is_a_load_tr_v, is_b_load_tr_v); + block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); HotLoopScheduler(); __builtin_amdgcn_sched_barrier(0); @@ -574,27 +600,26 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_block_tile); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_block_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_block_tile); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_block_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window, elementwise_Bs_res); } block_sync_lds(); - block_gemm.LocalPrefetch( - a_lds_gemm_window, b_lds_gemm_window, is_a_load_tr_v, is_b_load_tr_v); + block_gemm.LocalPrefetch(a_lds_gemm_window, b_lds_gemm_window); block_gemm(c_block_tile, a_lds_gemm_window, b_lds_gemm_window); } // __builtin_amdgcn_sched_barrier(0); @@ -602,13 +627,16 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 } }; - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const @@ -628,9 +656,13 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 * @note This is used by the persistent gemm kernel variants that don't determine * hot loop and tail number on the host side, e.g. grouped gemm kernel. */ - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, bool has_hot_loop, TailNumber tail_number, @@ -639,7 +671,7 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 const auto RunPipeline = [&](auto hot_loop_, auto tail_num_) { constexpr bool hot_loop = hot_loop_.value; constexpr auto tail_num = tail_num_.value; - constexpr auto PassThrough = [](const auto& x) { return x; }; + constexpr auto PassThrough = [](auto& e, const auto& x) { e = x; }; return PipelineImpl{}.template operator()( a_dram_block_window_tmp, PassThrough, @@ -658,20 +690,97 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 * @note This is used by the kernel variants that are able to determine * hot loop and tail number on the host side, e.g. non-persistent gemm kernel. */ - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, void* p_smem) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, - [](const ADataType& a) { return a; }, + [](auto& e, const ADataType& a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType& b) { return b; }, + [](auto& e, const BDataType& b) { e = b; }, num_loop, p_smem); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const AElementFunction& a_element_func, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, + const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + a_element_func, + ck_tile::make_tuple(b_dram_block_window_tmp), + b_element_func, + num_loop, + p_smem); + } + + /** + * @brief Quant operator(), single input: This function runs the pipeline by wrapping it with + * the tail handler. + * + * @note This is used by the persistent gemm kernel variants that don't determine + * hot loop and tail number on the host side, e.g. grouped gemm kernel. + */ + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + bool has_hot_loop, + TailNumber tail_number, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + has_hot_loop, + tail_number, + p_smem); + } + + /** + * @brief Quant operator(), single input: This function runs the pipeline using compile-time + * known hot loop and tail number. + * @param num_loop The number of loop iterations. This is determined at runtime due to e.g. + * SplitK. + * @note This is used by the kernel variants that are able to determine + * hot loop and tail number on the host side, e.g. non-persistent gemm kernel. + */ + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp index c835809b5d..b362f751c6 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp @@ -97,11 +97,24 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 using Base = BaseGemmPipelineAgBgCrCompV4; using PipelineImplBase = GemmPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; using CDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; + static_assert(!std::is_same_v, "Not implemented"); static constexpr index_t APackedSize = @@ -109,10 +122,6 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 static constexpr index_t BPackedSize = ck_tile::numeric_traits>::PackedSize; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; - using BlockGemm = remove_cvref_t())>; using I0 = number<0>; using I1 = number<1>; @@ -244,18 +253,26 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* __restrict__ p_smem_0, void* __restrict__ p_smem_1) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v KPerBlock == BDramBlockWindowTmp{}.get_window_lengths()[I1{}]), "B block window has incorrect lengths for defined BLayout!"); - ////////////// global window & register ///////////////// - // A DRAM tile window for load - auto a_copy_dram_window = - make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - a_dram_block_window_tmp.get_window_origin(), - Policy::template MakeADramTileDistribution()); - - // B DRAM tile window for load - auto b_copy_dram_window = - make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - b_dram_block_window_tmp.get_window_origin(), - Policy::template MakeBDramTileDistribution()); - - // A register tile for global load - constexpr auto ABlockTileDistr = a_copy_dram_window.get_tile_distribution(); - constexpr auto BBlockTileDistr = b_copy_dram_window.get_tile_distribution(); - using ABlockTile = decltype(make_static_distributed_tensor(ABlockTileDistr)); - using BBlockTile = decltype(make_static_distributed_tensor(BBlockTileDistr)); - ABlockTile a_global_load_tile; - BBlockTile b_global_load_tile; - using ADramTileWindowStep = typename ADramBlockWindowTmp::BottomTensorIndex; using BDramTileWindowStep = typename BDramBlockWindowTmp::BottomTensorIndex; @@ -312,8 +306,7 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 // global prefetch 0 // global read 0 - Base::GlobalPrefetch(a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + ////////////// LDS desc, window & register ///////////////// auto&& [a_lds_block0, b_lds_block0] = Base::GetABLdsTensorViews(p_smem_0); auto&& [a_lds_block1, b_lds_block1] = Base::GetABLdsTensorViews(p_smem_1); @@ -343,34 +336,75 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); + // Generating a tuple with tile_windows for values A0, A1, ... AN + auto a_tile_windows = generate_tuple( + [&](auto idx) { + return make_tile_window( + a_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + a_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeADramTileDistribution()); + }, + number{}); + + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + auto elementwise_As_res = load_tile_with_elementwise(a_tile_windows, a_element_func); + + // Move each A — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(a_tile_windows, a_dram_tile_window_step); + + // Generating a tuple with tile_windows for values B0, B1, ... BN + auto b_tile_windows = generate_tuple( + [&](auto idx) { + return make_tile_window( + b_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + b_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeBDramTileDistribution()); + }, + number{}); + + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + auto elementwise_Bs_res = load_tile_with_elementwise(b_tile_windows, b_element_func); + + // Move each B — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(b_tile_windows, b_dram_tile_window_step); + // LDS write 0 if constexpr(is_a_col_major && !is_a_load_tr_v()) { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window0, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window0, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window0, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window0, elementwise_Bs_res); } // global read 1 - Base::GlobalPrefetch(a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + elementwise_As_res = load_tile_with_elementwise(a_tile_windows, a_element_func); + move_tile_window(a_tile_windows, a_dram_tile_window_step); + + elementwise_Bs_res = load_tile_with_elementwise(b_tile_windows, b_element_func); + move_tile_window(b_tile_windows, b_dram_tile_window_step); block_sync_lds(); constexpr auto ALdsTileDistr = @@ -423,27 +457,32 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window1, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window1, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window1, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window1, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window1, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window1, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window1, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window1, elementwise_Bs_res); } - Base::GlobalPrefetch(a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch(b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + elementwise_As_res = load_tile_with_elementwise(a_tile_windows, a_element_func); + move_tile_window(a_tile_windows, a_dram_tile_window_step); + + elementwise_Bs_res = load_tile_with_elementwise(b_tile_windows, b_element_func); + move_tile_window(b_tile_windows, b_dram_tile_window_step); if(HasHotLoop) { @@ -461,31 +500,32 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp); } else { - Base::LocalPrefill( - a_copy_lds_window0, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window0, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp); } else { - Base::LocalPrefill( - b_copy_lds_window0, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window0, elementwise_Bs_res); } - Base::GlobalPrefetch( - a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch( - b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + elementwise_As_res = + load_tile_with_elementwise(a_tile_windows, a_element_func); + move_tile_window(a_tile_windows, a_dram_tile_window_step); + + elementwise_Bs_res = + load_tile_with_elementwise(b_tile_windows, b_element_func); + move_tile_window(b_tile_windows, b_dram_tile_window_step); // gemm block_gemm(c_block_tile, a_block_tile0, b_block_tile0); HotLoopScheduler(); @@ -501,32 +541,34 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window1, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window1, a_shuffle_tmp); } else { - Base::LocalPrefill( - a_copy_lds_window1, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window1, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window1, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window1, b_shuffle_tmp); } else { - Base::LocalPrefill( - b_copy_lds_window1, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window1, elementwise_Bs_res); } block_sync_lds(); - Base::GlobalPrefetch( - a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch( - b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + elementwise_As_res = + load_tile_with_elementwise(a_tile_windows, a_element_func); + move_tile_window(a_tile_windows, a_dram_tile_window_step); + + elementwise_Bs_res = + load_tile_with_elementwise(b_tile_windows, b_element_func); + move_tile_window(b_tile_windows, b_dram_tile_window_step); + // gemm block_gemm(c_block_tile, a_block_tile1, b_block_tile1); HotLoopScheduler(); @@ -548,23 +590,23 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window0, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window0, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window0, elementwise_As_res); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window0, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window0, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window0, elementwise_Bs_res); } block_gemm(c_block_tile, a_block_tile0, b_block_tile0); } @@ -606,13 +648,17 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 } }; - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem_0, @@ -628,27 +674,34 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 p_smem_1); } - public: - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const index_t num_loop, void* __restrict__ p_smem_0, void* __restrict__ p_smem_1) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, - [](const ADataType& a) { return a; }, + [](auto& e, const ADataType& a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType& b) { return b; }, + [](auto& e, const BDataType& b) { e = b; }, num_loop, p_smem_0, p_smem_1); } - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, bool has_hot_loop, TailNumber tail_number, @@ -658,7 +711,7 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 const auto RunPipeline = [&](auto hot_loop_, auto tail_num_) { constexpr bool hot_loop = hot_loop_.value; constexpr auto tail_num = tail_num_.value; - constexpr auto PassThrough = [](const auto& x) { return x; }; + constexpr auto PassThrough = [](auto& e, const auto& x) { e = x; }; return PipelineImpl{}.template operator()( a_dram_block_window_tmp, PassThrough, @@ -670,5 +723,69 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 }; return Base::TailHandler(RunPipeline, has_hot_loop, tail_number); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const AElementFunction& a_element_func, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem_0, + void* p_smem_1) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + a_element_func, + ck_tile::make_tuple(b_dram_block_window_tmp), + b_element_func, + num_loop, + p_smem_0, + p_smem_1); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + const index_t num_loop, + void* __restrict__ p_smem_0, + void* __restrict__ p_smem_1) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem_0, + p_smem_1); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + bool has_hot_loop, + TailNumber tail_number, + void* __restrict__ p_smem_0, + void* __restrict__ p_smem_1) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + has_hot_loop, + tail_number, + p_smem_0, + p_smem_1); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp index b83d37a790..474d1a5a21 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp @@ -41,15 +41,24 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 using Base = BaseGemmPipelineAgBgCrCompV5; using PipelineImplBase = GemmPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; using CDataType = remove_cvref_t; using ComputeDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; + + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; static constexpr index_t NumWaveGroups = Problem::NumWaveGroups; @@ -121,17 +130,25 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BsDramBlockWindowTmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* __restrict__ p_smem_0) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v BGemmTile b_tile_0, b_tile_1; // Register tile for A and B. - using ABlockTileDistr = decltype(a_copy_dram_window.get_tile_distribution()); - using BBlockTileDistr = decltype(b_copy_dram_window.get_tile_distribution()); + using ABlockTileDistr = + decltype(a_copy_dram_window[number<0>{}].get_tile_distribution()); + using BBlockTileDistr = + decltype(b_copy_dram_window[number<0>{}].get_tile_distribution()); using ABlockTile = decltype(make_static_distributed_tensor(ABlockTileDistr{})); using BBlockTile = decltype(make_static_distributed_tensor(BBlockTileDistr{})); - ABlockTile a_global_load_tile; - BBlockTile b_global_load_tile; + ABlockTile elementwise_As_res; + BBlockTile elementwise_Bs_res; // Block GEMM auto block_gemm = BlockGemm(); @@ -248,33 +267,45 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 // define ping, pong steps here as lambda functions. auto MemoryOpsStep = [&](auto idx) { // Memory read half here. - Base::GlobalPrefetch( - a_global_load_tile, a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch( - b_global_load_tile, b_copy_dram_window, b_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each + // A0, A1, … AN. The values A0, A1, … AN are read by the same thread. + elementwise_As_res = load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + // Move each A — the enhanced function move_tile_window is executed, which takes a + // tuple as input. + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each + // B0, B1, … BN. The values B0, B1, … BN are read by the same thread. + elementwise_Bs_res = load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + // Move each B — the enhanced function move_tile_window is executed, which takes a + // tuple as input. + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); if constexpr(is_a_col_major) { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_global_load_tile); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_global_load_tile, a_element_func); + Base::LocalPrefill(a_copy_lds_window, elementwise_As_res); } if constexpr(is_b_row_major) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_global_load_tile); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_global_load_tile, b_element_func); + Base::LocalPrefill(b_copy_lds_window, elementwise_Bs_res); } if(idx == 0) @@ -351,13 +382,17 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 } }; - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem_0) const @@ -371,21 +406,62 @@ struct GemmPipelineAgBgCrCompV5 : public BaseGemmPipelineAgBgCrCompV5 p_smem_0); } - public: - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const index_t num_loop, void* __restrict__ p_smem_0) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, - [](const ADataType& a) { return a; }, + [](auto& e, const ADataType& a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType& b) { return b; }, + [](auto& e, const BDataType& b) { e = b; }, num_loop, p_smem_0); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const AElementFunction& a_element_func, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem_0) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + a_element_func, + ck_tile::make_tuple(b_dram_block_window_tmp), + b_element_func, + num_loop, + p_smem_0); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + const index_t num_loop, + void* __restrict__ p_smem_0) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem_0); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp index e1acfebc47..9e522d4364 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp @@ -157,14 +157,23 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem using Base = BaseGemmPipelineAgBgCrMem; using PipelineImplBase = GemmPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; using BlockGemm = remove_cvref_t())>; @@ -236,17 +245,25 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v auto block_gemm = BlockGemm(); auto c_block_tile = block_gemm.MakeCBlockTile(); - using ABlockTileDistr = decltype(a_copy_dram_window.get_tile_distribution()); - using BBlockTileDistr = decltype(b_copy_dram_window.get_tile_distribution()); + using ABlockTileDistr = + decltype(a_copy_dram_window[number<0>{}].get_tile_distribution()); + using BBlockTileDistr = + decltype(b_copy_dram_window[number<0>{}].get_tile_distribution()); using ABlockTile = decltype(make_static_distributed_tensor(ABlockTileDistr{})); @@ -334,10 +353,21 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem // prefetch // global read 0 - Base::GlobalPrefetch( - a_block_tiles.get(I0{}), a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch( - b_block_tiles.get(I0{}), b_copy_dram_window, b_dram_tile_window_step); + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + a_block_tiles.at(I0{}) = load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + // Move each A — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + b_block_tiles.at(I0{}) = load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + // Move each B — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); @@ -348,32 +378,35 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(I0{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{}), a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(I0{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{}), b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{})); } // Global prefetch [1, PrefetchStages] static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) { - Base::GlobalPrefetch(a_block_tiles.get(number{}), - a_copy_dram_window, - a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tiles.get(number{}), - b_copy_dram_window, - b_dram_tile_window_step); + a_block_tiles.at(number{}) = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + b_block_tiles.at(number{}) = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); }); // main body @@ -397,14 +430,13 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem transpose_tile2d( a_shuffle_tmp, a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { Base::LocalPrefill( a_copy_lds_window, - a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}), - a_element_func); + a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { @@ -413,22 +445,23 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem transpose_tile2d( b_shuffle_tmp, b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { Base::LocalPrefill( b_copy_lds_window, - b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}), - b_element_func); + b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); } - Base::GlobalPrefetch(a_block_tiles.get(number{}), - a_copy_dram_window, - a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tiles.get(number{}), - b_copy_dram_window, - b_dram_tile_window_step); + a_block_tiles.at(number{}) = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + b_block_tiles.at(number{}) = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); }); i += PrefetchStages; @@ -450,26 +483,24 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(number{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { Base::LocalPrefill(a_copy_lds_window, - a_block_tiles.get(number{}), - a_element_func); + a_block_tiles.get(number{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(number{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { Base::LocalPrefill(b_copy_lds_window, - b_block_tiles.get(number{}), - b_element_func); + b_block_tiles.get(number{})); } }); @@ -526,17 +557,25 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v auto block_gemm = BlockGemm(); auto c_block_tile = block_gemm.MakeCBlockTile(); - using ABlockTileDistr = decltype(a_copy_dram_window.get_tile_distribution()); - using BBlockTileDistr = decltype(b_copy_dram_window.get_tile_distribution()); + using ABlockTileDistr = + decltype(a_copy_dram_window[number<0>{}].get_tile_distribution()); + using BBlockTileDistr = + decltype(b_copy_dram_window[number<0>{}].get_tile_distribution()); using ABlockTile = decltype(make_static_distributed_tensor(ABlockTileDistr{})); @@ -623,10 +664,22 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem // prefetch // global read 0 - Base::GlobalPrefetch( - a_block_tiles.get(I0{}), a_copy_dram_window, a_dram_tile_window_step); - Base::GlobalPrefetch( - b_block_tiles.get(I0{}), b_copy_dram_window, b_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + a_block_tiles.at(I0{}) = load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + // Move each A — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + b_block_tiles.at(I0{}) = load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + // Move each B — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); @@ -637,32 +690,35 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(I0{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { - Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{}), a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_block_tiles.get(I0{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(I0{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { - Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{}), b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{})); } // Global prefetch [1, PrefetchStages] static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) { - Base::GlobalPrefetch(a_block_tiles.get(number{}), - a_copy_dram_window, - a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tiles.get(number{}), - b_copy_dram_window, - b_dram_tile_window_step); + a_block_tiles.at(number{}) = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + b_block_tiles.at(number{}) = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); }); // main body @@ -687,14 +743,13 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem transpose_tile2d( a_shuffle_tmp, a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { Base::LocalPrefill( a_copy_lds_window, - a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}), - a_element_func); + a_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { @@ -703,22 +758,24 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem transpose_tile2d( b_shuffle_tmp, b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { Base::LocalPrefill( b_copy_lds_window, - b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{}), - b_element_func); + b_block_tiles.get(number<(prefetch_idx + 1) % PrefetchStages>{})); } - Base::GlobalPrefetch(a_block_tiles.get(number{}), - a_copy_dram_window, - a_dram_tile_window_step); - Base::GlobalPrefetch(b_block_tiles.get(number{}), - b_copy_dram_window, - b_dram_tile_window_step); + a_block_tiles.at(number{}) = + load_tile_with_elementwise(a_copy_dram_window, a_element_func); + + move_tile_window(a_copy_dram_window, a_dram_tile_window_step); + + b_block_tiles.at(number{}) = + load_tile_with_elementwise(b_copy_dram_window, b_element_func); + + move_tile_window(b_copy_dram_window, b_dram_tile_window_step); }); i += PrefetchStages; @@ -740,26 +797,24 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); transpose_tile2d(a_shuffle_tmp, a_block_tiles.get(number{})); - Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp, a_element_func); + Base::LocalPrefill(a_copy_lds_window, a_shuffle_tmp); } else { Base::LocalPrefill(a_copy_lds_window, - a_block_tiles.get(number{}), - a_element_func); + a_block_tiles.get(number{})); } if constexpr(is_b_row_major && !is_b_load_tr_v()) { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); transpose_tile2d(b_shuffle_tmp, b_block_tiles.get(number{})); - Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp, b_element_func); + Base::LocalPrefill(b_copy_lds_window, b_shuffle_tmp); } else { Base::LocalPrefill(b_copy_lds_window, - b_block_tiles.get(number{}), - b_element_func); + b_block_tiles.get(number{})); } }); @@ -813,13 +868,16 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem } }; - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const @@ -833,9 +891,13 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem p_smem); } - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, bool has_hot_loop, TailNumber tail_number, @@ -844,7 +906,7 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem const auto RunPipeline = [&](auto hot_loop_, auto tail_num_) { constexpr bool hot_loop = hot_loop_.value; constexpr auto tail_num = tail_num_.value; - constexpr auto PassThrough = [](const auto& x) { return x; }; + constexpr auto PassThrough = [](auto& e, const auto& x) { e = x; }; return PipelineImpl{}.template operator()( a_dram_block_window_tmp, PassThrough, @@ -856,20 +918,82 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem return Base::TailHandler(RunPipeline, has_hot_loop, tail_number); } - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, void* p_smem) const { return PipelineImpl{}.template operator()( a_dram_block_window_tmp, - [](const ADataType& a) { return a; }, + [](auto& e, const ADataType& a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType& b) { return b; }, + [](auto& e, const ADataType& a) { e = a; }, num_loop, p_smem); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const AElementFunction& a_element_func, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + a_element_func, + ck_tile::make_tuple(b_dram_block_window_tmp), + b_element_func, + num_loop, + p_smem); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + bool has_hot_loop, + TailNumber tail_number, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + has_hot_loop, + tail_number, + p_smem); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp index e3b4863392..eb363d59b8 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp @@ -15,14 +15,23 @@ namespace ck_tile { template struct GemmPipelineAGmemBGmemCRegV1 { - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; using BlockGemm = remove_cvref_t())>; @@ -81,17 +90,25 @@ struct GemmPipelineAGmemBGmemCRegV1 return Policy::template GetSmemSize(); } - template - CK_TILE_HOST_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_HOST_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const { + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v>, @@ -133,22 +150,30 @@ struct GemmPipelineAGmemBGmemCRegV1 auto b_lds_block = make_tensor_view(p_b_lds, b_lds_block_desc); // A DRAM tile window for load - auto a_copy_dram_window = - make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - a_dram_block_window_tmp.get_window_origin(), - Policy::template MakeADramTileDistribution()); + auto as_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + a_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + a_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeADramTileDistribution()); + }, + number{}); // A LDS tile window for store auto a_copy_lds_window = make_tile_window( a_lds_block, make_tuple(number{}, number{}), {0, 0}); // B DRAM tile window for load - auto b_copy_dram_window = - make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - b_dram_block_window_tmp.get_window_origin(), - Policy::template MakeBDramTileDistribution()); + auto bs_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + b_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + b_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeBDramTileDistribution()); + }, + number{}); // B LDS tile window for store auto b_copy_lds_window = make_tile_window( @@ -182,13 +207,22 @@ struct GemmPipelineAGmemBGmemCRegV1 // prefetch // global read 0 - auto a_block_tile = load_tile(a_copy_dram_window); - auto b_block_tile = load_tile(b_copy_dram_window); + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + auto elementwise_As_res = load_tile_with_elementwise(as_copy_dram_window, a_element_func); + + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + auto elementwise_Bs_res = load_tile_with_elementwise(bs_copy_dram_window, b_element_func); { // move to 1 - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - move_tile_window(b_copy_dram_window, {0, kKPerBlock}); + // Move each A — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(as_copy_dram_window, {0, kKPerBlock}); + // Move each B — the enhanced function move_tile_window is executed, which takes a tuple + // as input. + move_tile_window(bs_copy_dram_window, {0, kKPerBlock}); // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); @@ -198,13 +232,12 @@ struct GemmPipelineAGmemBGmemCRegV1 { auto a_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp, a_block_tile); - const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_shuffle_tmp); - store_tile(a_copy_lds_window, a_block_tile_tmp); + transpose_tile2d(a_shuffle_tmp, elementwise_As_res); + store_tile(a_copy_lds_window, a_shuffle_tmp); } else { - store_tile(a_copy_lds_window, tile_elementwise_in(a_element_func, a_block_tile)); + store_tile(a_copy_lds_window, elementwise_As_res); } // LDS write 0 @@ -212,13 +245,12 @@ struct GemmPipelineAGmemBGmemCRegV1 { auto b_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp, b_block_tile); - const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_shuffle_tmp); - store_tile(b_copy_lds_window, b_block_tile_tmp); + transpose_tile2d(b_shuffle_tmp, elementwise_Bs_res); + store_tile(b_copy_lds_window, b_shuffle_tmp); } else { - store_tile(b_copy_lds_window, tile_elementwise_in(b_element_func, b_block_tile)); + store_tile(b_copy_lds_window, elementwise_Bs_res); } } @@ -226,8 +258,8 @@ struct GemmPipelineAGmemBGmemCRegV1 while(iCounter > 0) { // global read i + 1 - a_block_tile = load_tile(a_copy_dram_window); - b_block_tile = load_tile(b_copy_dram_window); + elementwise_As_res = load_tile_with_elementwise(as_copy_dram_window, a_element_func); + elementwise_Bs_res = load_tile_with_elementwise(bs_copy_dram_window, b_element_func); block_sync_lds(); @@ -237,22 +269,20 @@ struct GemmPipelineAGmemBGmemCRegV1 block_sync_lds(); // move to i + 2 - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - move_tile_window(b_copy_dram_window, {0, kKPerBlock}); + move_tile_window(as_copy_dram_window, {0, kKPerBlock}); + move_tile_window(bs_copy_dram_window, {0, kKPerBlock}); // LDS write i + 1 if constexpr(is_a_col_major) { auto a_shuffle_tmp_loop = make_static_distributed_tensor( Policy::template MakeShuffledARegTileDistribution()); - transpose_tile2d(a_shuffle_tmp_loop, a_block_tile); - store_tile(a_copy_lds_window, - tile_elementwise_in(a_element_func, a_shuffle_tmp_loop)); + transpose_tile2d(a_shuffle_tmp_loop, elementwise_As_res); + store_tile(a_copy_lds_window, a_shuffle_tmp_loop); } else { - const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_block_tile); - store_tile(a_copy_lds_window, a_block_tile_tmp); + store_tile(a_copy_lds_window, elementwise_As_res); } // LDS write i + 1 @@ -260,14 +290,12 @@ struct GemmPipelineAGmemBGmemCRegV1 { auto b_shuffle_tmp_loop = make_static_distributed_tensor( Policy::template MakeShuffledBRegTileDistribution()); - transpose_tile2d(b_shuffle_tmp_loop, b_block_tile); - store_tile(b_copy_lds_window, - tile_elementwise_in(b_element_func, b_shuffle_tmp_loop)); + transpose_tile2d(b_shuffle_tmp_loop, elementwise_Bs_res); + store_tile(b_copy_lds_window, b_shuffle_tmp_loop); } else { - const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_block_tile); - store_tile(b_copy_lds_window, b_block_tile_tmp); + store_tile(b_copy_lds_window, elementwise_Bs_res); } iCounter--; @@ -284,20 +312,40 @@ struct GemmPipelineAGmemBGmemCRegV1 return c_block_tile; } - template - CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, index_t num_loop, void* p_smem) const { return operator()( a_dram_block_window_tmp, - [](const ADataType & a) { return a; }, + [](auto& e, const ADataType & a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType & b) { return b; }, + [](auto& e, const BDataType & b) { e = b; }, num_loop, p_smem); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2.hpp index b151cd6782..c309f8908a 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2.hpp @@ -15,30 +15,66 @@ namespace ck_tile { template struct GemmPipelineAGmemBGmemCRegV2 { - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; + static constexpr index_t APackedSize = ck_tile::numeric_traits>::PackedSize; static constexpr index_t BPackedSize = ck_tile::numeric_traits>::PackedSize; - static constexpr index_t kBlockSize = Problem::kBlockSize; + static constexpr index_t BlockSize = Problem::kBlockSize; static constexpr index_t kMPerBlock = BlockGemmShape::kM; static constexpr index_t kNPerBlock = BlockGemmShape::kN; static constexpr index_t kKPerBlock = BlockGemmShape::kK; + template + static constexpr index_t GetVectorSizeA() + { + return Problem::VectorSizeA; + } + template + static constexpr index_t GetVectorSizeB() + { + return Problem::VectorSizeB; + } + static constexpr index_t GetVectorSizeC() { return Problem::VectorSizeC; } + static constexpr index_t GetSmemPackA() { return Policy::template GetSmemPackA(); } static constexpr index_t GetSmemPackB() { return Policy::template GetSmemPackB(); } + static constexpr bool kPadM = Problem::kPadM; + static constexpr bool kPadN = Problem::kPadN; + static constexpr bool kPadK = Problem::kPadK; + + static constexpr bool Preshuffle = Problem::Preshuffle; + + static constexpr index_t NumWaveGroups = Problem::NumWaveGroups; + + // For the basic gemm pipelien DoubleSmemBuffer set to be false naturally. + static constexpr bool DoubleSmemBuffer = false; + [[nodiscard]] CK_TILE_HOST static const std::string GetName() { // clang-format off return concat('_', "pipeline_AGmemBGmemCRegV2", - concat('x', kMPerBlock, kNPerBlock, kKPerBlock, kBlockSize)); + concat('x', kMPerBlock, kNPerBlock, kKPerBlock, BlockSize)); // clang-format on } CK_TILE_HOST_DEVICE static constexpr auto TransposeC() { return Problem::TransposeC; } @@ -56,17 +92,31 @@ struct GemmPipelineAGmemBGmemCRegV2 BPackedSize; } - template (); + } + + template - CK_TILE_HOST_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + typename BElementFunction, + typename std::enable_if_t::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_HOST_DEVICE auto operator()(const AsDramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, - const BDramBlockWindowTmp& b_dram_block_window_tmp, + const BsDramBlockWindowTmp& b_dram_block_window_tmp, const BElementFunction& b_element_func, index_t num_loop, void* p_smem) const { + + using ADramBlockWindowTmp = + remove_cvref_t{}, AsDramBlockWindowTmp>>; + using BDramBlockWindowTmp = + remove_cvref_t{}, BsDramBlockWindowTmp>>; + static_assert( std::is_same_v> && std::is_same_v>, @@ -98,32 +148,40 @@ struct GemmPipelineAGmemBGmemCRegV2 auto b_lds_block = make_tensor_view(p_b_lds, b_lds_block_desc); // A DRAM tile window for load - auto a_copy_dram_window = - make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - a_dram_block_window_tmp.get_window_origin(), - Policy::template MakeADramTileDistribution()); + auto as_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + a_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + a_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeADramTileDistribution()); + }, + number{}); // A LDS tile window for store auto a_copy_lds_window = make_tile_window(a_lds_block, make_tuple(number{}, number{}), {0, 0}, - a_copy_dram_window.get_tile_distribution()); + as_copy_dram_window[number<0>{}].get_tile_distribution()); // B DRAM tile window for load - auto b_copy_dram_window = - make_tile_window(b_dram_block_window_tmp.get_bottom_tensor_view(), - make_tuple(number{}, number{}), - b_dram_block_window_tmp.get_window_origin(), - Policy::template MakeBDramTileDistribution()); + auto bs_copy_dram_window = generate_tuple( + [&](auto idx) { + return make_tile_window( + b_dram_block_window_tmp[number{}].get_bottom_tensor_view(), + make_tuple(number{}, number{}), + b_dram_block_window_tmp[number{}].get_window_origin(), + Policy::template MakeBDramTileDistribution()); + }, + number{}); // B LDS tile window for store auto b_copy_lds_window = make_tile_window(b_lds_block, make_tuple(number{}, number{}), {0, 0}, - b_copy_dram_window.get_tile_distribution()); + bs_copy_dram_window[number<0>{}].get_tile_distribution()); // Block GEMM constexpr auto block_gemm = Policy::template GetBlockGemm(); @@ -153,28 +211,30 @@ struct GemmPipelineAGmemBGmemCRegV2 // prefetch // global read 0 - auto a_block_tile = load_tile(a_copy_dram_window); - auto b_block_tile = load_tile(b_copy_dram_window); + // Load tile — during value loading, an elementwise function is executed for each A0, + // A1, … AN. The values A0, A1, … AN are read by the same thread. + auto elementwise_As_res = load_tile_with_elementwise(as_copy_dram_window, a_element_func); + // Load tile — during value loading, an elementwise function is executed for each B0, + // B1, … BN. The values B0, B1, … BN are read by the same thread. + auto elementwise_Bs_res = load_tile_with_elementwise(bs_copy_dram_window, b_element_func); { // move to 1 - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - move_tile_window(b_copy_dram_window, {0, kKPerBlock}); + move_tile_window(as_copy_dram_window, {0, kKPerBlock}); + move_tile_window(bs_copy_dram_window, {0, kKPerBlock}); // initialize C tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile); // LDS write 0 - const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_block_tile); - store_tile(a_copy_lds_window, a_block_tile_tmp); + store_tile(a_copy_lds_window, elementwise_As_res); // global read 1 - a_block_tile = load_tile(a_copy_dram_window); + elementwise_As_res = load_tile_with_elementwise(as_copy_dram_window, a_element_func); // LDS write 0 - const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_block_tile); - store_tile(b_copy_lds_window, b_block_tile_tmp); + store_tile(b_copy_lds_window, elementwise_Bs_res); // global read 1 - b_block_tile = load_tile(b_copy_dram_window); + elementwise_Bs_res = load_tile_with_elementwise(bs_copy_dram_window, b_element_func); } index_t iCounter = num_loop - 2; @@ -189,20 +249,18 @@ struct GemmPipelineAGmemBGmemCRegV2 block_sync_lds(); // move to i + 2 - move_tile_window(a_copy_dram_window, {0, kKPerBlock}); - move_tile_window(b_copy_dram_window, {0, kKPerBlock}); + move_tile_window(as_copy_dram_window, {0, kKPerBlock}); + move_tile_window(bs_copy_dram_window, {0, kKPerBlock}); // LDS write i + 1 - const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_block_tile); - store_tile(a_copy_lds_window, a_block_tile_tmp); + store_tile(a_copy_lds_window, elementwise_As_res); // global read i + 2 - a_block_tile = load_tile(a_copy_dram_window); + elementwise_As_res = load_tile_with_elementwise(as_copy_dram_window, a_element_func); // LDS write i + 1 - const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_block_tile); - store_tile(b_copy_lds_window, b_block_tile_tmp); + store_tile(b_copy_lds_window, elementwise_Bs_res); // global read i + 2 - b_block_tile = load_tile(b_copy_dram_window); + elementwise_Bs_res = load_tile_with_elementwise(bs_copy_dram_window, b_element_func); iCounter--; @@ -218,11 +276,9 @@ struct GemmPipelineAGmemBGmemCRegV2 block_sync_lds(); // LDS write num_loop - 1 - const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_block_tile); - store_tile(a_copy_lds_window, a_block_tile_tmp); + store_tile(a_copy_lds_window, elementwise_As_res); - const auto b_block_tile_tmp = tile_elementwise_in(b_element_func, b_block_tile); - store_tile(b_copy_lds_window, b_block_tile_tmp); + store_tile(b_copy_lds_window, elementwise_Bs_res); block_sync_lds(); @@ -241,12 +297,28 @@ struct GemmPipelineAGmemBGmemCRegV2 { return operator()( a_dram_block_window_tmp, - [](const ADataType & a) { return a; }, + [](auto& e, const ADataType & a) { e = a; }, b_dram_block_window_tmp, - [](const BDataType & b) { return b; }, + [](auto& e, const BDataType & b) { e = b; }, num_loop, p_smem); } + + template ::value && + !is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + const BDramBlockWindowTmp& b_dram_block_window_tmp, + index_t num_loop, + void* p_smem) const + { + return operator()(ck_tile::make_tuple(a_dram_block_window_tmp), + ck_tile::make_tuple(b_dram_block_window_tmp), + num_loop, + p_smem); + } }; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp index 52bd07c9e2..c73fa29245 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp @@ -5,16 +5,19 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" +#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" #include "ck_tile/host/concat.hpp" namespace ck_tile { -template @@ -22,18 +25,49 @@ struct GemmPipelineProblemBase { using Traits = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; // actually AccDataType - using ComputeDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; // actually AccDataType static constexpr bool FixedVectorSize = FixedVectorSize_; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; + + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + static constexpr bool ComputeDataTypeIsTuple = is_detected::value; + static constexpr bool ADataTypeIsTuple = is_detected::value; + static constexpr bool BDataTypeIsTuple = is_detected::value; + + static constexpr bool ALayoutIsTuple = is_detected::value; + static constexpr bool BLayoutIsTuple = is_detected::value; + + using ComputeDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + using AsLayoutTuple = std:: + conditional_t, remove_cvref_t>>; + using BsLayoutTuple = std:: + conditional_t, remove_cvref_t>>; + + using AsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using BsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using ComputeDataType = remove_cvref_t{}, ComputeDataTypeTuple>>; + using ADataType = remove_cvref_t{}, AsDataTypeTuple>>; + using ALayout = remove_cvref_t{}, AsLayoutTuple>>; + using BDataType = remove_cvref_t{}, BsDataTypeTuple>>; + using BLayout = remove_cvref_t{}, BsLayoutTuple>>; static constexpr bool TransposeC = Traits::TransposeC; static constexpr index_t NumWaveGroups = Traits::NumWaveGroups; @@ -66,7 +100,7 @@ struct GemmPipelineProblemBase { constexpr index_t PackedSize = ck_tile::numeric_traits>::PackedSize; - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) { constexpr index_t pixels_per_thread = BlockGemmShape::kM * BlockGemmShape::kK / kBlockSize; @@ -84,7 +118,7 @@ struct GemmPipelineProblemBase { constexpr index_t PackedSize = ck_tile::numeric_traits>::PackedSize; - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) { constexpr index_t pixels_per_thread = BlockGemmShape::kN * BlockGemmShape::kK / kBlockSize; @@ -125,7 +159,7 @@ struct GemmPipelineProblemBase { return VectorSizeA_; } - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) { return kPadK ? 1 : GetAlignmentA(); } @@ -140,7 +174,7 @@ struct GemmPipelineProblemBase { return VectorSizeB_; } - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) { return kPadN ? 1 : GetAlignmentB(); } @@ -161,35 +195,40 @@ struct GemmPipelineProblemBase }(); }; -// Alias for GemmPipelineProblem -template -using GemmPipelineProblem = GemmPipelineProblemBase; -template @@ -197,18 +236,48 @@ struct UniversalGemmPipelineProblem { using Traits = remove_cvref_t; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; // actually AccDataType - using ComputeDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; // actually AccDataType + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; static constexpr bool FixedVectorSize = FixedVectorSize_; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + static constexpr bool ComputeDataTypeIsTuple = is_detected::value; + static constexpr bool ADataTypeIsTuple = is_detected::value; + static constexpr bool BDataTypeIsTuple = is_detected::value; + + static constexpr bool ALayoutIsTuple = is_detected::value; + static constexpr bool BLayoutIsTuple = is_detected::value; + + using ComputeDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + using AsLayoutTuple = std:: + conditional_t, remove_cvref_t>>; + using BsLayoutTuple = std:: + conditional_t, remove_cvref_t>>; + + using AsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using BsDataTypeTuple = std::conditional_t, + remove_cvref_t>>; + + using ComputeDataType = remove_cvref_t{}, ComputeDataTypeTuple>>; + using ADataType = remove_cvref_t{}, AsDataTypeTuple>>; + using ALayout = remove_cvref_t{}, AsLayoutTuple>>; + using BDataType = remove_cvref_t{}, BsDataTypeTuple>>; + using BLayout = remove_cvref_t{}, BsLayoutTuple>>; static constexpr bool TransposeC = Traits::TransposeC; static constexpr index_t NumWaveGroups = Traits::NumWaveGroups; diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp index 8d47ab878e..c8f874acd6 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp @@ -356,11 +356,14 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeA() { - using ALayout = remove_cvref_t; - using ADataType = remove_cvref_t; + using AsLayout = remove_cvref_t; + using AsDataType = remove_cvref_t; constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + using ALayout = remove_cvref_t{}, AsLayout>>; + using ADataType = remove_cvref_t{}, AsDataType>>; + if constexpr(std::is_same_v) { return GetGlobalVectorLoadSize CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeB() { - using BLayout = remove_cvref_t; - using BDataType = remove_cvref_t; + using BsLayout = remove_cvref_t; + using BsDataType = remove_cvref_t; constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + using BLayout = remove_cvref_t{}, BsLayout>>; + using BDataType = remove_cvref_t{}, BsDataType>>; + if constexpr(std::is_same_v) { return GetGlobalVectorLoadSize CK_TILE_HOST_DEVICE static constexpr auto MakeADramTileDistribution() { - using ALayout = remove_cvref_t; - constexpr index_t BlockSize = Problem::kBlockSize; constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; @@ -491,6 +495,8 @@ struct UniversalGemmBasePolicy Problem::FixedVectorSize ? Problem::VectorSizeA : GetVectorSizeA(); constexpr index_t NumWaveGroups = Problem::NumWaveGroups; + using ALayout = remove_cvref_t< + std::tuple_element_t{}, remove_cvref_t>>; // Tile: MPerBlock X KPerBlock if constexpr(std::is_same_v) { @@ -518,8 +524,6 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeBDramTileDistribution() { - using BLayout = remove_cvref_t; - constexpr index_t BlockSize = Problem::kBlockSize; constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; @@ -527,6 +531,8 @@ struct UniversalGemmBasePolicy Problem::FixedVectorSize ? Problem::VectorSizeB : GetVectorSizeB(); constexpr index_t NumWaveGroups = Problem::NumWaveGroups; + using BLayout = remove_cvref_t< + std::tuple_element_t{}, remove_cvref_t>>; // Tile: KPerBlock X NPerBlock if constexpr(std::is_same_v) { @@ -554,7 +560,8 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledARegTileDistribution() { - using ALayout = remove_cvref_t; + using ALayout = remove_cvref_t< + std::tuple_element_t{}, remove_cvref_t>>; static_assert(std::is_same_v); constexpr index_t BlockSize = Problem::kBlockSize; constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; @@ -574,7 +581,8 @@ struct UniversalGemmBasePolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeShuffledBRegTileDistribution() { - using BLayout = remove_cvref_t; + using BLayout = remove_cvref_t< + std::tuple_element_t{}, remove_cvref_t>>; static_assert(std::is_same_v); constexpr index_t BlockSize = Problem::kBlockSize; constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; diff --git a/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp b/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp index 64900c9a97..96203b2cd2 100644 --- a/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp +++ b/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp @@ -10,8 +10,8 @@ namespace ck_tile { template struct TileGemmTraits @@ -23,9 +23,9 @@ struct TileGemmTraits // TODO this can't be hardcoded here! Should be in policy! static constexpr int _VectorSize = 16; - using ALayout = ALayout_; - using BLayout = BLayout_; - using CLayout = CLayout_; + using AsLayout = AsLayout_; + using BsLayout = BsLayout_; + using CLayout = CLayout_; static constexpr bool TransposeC = false; static constexpr bool UseStructuredSparsity = false; @@ -36,8 +36,8 @@ template @@ -76,8 +76,8 @@ using PersistentTileGemmUniversalTraits = TileGemmUniversalTraits; + using BTypeToUse = + std::conditional_t, + typename Problem::ADataType, + typename Problem::BDataType>; + using WarpGemm = WarpGemmDispatcher; using BlockWeightPreshufflePolicy = BlockWeightPreshuffleASmemBSmemCRegV1CustomPolicy { - using Base = BaseWeightPreshufflePipelineAGmemBGmemCRegV1; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using Base = BaseWeightPreshufflePipelineAGmemBGmemCRegV1; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; using BlockWeightPreshuffle = remove_cvref_t())>; @@ -188,7 +198,13 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 } } - template + template ::value && + !is_detected::value, + bool>* = nullptr, + index_t UnaryOpSize_ = 8> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, @@ -296,14 +312,14 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 NIterPerWarp> b_flat_dram_windows; - statically_indexed_array< - statically_indexed_array, - NIterPerWarp> + using BTypeToUse = + std::conditional_t, ADataType, BDataType>; + using BTileType = decltype(make_static_distributed_tensor(b_flat_distribution)); + + statically_indexed_array, NIterPerWarp> b_warp_tensor; - statically_indexed_array< - statically_indexed_array, - NIterPerWarp> + statically_indexed_array, NIterPerWarp> b_warp_tensor_2; static_for<0, NIterPerWarp, 1>{}([&](auto nIter) { @@ -313,7 +329,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -361,7 +378,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_2(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_2(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -394,7 +412,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -431,7 +450,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_2(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_2(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -455,7 +475,33 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 return c_block_tile; } - template + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + [[maybe_unused]] const AElementFunction& a_element_func, + const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, + [[maybe_unused]] const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem) const + { + return operator()( + a_dram_block_window_tmp[number<0>{}], + [](const ADataType & a) { return a; }, + b_flat_dram_block_window_tmp[number<0>{}], + num_loop, + p_smem); + } + + template ::value && + !is_detected::value, + bool>* = nullptr> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, index_t num_loop, @@ -463,7 +509,7 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV1 { return operator()( a_dram_block_window_tmp, - [](const ADataType & a) { return a; }, + [](auto& e, const ADataType & a) { e = a; }, b_flat_dram_block_window_tmp, num_loop, p_smem); diff --git a/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp b/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp index 129eac6557..670f4b0575 100644 --- a/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp +++ b/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp @@ -4,6 +4,7 @@ #pragma once #include "ck_tile/core.hpp" +#include "ck_tile/ops/common/load_interleaved_pk_type.hpp" #include "ck_tile/host/concat.hpp" #include "ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp" @@ -53,14 +54,23 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 { using Base = BaseWeightPreshufflePipelineAGmemBGmemCRegV2; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; + using AsDataType = remove_cvref_t; + using BsDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + + using AElementWise = remove_cvref_t; + using BElementWise = remove_cvref_t; using BlockGemmShape = remove_cvref_t; // TileFlatmmShape - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using CLayout = remove_cvref_t; + using AsLayout = remove_cvref_t; + using BsLayout = remove_cvref_t; + using CLayout = remove_cvref_t; + + using ALayout = remove_cvref_t>; + using BLayout = remove_cvref_t>; + + using ADataType = remove_cvref_t>; + using BDataType = remove_cvref_t>; using BlockWeightPreshuffle = remove_cvref_t())>; @@ -502,7 +512,11 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 template + typename AElementFunction, + typename std::enable_if_t::value && + !is_detected::value, + bool>* = nullptr, + index_t UnaryOpSize_ = 8> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const AElementFunction& a_element_func, const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, @@ -619,19 +633,19 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 b_flat_distribution); // pingpong buffer for B + using BTypeToUse = + std::conditional_t, ADataType, BDataType>; + using BTileType = decltype(make_static_distributed_tensor(b_flat_distribution)); + statically_indexed_array< statically_indexed_array, NIterPerWarp> b_flat_dram_windows; - statically_indexed_array< - statically_indexed_array, - NIterPerWarp> + statically_indexed_array, NIterPerWarp> b_warp_tensor_ping; - statically_indexed_array< - statically_indexed_array, - NIterPerWarp> + statically_indexed_array, NIterPerWarp> b_warp_tensor_pong; // Prefetch A0 @@ -647,7 +661,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_ping(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_ping(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); // move B window to next flat K @@ -694,7 +709,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_pong(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_pong(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -770,7 +786,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_ping(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_ping(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -850,7 +867,8 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 move_tile_window(b_flat_dram_windows(nIter)(kIter), {nIter * NFlatPerBlockPerIter, kIter * KFlatPerBlockPerIter}); - b_warp_tensor_pong(nIter)(kIter) = load_tile(b_flat_dram_windows(nIter)(kIter)); + load_int4_tile( + b_warp_tensor_pong(nIter)(kIter), b_flat_dram_windows(nIter)(kIter)); }); }); @@ -1001,8 +1019,37 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 return c_block_tile; } + // called from universal gemm kernel + template ::value && + is_detected::value, + bool>* = nullptr> + CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + [[maybe_unused]] const AElementFunction& a_element_func, + const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, + [[maybe_unused]] const BElementFunction& b_element_func, + index_t num_loop, + void* p_smem_ping, + void* p_smem_pong) const + { + return operator()( + a_dram_block_window_tmp[number<0>{}], + [](const ADataType& a) { return a; }, + b_flat_dram_block_window_tmp[number<0>{}], + num_loop, + p_smem_ping, + p_smem_pong); + } + // called from general gemm kernel - template + template ::value && + !is_detected::value, + bool>* = nullptr> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, index_t num_loop, @@ -1019,9 +1066,13 @@ struct WeightPreshufflePipelineAGmemBGmemCRegV2 } // called from grouped gemm kernel - template + template ::value && + !is_detected::value, + bool>* = nullptr> CK_TILE_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, - const BDramBlockWindowTmp& b_flat_dram_block_window_tmp, + const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp, index_t num_loop, TailNumber tail_number, void* __restrict__ p_smem_0, diff --git a/include/ck_tile/ops/gemm_group_quant.hpp b/include/ck_tile/ops/gemm_group_quant.hpp deleted file mode 100644 index 94b5ab8c3b..0000000000 --- a/include/ck_tile/ops/gemm_group_quant.hpp +++ /dev/null @@ -1,21 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include "ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp" -#include "ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp" -#include "ck_tile/ops/gemm_group_quant/kernel/gemm_quant_kernel.hpp" -#include "ck_tile/ops/gemm_group_quant/kernel/grouped_gemm_quant_kernel.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_group_quant_utils.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_quant_pipeline_problem.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp" -#include "ck_tile/ops/common/generic_2d_block_shape.hpp" -#include "ck_tile/ops/common/tensor_layout.hpp" -#include "ck_tile/ops/common/utils.hpp" diff --git a/include/ck_tile/ops/gemm_quant.hpp b/include/ck_tile/ops/gemm_quant.hpp new file mode 100644 index 0000000000..9f90050899 --- /dev/null +++ b/include/ck_tile/ops/gemm_quant.hpp @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp" +#include "ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp" +#include "ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp" +#include "ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_quant_pipeline_problem.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp" +#include "ck_tile/ops/common/generic_2d_block_shape.hpp" +#include "ck_tile/ops/common/tensor_layout.hpp" +#include "ck_tile/ops/common/utils.hpp" diff --git a/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp similarity index 96% rename from include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp rename to include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp index 182d9251b1..f75d02f1a6 100644 --- a/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp +++ b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp @@ -5,19 +5,19 @@ #include "ck_tile/core.hpp" #include "ck_tile/core/arch/arch.hpp" +#include "ck_tile/ops/common/load_interleaved_pk_type.hpp" #include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/elementwise.hpp" namespace ck_tile { -template +template struct BlockGemmAQuantBase { using AQDataType = remove_cvref_t; using ComputeDataType = remove_cvref_t; - static constexpr index_t UnaryOpSize = UnaryOpSize_; template CK_TILE_DEVICE static float cvt_scale_to_fp32(T scale) { @@ -42,23 +42,6 @@ struct BlockGemmAQuantBase } return scale_reg_f; } - - template - CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, - const WarpWindow& warp_window) - { - const element_wise::PassThroughPack8 elementwise_op{}; - - static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); - constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; - const auto in_dstr_tensors = load_tile(warp_window); - - using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); - static_for<0, thread_buffer_size, 1>{}([&](auto i) { - elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), - in_dstr_tensors.get_thread_buffer().template get_as()[i]); - }); - } }; // A is block window on shared memory @@ -66,7 +49,9 @@ struct BlockGemmAQuantBase // Consecutive kQuantGroupSize elements of A are quantized with a separate scale. // B is block window on shared memory // C is block distributed tensor -template +template struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase { private: @@ -172,6 +157,7 @@ struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase using Base = BlockGemmAQuantBase; + using Loader = remove_cvref_t>; using WarpGemm = remove_cvref_t; static constexpr index_t KIterPerWarp = Traits::KIterPerWarp; @@ -292,7 +278,7 @@ struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase { static_assert(std::is_same_v || std::is_same_v); - Base::load_interleaved_pk_type(a_warp_tile_, a_block_window); + Loader::load_interleaved_pk_type(a_warp_tile_, a_block_window); } else { @@ -302,7 +288,7 @@ struct AQuantBlockUniversalGemmAsBsCr : public BlockGemmAQuantBase { static_assert(std::is_same_v || std::is_same_v); - Base::load_interleaved_pk_type(b_warp_tile_, b_block_window); + Loader::load_interleaved_pk_type(b_warp_tile_, b_block_window); } else { diff --git a/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp similarity index 93% rename from include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp rename to include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp index 7e28ea8fa9..077d0d8fe2 100644 --- a/include/ck_tile/ops/gemm_group_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp +++ b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp @@ -5,19 +5,19 @@ #include "ck_tile/core.hpp" #include "ck_tile/core/arch/arch.hpp" +#include "ck_tile/ops/common/load_interleaved_pk_type.hpp" #include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_default_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" #include "ck_tile/ops/elementwise.hpp" namespace ck_tile { -template +template struct BlockGemmBQuantBase { using BQDataType = remove_cvref_t; using ComputeDataType = remove_cvref_t; - static constexpr index_t UnaryOpSize = UnaryOpSize_; template CK_TILE_DEVICE static float cvt_scale_to_fp32(T scale) { @@ -42,24 +42,6 @@ struct BlockGemmBQuantBase } return scale_reg_f; } - - // can be inherited from A - template - CK_TILE_DEVICE static void load_interleaved_pk_type(WarpTile& warp_tile, - const WarpWindow& warp_window) - { - const element_wise::PassThroughPack8 elementwise_op{}; - - static_assert(WarpTile::get_thread_buffer_size() % UnaryOpSize == 0); - constexpr index_t thread_buffer_size = WarpTile::get_thread_buffer_size() / UnaryOpSize; - const auto in_dstr_tensors = load_tile(warp_window); - - using ComputeVectorType = ComputeDataType __attribute__((ext_vector_type(UnaryOpSize))); - static_for<0, thread_buffer_size, 1>{}([&](auto i) { - elementwise_op(warp_tile.get_thread_buffer().template get_as()(i), - in_dstr_tensors.get_thread_buffer().template get_as()[i]); - }); - } }; // A is block window on shared memory @@ -67,7 +49,9 @@ struct BlockGemmBQuantBase // Consecutive kQuantGroupSize elements of B are quantized with a separate scale. // B is block window on shared memory // C is block distributed tensor -template +template struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase { private: @@ -170,6 +154,7 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase using Base = BlockGemmBQuantBase; + using Loader = remove_cvref_t>; using WarpGemm = remove_cvref_t; static constexpr index_t KIterPerWarp = Traits::KIterPerWarp; @@ -291,7 +276,7 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase { static_assert(std::is_same_v || std::is_same_v); - Base::load_interleaved_pk_type(a_warp_tile_, a_block_window); + Loader::load_interleaved_pk_type(a_warp_tile_, a_block_window); } else { @@ -301,7 +286,7 @@ struct BQuantBlockUniversalGemmAsBsCr : public BlockGemmBQuantBase { static_assert(std::is_same_v || std::is_same_v); - Base::load_interleaved_pk_type(b_warp_tile_, b_block_window); + Loader::load_interleaved_pk_type(b_warp_tile_, b_block_window); } else { diff --git a/include/ck_tile/ops/gemm_group_quant/kernel/gemm_quant_kernel.hpp b/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp similarity index 97% rename from include/ck_tile/ops/gemm_group_quant/kernel/gemm_quant_kernel.hpp rename to include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp index 13fa0b8dfb..82bf75a9e3 100644 --- a/include/ck_tile/ops/gemm_group_quant/kernel/gemm_quant_kernel.hpp +++ b/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp @@ -12,7 +12,7 @@ #include "ck_tile/core/numeric/integer.hpp" #include "ck_tile/core/numeric/math.hpp" #include "ck_tile/host/concat.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp" namespace ck_tile { @@ -330,7 +330,6 @@ struct QuantGemmKernel } } - // NOTE: no kernel currently uses BQuant like this: if constexpr(kQuantType == QuantType::BQuantGrouped) { static_assert(std::is_same_v); @@ -890,6 +889,7 @@ struct QuantGemmKernel * @param a_ptr input A pointer * @param b_ptr input B pointer * @param aq_ptr input AQ pointer + * @param bq_ptr input BQ pointer * @param c_ptr output C pointer * @param smem_ptr_0 The start memory pointer of the shared memory block. * @param kargs GEMM kernel arguments @@ -938,7 +938,8 @@ struct QuantGemmKernel return GemmPipeline{}.template operator()( a_block_window, b_block_window, bq_block_window, num_loop, smem_ptr_0); } - else if constexpr(kQuantType == QuantType::RowColQuant) + else if constexpr(kQuantType == QuantType::RowColQuant || + kQuantType == QuantType::TensorQuant) { return GemmPipeline{}.template operator()( a_block_window, b_block_window, num_loop, smem_ptr_0); @@ -964,6 +965,18 @@ struct QuantGemmKernel aq_block_window, bq_block_window); } + else if constexpr(kQuantType == QuantType::TensorQuant) + { + // TODO: why doesn't readfirstlane work here? + // const AccDataType aq_scale = + // __builtin_amdgcn_readfirstlane(type_convert(*aq_ptr)); + // const AccDataType bq_scale = + // __builtin_amdgcn_readfirstlane(type_convert(*bq_ptr)); + const AccDataType aq_scale = type_convert(*aq_ptr); + const AccDataType bq_scale = type_convert(*bq_ptr); + EpiloguePipeline{}( + c_block_window, c_block_tile, c_block_window, smem_ptr_0, aq_scale, bq_scale); + } } CK_TILE_DEVICE void operator()(QuantGemmKernelArgs kargs) const diff --git a/include/ck_tile/ops/gemm_group_quant/kernel/grouped_gemm_quant_kernel.hpp b/include/ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp similarity index 99% rename from include/ck_tile/ops/gemm_group_quant/kernel/grouped_gemm_quant_kernel.hpp rename to include/ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp index 925ea42678..07c45117e2 100644 --- a/include/ck_tile/ops/gemm_group_quant/kernel/grouped_gemm_quant_kernel.hpp +++ b/include/ck_tile/ops/gemm_quant/kernel/grouped_gemm_quant_kernel.hpp @@ -9,7 +9,7 @@ #include "ck_tile/host/stream_utils.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" -#include "ck_tile/ops/gemm_group_quant/kernel/gemm_quant_kernel.hpp" +#include "ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp" #include "ck_tile/host.hpp" #include diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp similarity index 99% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp index 5ce4268dca..24254013a4 100644 --- a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp @@ -9,7 +9,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/core/numeric/math.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp" #include "ck_tile/host/concat.hpp" namespace ck_tile { diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp similarity index 99% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp index 8f191f0f94..c27fbf5b50 100644 --- a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp @@ -9,7 +9,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp" #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp" -#include "ck_tile/ops/gemm_group_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp" +#include "ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_base.hpp" #include "ck_tile/host/concat.hpp" namespace ck_tile { diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_group_quant_utils.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp similarity index 100% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_group_quant_utils.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_quant_pipeline_problem.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_quant_pipeline_problem.hpp similarity index 87% rename from include/ck_tile/ops/gemm_group_quant/pipeline/gemm_quant_pipeline_problem.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/gemm_quant_pipeline_problem.hpp index a2cef2d994..d49204c64d 100644 --- a/include/ck_tile/ops/gemm_group_quant/pipeline/gemm_quant_pipeline_problem.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_quant_pipeline_problem.hpp @@ -168,17 +168,18 @@ template -using GemmRowColQuantPipelineProblem = GemmQuantPipelineProblemBase; +using GemmRowColTensorQuantPipelineProblem = + GemmQuantPipelineProblemBase; } // namespace ck_tile diff --git a/include/ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp b/include/ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp similarity index 74% rename from include/ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp rename to include/ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp index 44c6cd66c6..e97eeffb9b 100644 --- a/include/ck_tile/ops/gemm_group_quant/pipeline/tile_gemm_quant_traits.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp @@ -12,9 +12,22 @@ enum struct QuantType : std::uint16_t { AQuantGrouped = 0, BQuantGrouped = 1, - RowColQuant = 2 + RowColQuant = 2, + TensorQuant = 3 }; +std::string quant_type_to_string(QuantType quant_type) +{ + switch(quant_type) + { + case QuantType::AQuantGrouped: return "AQuantGrouped"; + case QuantType::BQuantGrouped: return "BQuantGrouped"; + case QuantType::RowColQuant: return "RowColQuant"; + case QuantType::TensorQuant: return "TensorQuant"; + default: return "Unknown"; + } +} + template ; + GroupedConvTraitsType_::ConvSpecialization, + GroupedConvTraitsType_::VectorSizeA, + GroupedConvTraitsType_::VectorSizeB, + GroupedConvTraitsType_::VectorSizeC>; static constexpr index_t NumDTensor = GroupedConvTraitsType_::NumDTensor; static constexpr auto I0 = number<0>(); @@ -468,6 +471,10 @@ template struct GroupedConvolutionBackwardDataKernel { + // Todo: Enable Vector Load Size > 1 + static_assert(GroupedConvTraitsType_::VectorSizeA == 1 && + GroupedConvTraitsType_::VectorSizeB == 1); + static constexpr index_t NDimSpatial = GroupedConvTraitsType_::NDimSpatial_; static constexpr ConvolutionSpecialization ConvSpecialization = GroupedConvTraitsType_::ConvSpecialization; @@ -509,10 +516,13 @@ struct GroupedConvolutionBackwardDataKernel static_assert(GemmPipeline::kPadM && GemmPipeline::kPadN && GemmPipeline::kPadK, "Not supported!"); - static_assert(std::is_same_v, - "Not supported A GEMM layout!"); - static_assert(std::is_same_v, - "Not supported B GEMM layout!"); + static_assert(std::is_same_v, "Not supported!"); + static_assert(std::is_same_v, "Not supported!"); + // TODO: Change to and enable vector load + // static_assert(std::is_same_v, + // "Not supported A GEMM layout!"); + // static_assert(std::is_same_v, + // "Not supported B GEMM layout!"); static_assert(std::is_same_v, "Not supported C GEMM layout!"); @@ -548,7 +558,7 @@ struct GroupedConvolutionBackwardDataKernel CK_TILE_HOST static bool IsSupportedArgument(const GroupedConvBwdDataKernelArgsSpecialized& kargs) { - if constexpr((EpiloguePipeline::GetVectorSizeC() % 2 != 0 && + if constexpr((GroupedConvTraitsType_::VectorSizeC % 2 != 0 && is_any_of::value) || !IsSplitKSupported) { @@ -625,7 +635,7 @@ struct GroupedConvolutionBackwardDataKernel std::is_same_v) { // Check access per C - if(ConvC % GemmPipeline::GetVectorSizeB() != 0) + if(ConvC % GroupedConvTraitsType_::VectorSizeB != 0) { CK_TILE_ERROR("Conv C is not a multiple of vector load size for input image!"); return false; @@ -637,13 +647,12 @@ struct GroupedConvolutionBackwardDataKernel return false; } - // check vector access of B // FIXME: layout if constexpr(std::is_same_v || std::is_same_v || std::is_same_v) { - if(ConvC % EpiloguePipeline::GetVectorSizeC() != 0) + if(ConvC % GroupedConvTraitsType_::VectorSizeC != 0) { CK_TILE_ERROR("Conv C is not a multiple of vector load size for weight!"); return false; @@ -655,12 +664,11 @@ struct GroupedConvolutionBackwardDataKernel return false; } - // check vector access of E if constexpr(std::is_same_v || std::is_same_v || std::is_same_v) { - if(ConvK % GemmPipeline::GetVectorSizeA() != 0) + if(ConvK % GroupedConvTraitsType_::VectorSizeA != 0) { CK_TILE_ERROR("Conv K is not a multiple of vector store size for output image!"); return false; @@ -957,7 +965,7 @@ struct GroupedConvolutionBackwardDataKernel { __shared__ char smem_ptr_1[GetSmemSize()]; if constexpr(!(EpiloguePipeline::MemoryOperation == memory_operation_enum::atomic_add && - EpiloguePipeline::GetVectorSizeC() % 2 != 0 && + GroupedConvTraitsType_::VectorSizeC % 2 != 0 && is_any_of::value)) { RunGemm2LDS(a_ptr, @@ -975,7 +983,7 @@ struct GroupedConvolutionBackwardDataKernel else { if constexpr(!(EpiloguePipeline::MemoryOperation == memory_operation_enum::atomic_add && - EpiloguePipeline::GetVectorSizeC() % 2 != 0 && + GroupedConvTraitsType_::VectorSizeC % 2 != 0 && is_any_of::value)) { RunGemm(a_ptr, b_ptr, kargs.ds_ptr, c_ptr, smem_ptr_0, kargs, i_m, i_n, group_id); diff --git a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp index 825aef3970..7bb3fedaf6 100644 --- a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp +++ b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp @@ -23,7 +23,10 @@ struct GroupedConvBwdWeightKernelArgs using ConvToGemmTransformer = TransformConvBwdWeightToGemm; + GroupedConvTraitsType_::ConvSpecialization, + GroupedConvTraitsType_::VectorSizeA, + GroupedConvTraitsType_::VectorSizeB, + GroupedConvTraitsType_::VectorSizeC>; static constexpr index_t NumDTensor = GroupedConvTraitsType_::NumDTensor; template < @@ -335,6 +338,10 @@ template struct GroupedConvolutionBackwardWeightKernel { + // Todo: Enable Vector Load Size > 1 + static_assert(GroupedConvTraitsType_::VectorSizeA == 1 && + GroupedConvTraitsType_::VectorSizeB == 1); + static constexpr index_t NDimSpatial = GroupedConvTraitsType_::NDimSpatial_; static constexpr ConvolutionSpecialization ConvSpecialization = GroupedConvTraitsType_::ConvSpecialization; @@ -355,11 +362,10 @@ struct GroupedConvolutionBackwardWeightKernel static constexpr index_t kBlockSize = GemmPipeline::BlockSize; - using InDataType = remove_cvref_t; - using WeiDataType = remove_cvref_t; + using OutDataType = remove_cvref_t; + using InDataType = remove_cvref_t; using DsDataType = remove_cvref_t; - // Below type is actually accumulation data type - the output of block GEMM. - using OutDataType = remove_cvref_t; + using WeiDataType = remove_cvref_t; using GroupedConvBwdWeightKernelArgsSpecialized = GroupedConvBwdWeightKernelArgs; @@ -376,6 +382,10 @@ struct GroupedConvolutionBackwardWeightKernel "Not supported!"); static_assert(std::is_same_v, "Not supported!"); static_assert(std::is_same_v, "Not supported!"); + // TODO: Change to and enable vector load + // static_assert(std::is_same_v, "Not + // supported!"); static_assert(std::is_same_v, "Not + // supported!"); static_assert(std::is_same_v, "Not supported!"); [[nodiscard]] CK_TILE_HOST static const std::string GetName() @@ -453,8 +463,8 @@ struct GroupedConvolutionBackwardWeightKernel CK_TILE_HOST static bool IsSupportedArgument(const GroupedConvBwdWeightKernelArgsSpecialized& kargs) { - if constexpr((EpiloguePipeline::GetVectorSizeC() % 2 != 0 && - is_any_of::value) || + if constexpr((GroupedConvTraitsType_::VectorSizeC % 2 != 0 && + is_any_of::value) || !IsSplitKSupported) { if(kargs.k_batch != 1) @@ -525,7 +535,7 @@ struct GroupedConvolutionBackwardWeightKernel std::is_same_v) { // Check access per C - if(ConvC % GemmPipeline::GetVectorSizeB() != 0) + if(ConvC % GroupedConvTraitsType_::VectorSizeB != 0) { CK_TILE_ERROR("Conv C is not a multiple of vector load size for input image!"); return false; @@ -537,13 +547,11 @@ struct GroupedConvolutionBackwardWeightKernel return false; } - // check vector access of B - // FIXME: layout if constexpr(std::is_same_v || std::is_same_v || std::is_same_v) { - if(ConvC % EpiloguePipeline::GetVectorSizeC() != 0) + if(ConvC % GroupedConvTraitsType_::VectorSizeC != 0) { CK_TILE_ERROR("Conv C is not a multiple of vector load size for weight!"); return false; @@ -555,12 +563,11 @@ struct GroupedConvolutionBackwardWeightKernel return false; } - // check vector access of E if constexpr(std::is_same_v || std::is_same_v || std::is_same_v) { - if(ConvK % GemmPipeline::GetVectorSizeA() != 0) + if(ConvK % GroupedConvTraitsType_::VectorSizeA != 0) { CK_TILE_ERROR("Conv K is not a multiple of vector store size for output image!"); return false; @@ -596,9 +603,8 @@ struct GroupedConvolutionBackwardWeightKernel }(); const auto& c_tensor_view = [&]() { - return make_tensor_view( - c_ptr, - kargs.c_grid_desc_m_n); // B: in + return make_tensor_view(c_ptr, + kargs.c_grid_desc_m_n); }(); const auto& ds_tensor_view = generate_tuple( @@ -607,11 +613,11 @@ struct GroupedConvolutionBackwardWeightKernel "Not supported!"); static_assert(std::is_same_v, "Not supported!"); - static_assert(std::is_same_v, OutDataType>, + static_assert(std::is_same_v, WeiDataType>, "Not supported!"); return make_tensor_view( - static_cast(ds_ptr[i]), kargs.c_grid_desc_m_n); + static_cast(ds_ptr[i]), kargs.c_grid_desc_m_n); }, number{}); @@ -829,8 +835,8 @@ struct GroupedConvolutionBackwardWeightKernel { __shared__ char smem_ptr_1[GetSmemSize()]; if constexpr(!(EpiloguePipeline::MemoryOperation == memory_operation_enum::atomic_add && - EpiloguePipeline::GetVectorSizeC() % 2 != 0 && - is_any_of::value)) + GroupedConvTraitsType_::VectorSizeC % 2 != 0 && + is_any_of::value)) { RunGemm2LDS(a_ptr, b_ptr, @@ -848,8 +854,8 @@ struct GroupedConvolutionBackwardWeightKernel else { if constexpr(!(EpiloguePipeline::MemoryOperation == memory_operation_enum::atomic_add && - EpiloguePipeline::GetVectorSizeC() % 2 != 0 && - is_any_of::value)) + GroupedConvTraitsType_::VectorSizeC % 2 != 0 && + is_any_of::value)) { RunGemm( a_ptr, b_ptr, kargs.ds_ptr, c_ptr, smem_ptr_0, kargs, num_loop, i_m, i_n, i_k); diff --git a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp index 6fcef5502e..d1eacd60cd 100644 --- a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp +++ b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp @@ -24,6 +24,9 @@ struct GroupedConvFwdKernelArgs using ConvToGemmFwdTransformer = TransformConvFwdToGemm; // Split N enabled static constexpr index_t NumDTensor = GroupedConvTraitsType_::NumDTensor; @@ -467,7 +470,7 @@ struct GroupedConvolutionForwardKernel CK_TILE_HOST static bool IsSupportedArgument(const GroupedConvFwdKernelArgsSpecialized& kargs) { - if constexpr((EpiloguePipeline::GetVectorSizeC() % 2 != 0 && + if constexpr((GroupedConvTraitsType_::VectorSizeC % 2 != 0 && is_any_of::value) || !IsSplitKSupported) { @@ -550,7 +553,7 @@ struct GroupedConvolutionForwardKernel std::is_same_v) { // Check access per C - if(ConvC % GemmPipeline::GetVectorSizeA() != 0) + if(ConvC % GroupedConvTraitsType_::VectorSizeA != 0) { CK_TILE_ERROR("Conv C is not a multiple of vector load size for input image!"); return false; @@ -568,7 +571,7 @@ struct GroupedConvolutionForwardKernel std::is_same_v || std::is_same_v) { - if(ConvC % GemmPipeline::GetVectorSizeB() != 0) + if(ConvC % GroupedConvTraitsType_::VectorSizeB != 0) { CK_TILE_ERROR("Conv C is not a multiple of vector load size for weight!"); return false; @@ -585,7 +588,7 @@ struct GroupedConvolutionForwardKernel std::is_same_v || std::is_same_v) { - if(ConvK % EpiloguePipeline::GetVectorSizeC() != 0) + if(ConvK % GroupedConvTraitsType_::VectorSizeC != 0) { CK_TILE_ERROR("Conv K is not a multiple of vector store size for output image!"); return false; @@ -858,7 +861,7 @@ struct GroupedConvolutionForwardKernel { __shared__ char smem_ptr_1[GetSmemSize()]; if constexpr(!(EpiloguePipeline::MemoryOperation == memory_operation_enum::atomic_add && - EpiloguePipeline::GetVectorSizeC() % 2 != 0 && + GroupedConvTraitsType_::VectorSizeC % 2 != 0 && is_any_of::value)) { RunGemm2LDS( @@ -868,7 +871,7 @@ struct GroupedConvolutionForwardKernel else { if constexpr(!(EpiloguePipeline::MemoryOperation == memory_operation_enum::atomic_add && - EpiloguePipeline::GetVectorSizeC() % 2 != 0 && + GroupedConvTraitsType_::VectorSizeC % 2 != 0 && is_any_of::value)) { RunGemm(a_ptr, b_ptr, kargs.ds_ptr, c_ptr, smem_ptr_0, kargs, i_m, i_n); diff --git a/include/ck_tile/ops/grouped_convolution/utils/grouped_convolution_utils.hpp b/include/ck_tile/ops/grouped_convolution/utils/grouped_convolution_utils.hpp index 3e5e87a975..c745aee622 100644 --- a/include/ck_tile/ops/grouped_convolution/utils/grouped_convolution_utils.hpp +++ b/include/ck_tile/ops/grouped_convolution/utils/grouped_convolution_utils.hpp @@ -49,7 +49,10 @@ template + typename OutLayout_, + index_t VectorSizeA_ = 1, + index_t VectorSizeB_ = 1, + index_t VectorSizeC_ = 1> struct GroupedConvTraits { private: @@ -67,14 +70,38 @@ struct GroupedConvTraits using WeiLayout = WeiLayout_; using DsLayout = DsLayout_; using OutLayout = OutLayout_; - using GroupedConvImplicitGemmTraits = TileGemmTraits; - static constexpr index_t NumDTensor = DsLayout::size(); - using ImplicitGemmDsLayout = decltype(generate_implicit_gemm_layout()); + using GroupedConvImplicitGemmTraitsFwd = + TileGemmTraits; + using GroupedConvImplicitGemmTraitsBwdData = + TileGemmTraits; + using GroupedConvImplicitGemmTraitsBwdWeight = + TileGemmTraits; + static constexpr ck_tile::index_t VectorSizeA = VectorSizeA_; + static constexpr ck_tile::index_t VectorSizeB = VectorSizeB_; + static constexpr ck_tile::index_t VectorSizeC = VectorSizeC_; + static constexpr index_t NumDTensor = DsLayout::size(); + using ImplicitGemmDsLayout = decltype(generate_implicit_gemm_layout()); }; } // namespace ck_tile diff --git a/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_data_to_gemm.hpp b/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_data_to_gemm.hpp index 972d05ff3e..c68a0a1400 100644 --- a/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_data_to_gemm.hpp +++ b/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_data_to_gemm.hpp @@ -10,6 +10,9 @@ namespace ck_tile { template 1 return make_naive_tensor_descriptor(make_tuple(N_, Wo_, K_), - make_tuple(NStride, WoStride, KStride)); + make_tuple(NStride, WoStride, KStride), + number{}, + I1); } template ::type = false> CK_TILE_HOST auto make_wei_grid_desc() const { // GKXC - return make_naive_tensor_descriptor_packed(make_tuple(K_, X_, C_)); + return make_naive_tensor_descriptor( + make_tuple(K_, X_, C_), make_tuple(X_ * C_, C_, I1), number{}, I1); } template ::type = false> @@ -462,7 +468,9 @@ struct TransformConvBwdDataToGemm // TODO Add support for NumGroupsToMerge > 1 return make_naive_tensor_descriptor(make_tuple(N_, Wi_, C_), - make_tuple(NStride, WiStride, CStride)); + make_tuple(NStride, WiStride, CStride), + number{}, + I1); } template ::type = false> @@ -477,7 +485,9 @@ struct TransformConvBwdDataToGemm // TODO Add support for NumGroupsToMerge > 1 return make_naive_tensor_descriptor(make_tuple(N_, Ho_, Wo_, K_), - make_tuple(NStride, HoStride, WoStride, KStride)); + make_tuple(NStride, HoStride, WoStride, KStride), + number{}, + I1); } template ::type = false> @@ -491,14 +501,19 @@ struct TransformConvBwdDataToGemm // TODO Add support for NumGroupsToMerge > 1 return make_naive_tensor_descriptor(make_tuple(N_, Hi_, Wi_, C_), - make_tuple(NStride, HiStride, WiStride, CStride)); + make_tuple(NStride, HiStride, WiStride, CStride), + number{}, + I1); } template ::type = false> CK_TILE_HOST auto make_wei_grid_desc() const { // GKYXC - return make_naive_tensor_descriptor_packed(make_tuple(K_, Y_, X_, C_)); + return make_naive_tensor_descriptor(make_tuple(K_, Y_, X_, C_), + make_tuple(C_ * X_ * Y_, C_ * X_, C_, I1), + number{}, + I1); } template ::type = false> @@ -514,7 +529,9 @@ struct TransformConvBwdDataToGemm // TODO Add support for NumGroupsToMerge > 1 return make_naive_tensor_descriptor( make_tuple(N_, Do_, Ho_, Wo_, K_), - make_tuple(NStride, DoStride, HoStride, WoStride, KStride)); + make_tuple(NStride, DoStride, HoStride, WoStride, KStride), + number{}, + I1); } template ::type = false> @@ -529,14 +546,20 @@ struct TransformConvBwdDataToGemm // TODO Add support for NumGroupsToMerge > 1 return make_naive_tensor_descriptor( make_tuple(N_, Di_, Hi_, Wi_, C_), - make_tuple(NStride, DiStride, HiStride, WiStride, CStride)); + make_tuple(NStride, DiStride, HiStride, WiStride, CStride), + number{}, + I1); } template ::type = false> CK_TILE_HOST auto make_wei_grid_desc() const { // GKZYXC - return make_naive_tensor_descriptor_packed(make_tuple(K_, Z_, Y_, X_, C_)); + return make_naive_tensor_descriptor( + make_tuple(K_, Z_, Y_, X_, C_), + make_tuple(C_ * X_ * Y_ * Z_, C_ * X_ * Y_, C_ * X_, C_, I1), + number{}, + I1); } // TODO: implement ck_tile::tensor_layout::convolution that describe packed/strided dimemsion as // properties diff --git a/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_weight_to_gemm.hpp b/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_weight_to_gemm.hpp index b2b7918810..010a8ac949 100644 --- a/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_weight_to_gemm.hpp +++ b/include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_weight_to_gemm.hpp @@ -10,6 +10,9 @@ namespace ck_tile { template 1 return make_naive_tensor_descriptor(make_tuple(K_, N_ * Wo_), - make_tuple(KStride, NDoHoWoStride)); + make_tuple(KStride, NDoHoWoStride), + number{}, + I1); } template ::type = false> @@ -433,7 +438,9 @@ struct TransformConvBwdWeightToGemm // TODO Add support for NumGroupsToMerge > 1 return make_naive_tensor_descriptor(make_tuple(N_, Wi_, C_), - make_tuple(NStride, WiStride, CStride)); + make_tuple(NStride, WiStride, CStride), + number{}, + I1); } template ::type = false> @@ -444,7 +451,8 @@ struct TransformConvBwdWeightToGemm constexpr auto CXStride = I1; // TODO Add support for NumGroupsToMerge > 1 - return make_naive_tensor_descriptor(make_tuple(K_, X_ * C_), make_tuple(KStride, CXStride)); + return make_naive_tensor_descriptor( + make_tuple(K_, X_ * C_), make_tuple(KStride, CXStride), number{}, I1); } template ::type = false> @@ -457,7 +465,9 @@ struct TransformConvBwdWeightToGemm // TODO Add support for NumGroupsToMerge > 1 return make_naive_tensor_descriptor(make_tuple(K_, N_ * Ho_ * Wo_), - make_tuple(KStride, NDoHoWoStride)); + make_tuple(KStride, NDoHoWoStride), + number{}, + I1); } template ::type = false> @@ -471,7 +481,9 @@ struct TransformConvBwdWeightToGemm // TODO Add support for NumGroupsToMerge > 1 return make_naive_tensor_descriptor(make_tuple(N_, Hi_, Wi_, C_), - make_tuple(NStride, HiStride, WiStride, CStride)); + make_tuple(NStride, HiStride, WiStride, CStride), + number{}, + I1); } template ::type = false> @@ -482,8 +494,8 @@ struct TransformConvBwdWeightToGemm constexpr auto CStride = I1; // TODO Add support for NumGroupsToMerge > 1 - return make_naive_tensor_descriptor(make_tuple(K_, Y_ * X_ * C_), - make_tuple(KStride, CStride)); + return make_naive_tensor_descriptor( + make_tuple(K_, Y_ * X_ * C_), make_tuple(KStride, CStride), number{}, I1); } template ::type = false> @@ -496,7 +508,9 @@ struct TransformConvBwdWeightToGemm // TODO Add support for NumGroupsToMerge > 1 return make_naive_tensor_descriptor(make_tuple(K_, N_ * Do_ * Ho_ * Wo_), - make_tuple(KStride, NDoHoWoStride)); + make_tuple(KStride, NDoHoWoStride), + number{}, + I1); } template ::type = false> @@ -511,7 +525,9 @@ struct TransformConvBwdWeightToGemm // TODO Add support for NumGroupsToMerge > 1 return make_naive_tensor_descriptor( make_tuple(N_, Di_, Hi_, Wi_, C_), - make_tuple(NStride, DiStride, HiStride, WiStride, CStride)); + make_tuple(NStride, DiStride, HiStride, WiStride, CStride), + number{}, + I1); } template ::type = false> @@ -523,7 +539,9 @@ struct TransformConvBwdWeightToGemm // TODO Add support for NumGroupsToMerge > 1 return make_naive_tensor_descriptor(make_tuple(K_, Z_ * Y_ * X_ * C_), - make_tuple(KStride, CStride)); + make_tuple(KStride, CStride), + number{}, + I1); } // TODO: implement ck_tile::tensor_layout::convolution that describe packed/strided dimemsion as diff --git a/include/ck_tile/ops/grouped_convolution/utils/transform_conv_fwd_to_gemm.hpp b/include/ck_tile/ops/grouped_convolution/utils/transform_conv_fwd_to_gemm.hpp index 2663d8a494..9e84973a89 100644 --- a/include/ck_tile/ops/grouped_convolution/utils/transform_conv_fwd_to_gemm.hpp +++ b/include/ck_tile/ops/grouped_convolution/utils/transform_conv_fwd_to_gemm.hpp @@ -10,6 +10,9 @@ namespace ck_tile { template {}, + I1); return transform_tensor_descriptor( in_gemmm_gemmk_desc, make_tuple(make_merge_transform(make_tuple(N_, Wo_)), @@ -458,7 +463,9 @@ struct TransformConvFwdToGemm { const auto in_gemmm_groups_gemmk_desc = make_naive_tensor_descriptor( make_tuple(N_, Wo_, NumGroupsToMerge, C_), - make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_, CStrideTensorA_)); + make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_, CStrideTensorA_), + number{}, + I1); return transform_tensor_descriptor( in_gemmm_groups_gemmk_desc, @@ -473,8 +480,11 @@ struct TransformConvFwdToGemm if constexpr(NumGroupsToMerge == 1) { - const auto in_n_wi_c_desc = make_naive_tensor_descriptor( - make_tuple(N_, Wi_), make_tuple(NStrideTensorA_, WiStride_)); + const auto in_n_wi_c_desc = + make_naive_tensor_descriptor(make_tuple(N_, Wi_), + make_tuple(NStrideTensorA_, WiStride_), + number{}, + I1); const auto in_n_wip_c_desc = transform_tensor_descriptor( in_n_wi_c_desc, @@ -502,7 +512,9 @@ struct TransformConvFwdToGemm { const auto in_n_wi_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Wi_, NumGroupsToMerge), - make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_)); + make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_), + number{}, + I1); const auto in_n_wip_c_desc = transform_tensor_descriptor( in_n_wi_c_desc, @@ -535,7 +547,9 @@ struct TransformConvFwdToGemm { const auto in_n_wi_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Wi_, C_), - make_tuple(NStrideTensorA_, WiStride_, CStrideTensorA_)); + make_tuple(NStrideTensorA_, WiStride_, CStrideTensorA_), + number{}, + I1); const auto in_n_wo_c_desc = transform_tensor_descriptor( in_n_wi_c_desc, @@ -556,7 +570,9 @@ struct TransformConvFwdToGemm { const auto in_n_wi_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Wi_, NumGroupsToMerge, C_), - make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_, CStrideTensorA_)); + make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_, CStrideTensorA_), + number{}, + I1); const auto in_n_wo_c_desc = transform_tensor_descriptor( in_n_wi_c_desc, @@ -581,7 +597,9 @@ struct TransformConvFwdToGemm { const auto in_n_wi_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Wi_, C_), - make_tuple(NStrideTensorA_, WiStride_, CStrideTensorA_)); + make_tuple(NStrideTensorA_, WiStride_, CStrideTensorA_), + number{}, + I1); const auto in_n_wip_c_desc = transform_tensor_descriptor( in_n_wi_c_desc, @@ -611,7 +629,9 @@ struct TransformConvFwdToGemm { const auto in_n_wi_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Wi_, NumGroupsToMerge, C_), - make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_, CStrideTensorA_)); + make_tuple(NStrideTensorA_, WiStride_, GStrideTensorA_, CStrideTensorA_), + number{}, + I1); const auto in_n_wip_c_desc = transform_tensor_descriptor( in_n_wi_c_desc, @@ -661,7 +681,9 @@ struct TransformConvFwdToGemm { const auto in_gemmm_gemmk_desc = make_naive_tensor_descriptor( make_tuple(N_, Ho_, Wo_, C_), - make_tuple(NStrideTensorA_, HiStride_, WiStride_, CStrideTensorA_)); + make_tuple(NStrideTensorA_, HiStride_, WiStride_, CStrideTensorA_), + number{}, + I1); return transform_tensor_descriptor( in_gemmm_gemmk_desc, @@ -675,7 +697,9 @@ struct TransformConvFwdToGemm const auto in_gemmm_groups_gemmk_desc = make_naive_tensor_descriptor( make_tuple(N_, Ho_, Wo_, NumGroupsToMerge, C_), make_tuple( - NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_, CStrideTensorA_)); + NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_, CStrideTensorA_), + number{}, + I1); return transform_tensor_descriptor( in_gemmm_groups_gemmk_desc, @@ -689,8 +713,11 @@ struct TransformConvFwdToGemm { if constexpr(NumGroupsToMerge == 1) { - const auto in_n_hi_wi_c_desc = make_naive_tensor_descriptor( - make_tuple(N_, Hi_, Wi_), make_tuple(NStrideTensorA_, HiStride_, WiStride_)); + const auto in_n_hi_wi_c_desc = + make_naive_tensor_descriptor(make_tuple(N_, Hi_, Wi_), + make_tuple(NStrideTensorA_, HiStride_, WiStride_), + number{}, + I1); const auto in_n_hip_wip_c_desc = transform_tensor_descriptor( in_n_hi_wi_c_desc, @@ -721,7 +748,9 @@ struct TransformConvFwdToGemm { const auto in_n_hi_wi_groups_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Hi_, Wi_, NumGroupsToMerge), - make_tuple(NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_)); + make_tuple(NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_), + number{}, + I1); const auto in_n_hip_wip_groups_c_desc = transform_tensor_descriptor( in_n_hi_wi_groups_c_desc, @@ -757,7 +786,9 @@ struct TransformConvFwdToGemm { const auto in_n_hi_wi_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Hi_, Wi_, C_), - make_tuple(NStrideTensorA_, HiStride_, WiStride_, CStrideTensorA_)); + make_tuple(NStrideTensorA_, HiStride_, WiStride_, CStrideTensorA_), + number{}, + I1); const auto in_n_ho_wo_c_desc = transform_tensor_descriptor( in_n_hi_wi_c_desc, @@ -780,7 +811,9 @@ struct TransformConvFwdToGemm const auto in_n_hi_wi_groups_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Hi_, Wi_, NumGroupsToMerge, C_), make_tuple( - NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_, CStrideTensorA_)); + NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_, CStrideTensorA_), + number{}, + I1); const auto in_n_ho_wo_groups_c_desc = transform_tensor_descriptor( in_n_hi_wi_groups_c_desc, @@ -808,7 +841,9 @@ struct TransformConvFwdToGemm { const auto in_n_hi_wi_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Hi_, Wi_, C_), - make_tuple(NStrideTensorA_, HiStride_, WiStride_, CStrideTensorA_)); + make_tuple(NStrideTensorA_, HiStride_, WiStride_, CStrideTensorA_), + number{}, + I1); const auto in_n_hip_wip_c_desc = transform_tensor_descriptor( in_n_hi_wi_c_desc, @@ -843,7 +878,9 @@ struct TransformConvFwdToGemm const auto in_n_hi_wi_groups_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Hi_, Wi_, NumGroupsToMerge, C_), make_tuple( - NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_, CStrideTensorA_)); + NStrideTensorA_, HiStride_, WiStride_, GStrideTensorA_, CStrideTensorA_), + number{}, + I1); const auto in_n_hip_wip_groups_c_desc = transform_tensor_descriptor( in_n_hi_wi_groups_c_desc, @@ -904,7 +941,9 @@ struct TransformConvFwdToGemm { const auto in_gemmm_gemmk_desc = make_naive_tensor_descriptor( make_tuple(N_, Do_, Ho_, Wo_, C_), - make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, CStrideTensorA_)); + make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, CStrideTensorA_), + number{}, + I1); return transform_tensor_descriptor( in_gemmm_gemmk_desc, @@ -922,7 +961,9 @@ struct TransformConvFwdToGemm HiStride_, WiStride_, GStrideTensorA_, - CStrideTensorA_)); + CStrideTensorA_), + number{}, + I1); return transform_tensor_descriptor( in_gemmm_groups_gemmk_desc, @@ -939,7 +980,9 @@ struct TransformConvFwdToGemm { const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Di_, Hi_, Wi_), - make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_)); + make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_), + number{}, + I1); const auto in_n_hip_wip_c_desc = transform_tensor_descriptor( in_n_di_hi_wi_c_desc, @@ -975,7 +1018,9 @@ struct TransformConvFwdToGemm { const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Di_, Hi_, Wi_, NumGroupsToMerge), - make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, GStrideTensorA_)); + make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, GStrideTensorA_), + number{}, + I1); const auto in_n_hip_wip_c_desc = transform_tensor_descriptor( in_n_di_hi_wi_c_desc, @@ -1022,7 +1067,9 @@ struct TransformConvFwdToGemm { const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Di_, Hi_, Wi_, C_), - make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, CStrideTensorA_)); + make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, CStrideTensorA_), + number{}, + I1); const auto in_n_do_ho_wo_c_desc = transform_tensor_descriptor( in_n_di_hi_wi_c_desc, @@ -1052,7 +1099,9 @@ struct TransformConvFwdToGemm HiStride_, WiStride_, GStrideTensorA_, - CStrideTensorA_)); + CStrideTensorA_), + number{}, + I1); const auto in_n_do_ho_wo_c_desc = transform_tensor_descriptor( in_n_di_hi_wi_c_desc, @@ -1090,7 +1139,9 @@ struct TransformConvFwdToGemm { const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor( make_tuple(N_, Di_, Hi_, Wi_, C_), - make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, CStrideTensorA_)); + make_tuple(NStrideTensorA_, DiStride_, HiStride_, WiStride_, CStrideTensorA_), + number{}, + I1); const auto in_n_hip_wip_c_desc = transform_tensor_descriptor( in_n_di_hi_wi_c_desc, @@ -1138,7 +1189,9 @@ struct TransformConvFwdToGemm HiStride_, WiStride_, GStrideTensorA_, - CStrideTensorA_)); + CStrideTensorA_), + number{}, + I1); const auto in_n_hip_wip_c_desc = transform_tensor_descriptor( in_n_di_hi_wi_c_desc, @@ -1217,14 +1270,19 @@ struct TransformConvFwdToGemm if constexpr(NumGroupsToMerge == 1) { - return make_naive_tensor_descriptor_packed(make_tuple(K_, FilterSizeNumType{})); + return make_naive_tensor_descriptor(make_tuple(K_, FilterSizeNumType{}), + make_tuple(FilterSizeNumType{}, I1), + number{}, + I1); } else { const auto wei_gemmn_groups_gemmk_desc = make_naive_tensor_descriptor( make_tuple(K_, NumGroupsToMerge, FilterSizeNumType{}), - make_tuple(KStrideTensorB_, GStrideTensorB_, CStrideTensorB_)); + make_tuple(KStrideTensorB_, GStrideTensorB_, CStrideTensorB_), + number{}, + I1); return transform_tensor_descriptor( wei_gemmn_groups_gemmk_desc, make_tuple(make_merge_transform(make_tuple(K_, NumGroupsToMerge)), @@ -1237,13 +1295,18 @@ struct TransformConvFwdToGemm { if constexpr(NumGroupsToMerge == 1) { - return make_naive_tensor_descriptor_packed(make_tuple(K_, ZYX_ * C_)); + return make_naive_tensor_descriptor(make_tuple(K_, ZYX_ * C_), + make_tuple(ZYX_ * C_, I1), + number{}, + I1); } else { const auto wei_gemmn_groups_gemmk_desc = make_naive_tensor_descriptor( make_tuple(K_, NumGroupsToMerge, ZYX_ * C_), - make_tuple(KStrideTensorB_, GStrideTensorB_, CStrideTensorB_)); + make_tuple(KStrideTensorB_, GStrideTensorB_, CStrideTensorB_), + number{}, + I1); return transform_tensor_descriptor( wei_gemmn_groups_gemmk_desc, make_tuple(make_merge_transform(make_tuple(K_, NumGroupsToMerge)), @@ -1270,14 +1333,18 @@ struct TransformConvFwdToGemm if constexpr(NumGroupsToMerge == 1) { return make_naive_tensor_descriptor(make_tuple(NDoHoWo, K_), - make_tuple(WoStride_, KStrideTensorC_)); + make_tuple(WoStride_, KStrideTensorC_), + number{}, + I1); } else { const auto nhwo_groups_k_1_desc = make_naive_tensor_descriptor( make_tuple(N_, Wo_, NumGroupsToMerge, K_, 1), make_tuple( - NStrideTensorC_, WoStride_, GStrideTensorC_, KStrideTensorC_, GStrideTensorC_)); + NStrideTensorC_, WoStride_, GStrideTensorC_, KStrideTensorC_, GStrideTensorC_), + number{}, + I1); // Padd 1 to NumGroupsToMerge const auto padded_desc = transform_tensor_descriptor( nhwo_groups_k_1_desc, @@ -1328,7 +1395,9 @@ struct TransformConvFwdToGemm if constexpr(NumGroupsToMerge == 1) { return make_naive_tensor_descriptor(make_tuple(NDoHoWo, K_), - make_tuple(WoStride_, KStrideTensorC_)); + make_tuple(WoStride_, KStrideTensorC_), + number{}, + I1); } else { @@ -1339,7 +1408,9 @@ struct TransformConvFwdToGemm WoStride_, GStrideTensorC_, KStrideTensorC_, - GStrideTensorC_)); + GStrideTensorC_), + number{}, + I1); // Padd 1 to NumGroupsToMerge const auto padded_desc = transform_tensor_descriptor( nhwo_groups_k_1_desc, @@ -1390,7 +1461,9 @@ struct TransformConvFwdToGemm if constexpr(NumGroupsToMerge == 1) { return make_naive_tensor_descriptor(make_tuple(NDoHoWo, K_), - make_tuple(WoStride_, KStrideTensorC_)); + make_tuple(WoStride_, KStrideTensorC_), + number{}, + I1); } else { @@ -1402,7 +1475,9 @@ struct TransformConvFwdToGemm WoStride_, GStrideTensorC_, KStrideTensorC_, - GStrideTensorC_)); + GStrideTensorC_), + number{}, + I1); // Padd 1 to NumGroupsToMerge const auto padded_desc = transform_tensor_descriptor( nhwo_groups_k_1_desc, diff --git a/script/launch_tests.sh b/script/launch_tests.sh index 5e71e25478..17a99e62a3 100755 --- a/script/launch_tests.sh +++ b/script/launch_tests.sh @@ -49,7 +49,7 @@ with open('$TEST_FILE', 'r') as f: if tests: # Extract just the filename after the last '/' clean_tests = [os.path.basename(test) for test in tests] - print('ctest -R \"' + '|'.join(clean_tests) + '\"') + print('ctest --output-on-failure -R \"' + '|'.join(clean_tests) + '\"') else: print('# No tests to run') ") @@ -57,5 +57,3 @@ with open('$TEST_FILE', 'r') as f: echo "$command" eval "$command" - - diff --git a/test/ck_tile/CMakeLists.txt b/test/ck_tile/CMakeLists.txt index 9314d4b795..b08f0d8316 100644 --- a/test/ck_tile/CMakeLists.txt +++ b/test/ck_tile/CMakeLists.txt @@ -5,6 +5,7 @@ add_subdirectory(batched_gemm) add_subdirectory(grouped_gemm) add_subdirectory(grouped_gemm_preshuffle) add_subdirectory(gemm_multi_d) +add_subdirectory(gemm_multi_abd) add_subdirectory(gemm_streamk) add_subdirectory(data_type) add_subdirectory(container) diff --git a/test/ck_tile/batched_gemm/test_batched_gemm_ut_cases.inc b/test/ck_tile/batched_gemm/test_batched_gemm_ut_cases.inc index 035377734b..8f24c9bfe1 100644 --- a/test/ck_tile/batched_gemm/test_batched_gemm_ut_cases.inc +++ b/test/ck_tile/batched_gemm/test_batched_gemm_ut_cases.inc @@ -29,7 +29,8 @@ TYPED_TEST(TestCkTileBatchedGemm, Basic) {256, 256, 64, 8}, {256, 256, 64, 16}}; - if(ck_tile::get_device_name() != "gfx950") { + if(ck_tile::get_device_name() != "gfx950") + { gemmParams.emplace_back(256, 256, 128, 2); } diff --git a/test/ck_tile/epilogue/test_cshuffle_epilogue.cpp b/test/ck_tile/epilogue/test_cshuffle_epilogue.cpp index 7b2071eb54..d2ae4c6adc 100644 --- a/test/ck_tile/epilogue/test_cshuffle_epilogue.cpp +++ b/test/ck_tile/epilogue/test_cshuffle_epilogue.cpp @@ -41,8 +41,8 @@ TEST_F(CShuffleEpilogueTest, BasicHalfTest) NPerXdl, KPerXdl>; - bool result = run_cshuffle_epilogue_test(); - EXPECT_TRUE(result) << "Basic CShuffleEpilogue test failed"; + auto result = run_cshuffle_epilogue_test(ScaleType::None); + EXPECT_FLOAT_EQ(result[0], 2.0F) << "Basic CShuffleEpilogue test failed"; } TEST_F(CShuffleEpilogueTest, BasicHalfTestWithScale) @@ -73,8 +73,45 @@ TEST_F(CShuffleEpilogueTest, BasicHalfTestWithScale) NPerXdl, KPerXdl>; - bool result = run_cshuffle_epilogue_test(true); - EXPECT_TRUE(result) << "Scale CShuffleEpilogue test failed"; + auto result = + run_cshuffle_epilogue_test(ScaleType::RowCol); + EXPECT_FLOAT_EQ(result[0], 2.0F) << "RowCol CShuffleEpilogue test failed: first element not 2"; + EXPECT_FLOAT_EQ(result[1], 4.0F) + << "RowCol CShuffleEpilogue test failed: second element not 2*2"; +} + +TEST_F(CShuffleEpilogueTest, BasicHalfTestWithTensorScale) +{ + // Basic test configuration with half_t data types + using ADataType = ck_tile::half_t; + using BDataType = ck_tile::half_t; + using AccDataType = float; + using ODataType = ck_tile::half_t; + + constexpr index_t kMPerBlock = 256; + constexpr index_t kNPerBlock = 256; + constexpr index_t MWave = 2; + constexpr index_t NWave = 2; + constexpr index_t MPerXdl = 32; + constexpr index_t NPerXdl = 32; + constexpr index_t KPerXdl = 8; + + using TestProblem = SimpleCShuffleEpilogueProblem; + + auto result = + run_cshuffle_epilogue_test(ScaleType::Tensor); + EXPECT_FLOAT_EQ(result[0], 4.0F) + << "TensorScale CShuffleEpilogue test failed: first element not 2*2=4"; } int main(int argc, char** argv) diff --git a/test/ck_tile/epilogue/test_cshuffle_epilogue_util.hpp b/test/ck_tile/epilogue/test_cshuffle_epilogue_util.hpp index c23957d802..01e6c91c7c 100644 --- a/test/ck_tile/epilogue/test_cshuffle_epilogue_util.hpp +++ b/test/ck_tile/epilogue/test_cshuffle_epilogue_util.hpp @@ -19,8 +19,15 @@ namespace ck_tile { +enum class ScaleType +{ + None, + RowCol, + Tensor +}; + // Simple test kernel to invoke the CShuffleEpilogue -template +template __global__ void test_cshuffle_epilogue_kernel(typename Problem::ODataType* __restrict__ output_data, float* m_scale, float* n_scale) @@ -61,7 +68,7 @@ __global__ void test_cshuffle_epilogue_kernel(typename Problem::ODataType* __res auto empty_ds = make_tuple(); // Call the epilogue - if constexpr(UseScale) + if constexpr(Scale == ScaleType::RowCol) { const auto m_scale_window = make_tile_window( make_naive_tensor_view( @@ -75,6 +82,10 @@ __global__ void test_cshuffle_epilogue_kernel(typename Problem::ODataType* __res {0, 0}); Epilogue{}(output_tile_window, acc_tile, empty_ds, smem, m_scale_window, n_scale_window); } + else if constexpr(Scale == ScaleType::Tensor) + { + Epilogue{}(output_tile_window, acc_tile, empty_ds, smem, *m_scale, *n_scale); + } else { Epilogue{}(output_tile_window, acc_tile, empty_ds, smem); @@ -113,7 +124,7 @@ using SimpleCShuffleEpilogueProblem = memory_operation_enum::set>; template -bool run_cshuffle_epilogue_test(bool use_scale = false) +auto run_cshuffle_epilogue_test(ScaleType scale = ScaleType::None) { using ODataType = typename Problem::ODataType; @@ -142,7 +153,7 @@ bool run_cshuffle_epilogue_test(bool use_scale = false) dim3 gridSize(1, 1, 1); dim3 blockSize(kBlockSize, 1, 1); - if(use_scale) + if(scale == ScaleType::RowCol) { float* m_scale; float* n_scale; @@ -155,12 +166,25 @@ bool run_cshuffle_epilogue_test(bool use_scale = false) hipMemcpy(m_scale, h_m_scale.data(), M * sizeof(float), hipMemcpyHostToDevice)); HIP_CHECK_ERROR( hipMemcpy(n_scale, h_n_scale.data(), N * sizeof(float), hipMemcpyHostToDevice)); - test_cshuffle_epilogue_kernel + test_cshuffle_epilogue_kernel + <<>>(device_output, m_scale, n_scale); + } + else if(scale == ScaleType::Tensor) + { + float* m_scale; + float* n_scale; + std::vector h_m_scale(1, 2.0F); + std::vector h_n_scale(1, 1.0F); + HIP_CHECK_ERROR(hipMalloc(&m_scale, sizeof(float))); + HIP_CHECK_ERROR(hipMalloc(&n_scale, sizeof(float))); + HIP_CHECK_ERROR(hipMemcpy(m_scale, h_m_scale.data(), sizeof(float), hipMemcpyHostToDevice)); + HIP_CHECK_ERROR(hipMemcpy(n_scale, h_n_scale.data(), sizeof(float), hipMemcpyHostToDevice)); + test_cshuffle_epilogue_kernel <<>>(device_output, m_scale, n_scale); } else { - test_cshuffle_epilogue_kernel + test_cshuffle_epilogue_kernel <<>>(device_output, nullptr, nullptr); } @@ -172,20 +196,10 @@ bool run_cshuffle_epilogue_test(bool use_scale = false) HIP_CHECK_ERROR(hipMemcpy( host_output.data(), device_output, output_size * sizeof(ODataType), hipMemcpyDeviceToHost)); - // Basic verification - just check that output has a 2, and 4 if using scaling - bool has_2 = - type_convert(host_output[0]) > 1.9F && type_convert(host_output[0]) < 2.1F; - bool scale_has_4 = true; - if(use_scale) - { - scale_has_4 = type_convert(host_output[1]) > 3.9F && - type_convert(host_output[1]) < 4.1F; - } - // Cleanup HIP_CHECK_ERROR(hipFree(device_output)); - return has_2 && scale_has_4; + return host_output; } } // namespace ck_tile diff --git a/test/ck_tile/fmha/test_fmha_fwd.inc b/test/ck_tile/fmha/test_fmha_fwd.inc index f02ef1e55e..08abd3358d 100644 --- a/test/ck_tile/fmha/test_fmha_fwd.inc +++ b/test/ck_tile/fmha/test_fmha_fwd.inc @@ -32,9 +32,6 @@ const ck_tile::stream_config stream_config{ 1, // rotating_count_ }; -// range_q, range_k, range_v, range_p, range_o, squant -#define QUANT_ARGS 1, 1, 1, 1, 1, squant - #define COMMON_ARGS \ init_method, static_cast(ck_tile::EnvValue(CK_TILE_ENV(CK_TILE_TEST_SEED))), 1, \ stream_config @@ -117,7 +114,7 @@ TEST_P(AllLong, Test) 1024, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -179,7 +176,7 @@ TEST_P(HDimPadding, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -236,7 +233,7 @@ TEST_P(ElementwiseBias, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -292,7 +289,7 @@ TEST_P(Alibi, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -350,7 +347,7 @@ TEST_P(Dropout, Test) drop_offset, // drop_offset drop_prefs, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -410,7 +407,7 @@ TEST_P(PagedKV, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -476,7 +473,7 @@ TEST_P(SplitKV, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, true, // is_rotary_interleaved num_splits, // num_splits COMMON_ARGS); @@ -548,7 +545,7 @@ TEST_P(AppendKV, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, false, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); @@ -618,7 +615,7 @@ TEST_P(AppendKVRoPE, Test) 0, // drop_offset false, // drop_prefs mask_str, // mask_str - QUANT_ARGS, + squant, is_rotary_interleaved, // is_rotary_interleaved 1, // num_splits COMMON_ARGS); diff --git a/test/ck_tile/fmha/test_fmha_fwd_fp8.cpp b/test/ck_tile/fmha/test_fmha_fwd_fp8.cpp index 46ed8f4125..b99c304d1f 100644 --- a/test/ck_tile/fmha/test_fmha_fwd_fp8.cpp +++ b/test/ck_tile/fmha/test_fmha_fwd_fp8.cpp @@ -17,22 +17,21 @@ using DataTypeConfig = FmhaFwdFp8; // instances are added), however the corresponding tests are not disabled (they will be skipped) // in case such instances will be added in the future. -const auto HDimValues = Values(std::tuple{64, -1}, std::tuple{128, -1}, std::tuple{256, -1}); +const auto HDimValues = Values(std::tuple{64, -1}, std::tuple{128, -1}); -const auto SplitKVHDimValues = Values(std::tuple{64, -1}, std::tuple{128, -1}, std::tuple{256, -1}); +const auto SplitKVHDimValues = Values(std::tuple{64, -1}, std::tuple{128, -1}); -const auto AppendKVHDimValues = - Values(std::tuple{64, -1}, std::tuple{128, -1}, std::tuple{256, -1}); +const auto AppendKVHDimValues = Values(std::tuple{64, -1}, std::tuple{128, -1}); // There are no fp8 instances with seqlen padding (mode_enum::group requires it) const auto ModeValues = Values(mode_enum::batch); const auto IsVRowmajorValues = Values(false); -const bool squant = true; -const std::string init_method = "ufq"; +const auto squant = true; +const std::string init_method = "uf"; const bool def_lse = false; -const bool def_is_v_rowmajor = false; +const bool def_is_v_rowmajor = true; int adjust_seqlen(int seqlen) { diff --git a/test/ck_tile/gemm/test_gemm_pipeline_smoke_run_test.inc b/test/ck_tile/gemm/test_gemm_pipeline_smoke_run_test.inc index ab74e4e7b1..57feefceab 100644 --- a/test/ck_tile/gemm/test_gemm_pipeline_smoke_run_test.inc +++ b/test/ck_tile/gemm/test_gemm_pipeline_smoke_run_test.inc @@ -2,6 +2,8 @@ // Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. #pragma once +#include "ck_tile/host/permute_pk_int4.hpp" + template static constexpr inline auto is_row_major(Layout layout_) { @@ -91,61 +93,6 @@ void permute_tensor_b(Tensor& tensor) } } -template -void permute_vectors_i4x4_b(Tensor& tensor) -{ - const ck_tile::index_t K = tensor.get_length(0); - const ck_tile::index_t N = tensor.get_length(1); - // vector pk_i4x4 permute - for(int i = 0; i < N; i++) - { - for(int j = 0; j < K; j += 8) - { - int8_t input[8]; - - for(int k = 0; k < 4; k++) - { - int8_t i4x2 = tensor(j + k * 2, i).data; - input[k * 2 + 0] = (i4x2 >> 4) & 0xf; - input[k * 2 + 1] = (i4x2 >> 0) & 0xf; - } - - // permute 01234567->20643175 - { - int8_t hi = input[2]; - int8_t lo = input[0]; - int8_t i4x2 = (hi << 4) | lo; - - tensor(j + 0, i) = i4x2; - } - - { - int8_t hi = input[6]; - int8_t lo = input[4]; - int8_t i4x2 = (hi << 4) | lo; - - tensor(j + 2, i) = i4x2; - } - - { - int8_t hi = input[3]; - int8_t lo = input[1]; - int8_t i4x2 = (hi << 4) | lo; - - tensor(j + 4, i) = i4x2; - } - - { - int8_t hi = input[7]; - int8_t lo = input[5]; - int8_t i4x2 = (hi << 4) | lo; - - tensor(j + 6, i) = i4x2; - } - } - } -} - template - -#include "ck_tile/core.hpp" -#include "ck_tile/host/kernel_launch.hpp" -#include "ck_tile/ops/epilogue.hpp" -#include "ck_tile/ops/gemm.hpp" -#include "ck_tile/ops/gemm_group_quant.hpp" - -#define CK_TILE_PIPELINE_PREFILL 1 -#define CK_TILE_PIPELINE_DECODE 2 -#define CK_TILE_PIPELINE_PRESHUFFLEQUANT 3 - -template -constexpr ck_tile::index_t get_k_warp_tile() -{ -#if defined(CK_GFX950_SUPPORT) - constexpr bool is_8bit_float = - std::is_same_v || std::is_same_v; - if constexpr(M_Warp_Tile == 32) - return is_8bit_float ? 64 : 16; - else - return is_8bit_float ? 128 : 32; -#else - if constexpr(M_Warp_Tile == 32) - return 16; - else - return 32; -#endif -} - -template -auto calculate_rtol_atol(const ck_tile::index_t K, - const ck_tile::index_t kbatch, - const float max_accumulated_value) -{ - using ComputeType = - std::conditional_t; - // Calculate thresholds - const auto rtol = ck_tile::get_relative_threshold( - ck_tile::integer_divide_ceil(K, kbatch)); - const auto atol = ck_tile::get_absolute_threshold( - max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); - // Calculate error due to split_k accumulation - const auto rtol_split_k = - ck_tile::get_relative_threshold(kbatch); - const auto atol_split_k = ck_tile::get_absolute_threshold( - max_accumulated_value, kbatch); - // Use higher threshold - return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); -} - -class ArgumentsNotSupportedException : public std::logic_error -{ - public: - explicit ArgumentsNotSupportedException(const std::string& message) : logic_error(message) {} -}; - -struct GemmConfigBase -{ - static constexpr bool kPadM = false; - static constexpr bool kPadN = false; - static constexpr bool kPadK = false; - - static constexpr bool PermuteA = false; - static constexpr bool PermuteB = false; - - static constexpr bool TransposeC = false; - static constexpr bool UseStructuredSparsity = false; - - static constexpr int kBlockPerCu = 1; - static constexpr ck_tile::index_t TileParitionerGroupNum = 8; - static constexpr ck_tile::index_t TileParitionerM01 = 4; - static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; - static constexpr ck_tile::index_t NumWaveGroups = 1; - static constexpr bool PreshuffleQuant = false; - static constexpr bool DoubleSmemBuffer = true; -}; - -template -struct GemmConfigDecode : public GemmConfigBase -{ - static constexpr ck_tile::index_t M_Tile = 16; - static constexpr ck_tile::index_t N_Tile = 64; - static constexpr ck_tile::index_t K_Tile = 256 / sizeof(PrecType); - - static constexpr ck_tile::index_t M_Warp = 1; - static constexpr ck_tile::index_t N_Warp = 4; - static constexpr ck_tile::index_t K_Warp = 1; - - static constexpr ck_tile::index_t M_Warp_Tile = 16; - static constexpr ck_tile::index_t N_Warp_Tile = 16; - static constexpr ck_tile::index_t K_Warp_Tile = get_k_warp_tile(); - - static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Default; - static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_DECODE; -}; - -template -struct GemmConfigPrefill : public GemmConfigBase -{ - static constexpr ck_tile::index_t M_Tile = 128; - static constexpr ck_tile::index_t N_Tile = 128; - static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType); - - static constexpr ck_tile::index_t M_Warp = 1; - static constexpr ck_tile::index_t N_Warp = 4; - static constexpr ck_tile::index_t K_Warp = 1; - - static constexpr ck_tile::index_t M_Warp_Tile = 16; - static constexpr ck_tile::index_t N_Warp_Tile = 16; - static constexpr ck_tile::index_t K_Warp_Tile = get_k_warp_tile(); - - static constexpr int kBlockPerCu = 2; - static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Default; - static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_PREFILL; -}; - -template -struct GemmConfigPreshuffleQuant : public GemmConfigBase -{ - static constexpr ck_tile::index_t M_Tile = 16; - static constexpr ck_tile::index_t N_Tile = 64; - static constexpr ck_tile::index_t K_Tile = 256 / sizeof(PrecType); - - static constexpr ck_tile::index_t M_Warp = 1; - static constexpr ck_tile::index_t N_Warp = 4; - static constexpr ck_tile::index_t K_Warp = 1; - - static constexpr ck_tile::index_t M_Warp_Tile = 16; - static constexpr ck_tile::index_t N_Warp_Tile = 16; - static constexpr ck_tile::index_t K_Warp_Tile = - get_k_from_preshuffled_warp_tile(); - - static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Default; - static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_PRESHUFFLEQUANT; - static constexpr bool PreshuffleQuant = true; -}; - -template -struct GemmQuantTypeConfig -{ - using ADataType = ADataType_; - using QDataType = QDataType_; - using BDataType = BDataType_; - using AccDataType = float; - using CDataType = CDataType_; -}; - -template -struct DataTypeTraits; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "fp32"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "fp64"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "int32"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "fp16"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "bf16"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "fp8"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "bf8"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "pk_int4_t"; -}; - -template <> -struct DataTypeTraits -{ - static constexpr const char* name = "int8"; -}; - -auto create_args(int argc, char* argv[]) -{ - ck_tile::ArgParser arg_parser; - arg_parser.insert("m", "3840", "m dimension") - .insert("n", "4096", "n dimension") - .insert("k", "2048", "k dimension") - .insert("a_layout", "R", "A tensor data layout - Row by default") - .insert("aq_layout", "R", "Aq tensor data layout - Row by default") - .insert("b_layout", "C", "B tensor data layout - Column by default") - .insert("c_layout", "R", "C tensor data layout - Row by default") - .insert("stride_a", "0", "Tensor A stride") - .insert("stride_q", "0", "Tensor AQ stride") - .insert("stride_b", "0", "Tensor B stride") - .insert("stride_c", "0", "Tensor C stride") - .insert("v", "2", "0. No validation, 1. Validation on CPU, 2. Validation on GPU") - .insert("prec", "i4fp8", "data type. fp8/bf8/i4fp8/i4bf8/i4f32fp8/i4f32bf8") - .insert("warmup", "50", "number of iterations before benchmark the kernel") - .insert("repeat", "100", "number of iterations to benchmark the kernel") - .insert("timer", "gpu", "gpu:gpu timer, cpu:cpu timer") - .insert("split_k", "1", "splitK value") - .insert("init", "0", "0:random, 1:linear, 2:constant(1)") - .insert("persistent", "0", "0:non-persistent, 1:persistent") - .insert("as_br_cr", "false", "Choose between as_br_cr and as_bs_cr"); - - bool result = arg_parser.parse(argc, argv); - return std::make_tuple(result, arg_parser); -} - -// host API -float gemm_calc_aquant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& s); diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp b/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp new file mode 100644 index 0000000000..ed3231d140 --- /dev/null +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp @@ -0,0 +1,179 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/host/check_err.hpp" +#include "ck_tile/host/reference/reference_gemm.hpp" +#include "ck_tile/ops/epilogue.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm_quant.hpp" + +// Forward declarations for quant type-specific implementations +template +struct QuantTypeTraits; + +// Base class for common quant gemm functionality +template +class TestCkTileGemmQuantBase : public ::testing::Test +{ + protected: + using ALayout = std::tuple_element_t<0, Tuple>; + using BLayout = std::tuple_element_t<1, Tuple>; + using CLayout = std::tuple_element_t<2, Tuple>; + using ADataType = std::tuple_element_t<3, Tuple>; + using BDataType = std::tuple_element_t<4, Tuple>; + using QDataType = std::tuple_element_t<5, Tuple>; + using CDataType = std::tuple_element_t<6, Tuple>; + static constexpr auto QuantType = std::tuple_element_t<7, Tuple>::value; + using GemmConfig = std::tuple_element_t<8, Tuple>; + static constexpr uint32_t QuantGroupSize = std::tuple_element_t<9, Tuple>::value; + using AccDataType = float; // accumulate always in float + + // Get the quant-type specific data types from traits + using QuantTraits = QuantTypeTraits; + using ComputeDataType = typename QuantTraits::template ComputeDataType; + + static constexpr ck_tile::index_t M_Tile = GemmConfig::M_Tile; + static constexpr ck_tile::index_t N_Tile = GemmConfig::N_Tile; + static constexpr ck_tile::index_t K_Tile = GemmConfig::K_Tile; + + static constexpr ck_tile::index_t M_Warp = GemmConfig::M_Warp; + static constexpr ck_tile::index_t N_Warp = GemmConfig::N_Warp; + static constexpr ck_tile::index_t K_Warp = GemmConfig::K_Warp; + + static constexpr ck_tile::index_t M_Warp_Tile = GemmConfig::M_Warp_Tile; + static constexpr ck_tile::index_t N_Warp_Tile = GemmConfig::N_Warp_Tile; + static constexpr ck_tile::index_t K_Warp_Tile = GemmConfig::K_Warp_Tile; + + public: + void SetUp() override { static_cast(this)->SetUpQuantTypeSpecific(); } + + void TearDown() override { static_cast(this)->TearDownQuantTypeSpecific(); } + + // Common test execution logic + void invoke_quant_gemm(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& s) + { + constexpr bool kPadM = false; + constexpr bool kPadN = false; + constexpr bool kPadK = false; + constexpr bool kPreshuffle = false; + + using CodegenGemmShape = + ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + + using TilePartitioner = ck_tile::GemmTile1DPartitioner; + + using CodegenGemmTraits = ck_tile::TileGemmQuantTraits; + + // Let the derived class create the appropriate pipeline and epilogue + static_cast(this) + ->template run_quant_gemm_impl( + args, s); + } + + void RunTest(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + { + // Generate test data and run the kernel + static_cast(this)->run_test_with_validation(M, N, K); + } + + // Helper function to check layout + template + static constexpr auto is_row_major(Layout) + { + return ck_tile::bool_constant, + ck_tile::tensor_layout::gemm::RowMajor>>{}; + } + + // Tolerance calculation function for validation + template + auto calculate_rtol_atol(const ck_tile::index_t K, + const ck_tile::index_t kbatch, + const float max_accumulated_value) + { + using ComputeType = + std::conditional_t; + // Calculate thresholds + const auto rtol = ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(K, kbatch)); + const auto atol = ck_tile::get_absolute_threshold( + max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); + // Calculate error due to split_k accumulation + const auto rtol_split_k = + ck_tile::get_relative_threshold(kbatch); + const auto atol_split_k = + ck_tile::get_absolute_threshold( + max_accumulated_value, kbatch); + // Use higher threshold + return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); + } +}; + +// Define generic QuantTypeTraits template (will be specialized) +template +struct QuantTypeTraits +{ + static_assert(QT == ck_tile::QuantType::AQuantGrouped || + QT == ck_tile::QuantType::BQuantGrouped || + QT == ck_tile::QuantType::RowColQuant || + QT == ck_tile::QuantType::TensorQuant, + "Unsupported quantization type"); +}; + +// Specialization for AQuantGrouped +template <> +struct QuantTypeTraits +{ + template + using ComputeDataType = BDataType; // For AQuant, compute type is BDataType + + static constexpr const char* name = "aquant"; +}; + +// Specialization for BQuantGrouped +template <> +struct QuantTypeTraits +{ + template + using ComputeDataType = ADataType; // For BQuant, compute type is ADataType + + static constexpr const char* name = "bquant"; +}; + +// Specialization for RowColQuant +template <> +struct QuantTypeTraits +{ + template + using ComputeDataType = ADataType; // For RowColQuant, compute type is ADataType + + static constexpr const char* name = "rowcol"; +}; + +// Specialization for TensorQuant +template <> +struct QuantTypeTraits +{ + template + using ComputeDataType = ADataType; // For TensorQuant, compute type is ADataType + + static constexpr const char* name = "tensor"; +}; diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp b/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp new file mode 100644 index 0000000000..5fc6b2f15c --- /dev/null +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_fixtures.hpp @@ -0,0 +1,919 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "test_gemm_quant_base.hpp" +#include "ck_tile/host/permute_pk_int4.hpp" + +struct GemmConfigBase +{ + static constexpr bool kPadM = false; + static constexpr bool kPadN = false; + static constexpr bool kPadK = false; + + static constexpr bool PermuteA = false; + static constexpr bool PermuteB = false; + + static constexpr bool TransposeC = false; + static constexpr bool UseStructuredSparsity = false; + + static constexpr int kBlockPerCu = 1; + static constexpr ck_tile::index_t TileParitionerGroupNum = 8; + static constexpr ck_tile::index_t TileParitionerM01 = 4; + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; + static constexpr ck_tile::index_t NumWaveGroups = 1; + static constexpr bool PreshuffleQuant = false; + static constexpr bool DoubleSmemBuffer = false; + + // Default GEMM tile sizes for tests + static constexpr ck_tile::index_t M_Tile = 16; + static constexpr ck_tile::index_t N_Tile = 64; + static constexpr ck_tile::index_t K_Tile = 256; + + static constexpr ck_tile::index_t M_Warp = 1; + static constexpr ck_tile::index_t N_Warp = 4; + static constexpr ck_tile::index_t K_Warp = 1; + + static constexpr ck_tile::index_t M_Warp_Tile = 16; + static constexpr ck_tile::index_t N_Warp_Tile = 16; + static constexpr ck_tile::index_t K_Warp_Tile = 32; +}; + +template +class TestCkTileGemmAQuant : public TestCkTileGemmQuantBase> +{ + using Base = TestCkTileGemmQuantBase>; + friend Base; + + public: + using typename Base::AccDataType; + using typename Base::ADataType; + using typename Base::ALayout; + using typename Base::BDataType; + using typename Base::BLayout; + using typename Base::CDataType; + using typename Base::CLayout; + using typename Base::ComputeDataType; + using typename Base::QDataType; + + static constexpr auto QuantType = Base::QuantType; + static constexpr uint32_t QuantGroupSize = Base::QuantGroupSize; + + protected: + void SetUpQuantTypeSpecific() {} + void TearDownQuantTypeSpecific() {} + + // AQuant-specific data generation + void run_test_with_validation(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + { + const ck_tile::index_t stride_A = K; + const ck_tile::index_t stride_B = K; + const ck_tile::index_t stride_C = M; + + // AQuant uses grouped quantization for A matrix + const ck_tile::index_t AQK = ck_tile::integer_divide_ceil(K, QuantGroupSize); + const ck_tile::index_t stride_AQ = + ck_tile::get_default_stride(M, AQK, 0, this->is_row_major(ALayout{})); + + // Generate test data + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(M, K, stride_A, this->is_row_major(ALayout{}))); + ck_tile::HostTensor aq_m_aqk( + ck_tile::host_tensor_descriptor(M, AQK, stride_AQ, this->is_row_major(ALayout{}))); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(K, N, stride_B, this->is_row_major(BLayout{}))); + + // Initialize data with random values + if constexpr(std::is_same_v) + { + ck_tile::FillUniformDistribution{-5.0f, 5.0f}(a_m_k); + } + else + { + ck_tile::FillUniformDistribution{-2.0f, 3.0f}(a_m_k); + } + ck_tile::FillUniformDistribution{-5.0f, 5.0f}(b_k_n); + ck_tile::FillUniformDistribution{-2.0f, 2.0f}(aq_m_aqk); + + // Allocate device memory + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size() * sizeof(ADataType)); + ck_tile::DeviceMem aq_m_aqk_dev_buf(aq_m_aqk.get_element_space_size() * sizeof(QDataType)); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size() * sizeof(BDataType)); + ck_tile::DeviceMem c_m_n_dev_buf(M * N * sizeof(CDataType)); + + // Copy to device + if constexpr(std::is_same_v) + { + // Permute vector pk_i4x4 data for device implementation + ck_tile::HostTensor temp = a_m_k; + ck_tile::permute_vectors_i4x4_b(temp); + a_m_k_dev_buf.ToDevice(temp.data()); + } + else + { + a_m_k_dev_buf.ToDevice(a_m_k.data()); + } + aq_m_aqk_dev_buf.ToDevice(aq_m_aqk.data()); + b_k_n_dev_buf.ToDevice(b_k_n.data()); + + // Create args for kernel execution + ck_tile::QuantGemmHostArgs args{ + a_m_k_dev_buf.GetDeviceBuffer(), // a_ptr + b_k_n_dev_buf.GetDeviceBuffer(), // b_ptr + c_m_n_dev_buf.GetDeviceBuffer(), // c_ptr + aq_m_aqk_dev_buf.GetDeviceBuffer(), // aq_ptr (scales) + nullptr, // bq_ptr (not used for AQuant) + 1, // k_batch + M, + N, + K, // M, N, K + AQK, // QK_A + 0, // QK_B (not used for AQuant) + stride_A, + stride_B, + stride_C, + stride_AQ, + 0 // strides + }; + + // Run the kernel + ck_tile::stream_config stream_config{}; + this->invoke_quant_gemm(args, stream_config); + + // Validation using reference implementation + ck_tile::HostTensor c_m_n_host_ref( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_host_ref.SetZero(); + + // Run reference AQuant implementation + ck_tile::reference_gemm_quant(a_m_k, aq_m_aqk, b_k_n, c_m_n_host_ref); + + // Get device result + ck_tile::HostTensor c_m_n_dev_result( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.mData.data()); + + // Calculate error tolerances + const float max_accumulated_value = + *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); + const auto rtol_atol = + this->template calculate_rtol_atol( + K, 1, max_accumulated_value); + + // Validate results + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + EXPECT_TRUE(pass) << "AQuantGrouped validation failed with M=" << M << ", N=" << N + << ", K=" << K; + + if(!pass) + { + std::cout << "AQuantGrouped - Relative error threshold: " + << rtol_atol.at(ck_tile::number<0>{}) + << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + } + } + + private: + // AQuant-specific pipeline implementation + template + void run_quant_gemm_impl(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s) + { + using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase; + + using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; + + const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + constexpr bool transpose_c = false; + + using PipelineProblem = + ck_tile::GemmAQuantPipelineProblem; + + using GemmPipeline = ck_tile::AQuantGemmPipelineAgBgCrCompV3; + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem, + AccDataType, + CDataType, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + Base::M_Warp, + Base::N_Warp, + Base::M_Warp_Tile, + Base::N_Warp_Tile, + Base::K_Warp_Tile, + transpose_c, + ck_tile::memory_operation_enum::set>>; + + using Kernel = ck_tile::QuantGemmKernel; + + auto kargs = Kernel::MakeKernelArgs(args); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Arguments not supported for AQuant kernel"); + } + + ck_tile::launch_kernel(s, + ck_tile::make_kernel( + Kernel{}, grids, blocks, 0, kargs)); + }; + + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + } +}; + +// BQuant-specific test fixture +template +class TestCkTileGemmBQuant : public TestCkTileGemmQuantBase> +{ + using Base = TestCkTileGemmQuantBase>; + friend Base; + + public: + using typename Base::AccDataType; + using typename Base::ADataType; + using typename Base::ALayout; + using typename Base::BDataType; + using typename Base::BLayout; + using typename Base::CDataType; + using typename Base::CLayout; + using typename Base::ComputeDataType; + using typename Base::QDataType; + + static constexpr auto QuantType = Base::QuantType; + static constexpr uint32_t QuantGroupSize = Base::QuantGroupSize; + + protected: + void SetUpQuantTypeSpecific() {} + void TearDownQuantTypeSpecific() {} + + void run_test_with_validation(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + { + const ck_tile::index_t stride_A = K; + const ck_tile::index_t stride_B = K; + const ck_tile::index_t stride_C = M; + + // BQuant uses grouped quantization for B matrix + const ck_tile::index_t BQK = ck_tile::integer_divide_ceil(K, QuantGroupSize); + const ck_tile::index_t stride_BQ = BQK; + + // Generate test data + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(M, K, stride_A, this->is_row_major(ALayout{}))); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(K, N, stride_B, this->is_row_major(BLayout{}))); + ck_tile::HostTensor bq_bqk_n( + ck_tile::host_tensor_descriptor(BQK, N, stride_BQ, this->is_row_major(BLayout{}))); + + // Initialize data with random values + ck_tile::FillUniformDistribution{-0.5f, 0.5f}(a_m_k); + ck_tile::FillUniformDistribution{0.f, 1.f}(b_k_n); + ck_tile::FillUniformDistribution{0.001f, 0.01f}(bq_bqk_n); + + // Allocate device memory + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size() * sizeof(ADataType)); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size() * sizeof(BDataType)); + ck_tile::DeviceMem bq_bqk_n_dev_buf(bq_bqk_n.get_element_space_size() * sizeof(QDataType)); + ck_tile::DeviceMem c_m_n_dev_buf(M * N * sizeof(CDataType)); + + // Copy to device + a_m_k_dev_buf.ToDevice(a_m_k.data()); + if constexpr(std::is_same_v) + { + // Permute vector pk_i4x4 data for device implementation + ck_tile::HostTensor temp = b_k_n; + ck_tile::permute_vectors_i4x4_b(temp); + b_k_n_dev_buf.ToDevice(temp.data()); + } + else + { + b_k_n_dev_buf.ToDevice(b_k_n.data()); + } + bq_bqk_n_dev_buf.ToDevice(bq_bqk_n.data()); + + // Create args for kernel execution + ck_tile::QuantGemmHostArgs args{ + a_m_k_dev_buf.GetDeviceBuffer(), // a_ptr + b_k_n_dev_buf.GetDeviceBuffer(), // b_ptr + c_m_n_dev_buf.GetDeviceBuffer(), // c_ptr + nullptr, // aq_ptr (not used for BQuant) + bq_bqk_n_dev_buf.GetDeviceBuffer(), // bq_ptr (scales) + 1, // k_batch + M, + N, + K, // M, N, K + 0, // QK_A (not used for BQuant) + BQK, // QK_B + stride_A, + stride_B, + stride_C, + 0, + stride_BQ // strides + }; + + // Run the kernel + ck_tile::stream_config stream_config{}; + this->invoke_quant_gemm(args, stream_config); + + // Validation using reference implementation + ck_tile::HostTensor c_m_n_host_ref( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_host_ref.SetZero(); + + // Run reference BQuant implementation + ck_tile::reference_gemm_quant(a_m_k, bq_bqk_n, b_k_n, c_m_n_host_ref); + + // Get device result + ck_tile::HostTensor c_m_n_dev_result( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.mData.data()); + + // Calculate error tolerances + const float max_accumulated_value = + *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); + const auto rtol_atol = + this->template calculate_rtol_atol( + K, 1, max_accumulated_value); + + // Validate results + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + EXPECT_TRUE(pass) << "BQuantGrouped validation failed with M=" << M << ", N=" << N + << ", K=" << K; + + if(!pass) + { + std::cout << "BQuantGrouped - Relative error threshold: " + << rtol_atol.at(ck_tile::number<0>{}) + << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + } + } + + private: + // BQuant-specific pipeline implementation + template + void run_quant_gemm_impl(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s) + { + using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase; + + using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; + + const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + + using PipelineProblem = + ck_tile::GemmBQuantPipelineProblem; + + using GemmPipeline = ck_tile::BQuantGemmPipelineAgBgCrCompV3; + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem, + AccDataType, + CDataType, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + Base::M_Warp, + Base::N_Warp, + Base::M_Warp_Tile, + Base::N_Warp_Tile, + Base::K_Warp_Tile, + false, // transpose_c + ck_tile::memory_operation_enum::set>>; + + using Kernel = ck_tile::QuantGemmKernel; + + auto kargs = Kernel::MakeKernelArgs(args); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Arguments not supported for BQuant kernel"); + } + + ck_tile::launch_kernel(s, + ck_tile::make_kernel( + Kernel{}, grids, blocks, 0, kargs)); + }; + + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + } +}; + +// RowColQuant-specific test fixture +template +class TestCkTileGemmRowColQuant + : public TestCkTileGemmQuantBase> +{ + using Base = TestCkTileGemmQuantBase>; + friend Base; + + public: + using typename Base::AccDataType; + using typename Base::ADataType; + using typename Base::ALayout; + using typename Base::BDataType; + using typename Base::BLayout; + using typename Base::CDataType; + using typename Base::CLayout; + using typename Base::ComputeDataType; + using typename Base::QDataType; + + static constexpr auto QuantType = Base::QuantType; + static constexpr uint32_t QuantGroupSize = Base::QuantGroupSize; + + protected: + void SetUpQuantTypeSpecific() {} + void TearDownQuantTypeSpecific() {} + + void run_test_with_validation(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + { + const ck_tile::index_t stride_A = K; + const ck_tile::index_t stride_B = K; + const ck_tile::index_t stride_C = M; + + // RowColQuant uses per-row and per-column scales + const ck_tile::index_t stride_row_scales = 1; + const ck_tile::index_t stride_col_scales = 1; + + // Generate test data + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(M, K, stride_A, this->is_row_major(ALayout{}))); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(K, N, stride_B, this->is_row_major(BLayout{}))); + ck_tile::HostTensor row_scales_m(ck_tile::host_tensor_descriptor( + M, 1, stride_row_scales, ck_tile::bool_constant{})); + ck_tile::HostTensor col_scales_n(ck_tile::host_tensor_descriptor( + N, 1, stride_col_scales, ck_tile::bool_constant{})); + + // Initialize data with random values + ck_tile::FillUniformDistribution{-0.5f, 0.5f}(a_m_k); + ck_tile::FillUniformDistribution{-0.5f, 0.5f}(b_k_n); + ck_tile::FillUniformDistribution{0.001f, 0.01f}(row_scales_m); + ck_tile::FillUniformDistribution{0.001f, 0.01f}(col_scales_n); + + // Allocate device memory + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size() * sizeof(ADataType)); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size() * sizeof(BDataType)); + ck_tile::DeviceMem row_scales_dev_buf(row_scales_m.get_element_space_size() * + sizeof(QDataType)); + ck_tile::DeviceMem col_scales_dev_buf(col_scales_n.get_element_space_size() * + sizeof(QDataType)); + ck_tile::DeviceMem c_m_n_dev_buf(M * N * sizeof(CDataType)); + + // Copy to device + a_m_k_dev_buf.ToDevice(a_m_k.data()); + b_k_n_dev_buf.ToDevice(b_k_n.data()); + row_scales_dev_buf.ToDevice(row_scales_m.data()); + col_scales_dev_buf.ToDevice(col_scales_n.data()); + + // Create args for kernel execution + ck_tile::QuantGemmHostArgs args{ + a_m_k_dev_buf.GetDeviceBuffer(), // a_ptr + b_k_n_dev_buf.GetDeviceBuffer(), // b_ptr + c_m_n_dev_buf.GetDeviceBuffer(), // c_ptr + row_scales_dev_buf.GetDeviceBuffer(), // aq_ptr (row scales) + col_scales_dev_buf.GetDeviceBuffer(), // bq_ptr (col scales) + 1, // k_batch + M, + N, + K, // M, N, K + 1, // QK_A (row scales) + 1, // QK_B (col scales) + stride_A, + stride_B, + stride_C, + stride_row_scales, + stride_col_scales // strides + }; + + // Run the kernel + ck_tile::stream_config stream_config{}; + this->invoke_quant_gemm(args, stream_config); + + // Validation using reference implementation + ck_tile::HostTensor c_m_n_host_ref( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_host_ref.SetZero(); + + // Run reference RowColQuant implementation + ck_tile::reference_gemm_rowcol_quant( + a_m_k, row_scales_m, b_k_n, col_scales_n, c_m_n_host_ref); + + // Get device result + ck_tile::HostTensor c_m_n_dev_result( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.mData.data()); + + // Calculate error tolerances + const float max_accumulated_value = + *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); + const auto rtol_atol = + this->template calculate_rtol_atol( + K, 1, max_accumulated_value); + + // Validate results + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + EXPECT_TRUE(pass) << "RowColQuant validation failed with M=" << M << ", N=" << N + << ", K=" << K; + + if(!pass) + { + std::cout << "RowColQuant - Relative error threshold: " + << rtol_atol.at(ck_tile::number<0>{}) + << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + } + } + + private: + // RowColQuant-specific pipeline implementation + template + void run_quant_gemm_impl(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s) + { + using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase; + + using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; + + const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + constexpr bool transpose_c = false; + + using PipelineProblem = ck_tile::GemmRowColTensorQuantPipelineProblem< + ADataType, + BDataType, + AccDataType, + AccDataType, + CodegenGemmShape, + CodegenGemmTraits, + transpose_c, + ComputeDataType, + ck_tile::GemmPipelineScheduler::Intrawave, + has_hot_loop_v, + tail_number_v>; + + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3; + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem, + AccDataType, + CDataType, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + Base::M_Warp, + Base::N_Warp, + Base::M_Warp_Tile, + Base::N_Warp_Tile, + Base::K_Warp_Tile, + transpose_c, + ck_tile::memory_operation_enum::set>>; + + using Kernel = ck_tile::QuantGemmKernel; + + auto kargs = Kernel::MakeKernelArgs(args); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Arguments not supported for RowColQuant kernel"); + } + + ck_tile::launch_kernel(s, + ck_tile::make_kernel( + Kernel{}, grids, blocks, 0, kargs)); + }; + + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + } +}; + +// TensorQuant-specific test fixture +template +class TestCkTileGemmTensorQuant + : public TestCkTileGemmQuantBase> +{ + using Base = TestCkTileGemmQuantBase>; + friend Base; + + public: + using typename Base::AccDataType; + using typename Base::ADataType; + using typename Base::ALayout; + using typename Base::BDataType; + using typename Base::BLayout; + using typename Base::CDataType; + using typename Base::CLayout; + using typename Base::ComputeDataType; + using typename Base::QDataType; + + static constexpr auto QuantType = Base::QuantType; + static constexpr uint32_t QuantGroupSize = Base::QuantGroupSize; + + protected: + void SetUpQuantTypeSpecific() {} + void TearDownQuantTypeSpecific() {} + + void run_test_with_validation(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + { + const ck_tile::index_t stride_A = K; + const ck_tile::index_t stride_B = K; + const ck_tile::index_t stride_C = M; + + // TensorQuant uses single scalar scale for each tensor + const ck_tile::index_t stride_scale_a = 1; + const ck_tile::index_t stride_scale_b = 1; + + // Generate test data + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(M, K, stride_A, this->is_row_major(ALayout{}))); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(K, N, stride_B, this->is_row_major(BLayout{}))); + ck_tile::HostTensor scale_a( + ck_tile::host_tensor_descriptor(1, 1, stride_scale_a, ck_tile::bool_constant{})); + ck_tile::HostTensor scale_b( + ck_tile::host_tensor_descriptor(1, 1, stride_scale_b, ck_tile::bool_constant{})); + + // Initialize data with random values + ck_tile::FillUniformDistribution{-0.5f, 0.5f}(a_m_k); + ck_tile::FillUniformDistribution{-0.5f, 0.5f}(b_k_n); + ck_tile::FillUniformDistribution{0.001f, 0.01f}(scale_a); + ck_tile::FillUniformDistribution{0.001f, 0.01f}(scale_b); + + // Allocate device memory + ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size() * sizeof(ADataType)); + ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size() * sizeof(BDataType)); + ck_tile::DeviceMem scale_a_dev_buf(scale_a.get_element_space_size() * sizeof(QDataType)); + ck_tile::DeviceMem scale_b_dev_buf(scale_b.get_element_space_size() * sizeof(QDataType)); + ck_tile::DeviceMem c_m_n_dev_buf(M * N * sizeof(CDataType)); + + // Copy to device + a_m_k_dev_buf.ToDevice(a_m_k.data()); + b_k_n_dev_buf.ToDevice(b_k_n.data()); + scale_a_dev_buf.ToDevice(scale_a.data()); + scale_b_dev_buf.ToDevice(scale_b.data()); + + // Create args for kernel execution + ck_tile::QuantGemmHostArgs args{ + a_m_k_dev_buf.GetDeviceBuffer(), // a_ptr + b_k_n_dev_buf.GetDeviceBuffer(), // b_ptr + c_m_n_dev_buf.GetDeviceBuffer(), // c_ptr + scale_a_dev_buf.GetDeviceBuffer(), // aq_ptr (scale A) + scale_b_dev_buf.GetDeviceBuffer(), // bq_ptr (scale B) + 1, // k_batch + M, + N, + K, // M, N, K + 1, // QK_A (tensor scale) + 1, // QK_B (tensor scale) + stride_A, + stride_B, + stride_C, + stride_scale_a, + stride_scale_b // strides + }; + + // Run the kernel + ck_tile::stream_config stream_config{}; + this->invoke_quant_gemm(args, stream_config); + + // Validation using reference implementation + ck_tile::HostTensor c_m_n_host_ref( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_host_ref.SetZero(); + + // Run reference TensorQuant implementation + ck_tile::reference_gemm_tensor_quant( + a_m_k, scale_a, b_k_n, scale_b, c_m_n_host_ref); + + // Get device result + ck_tile::HostTensor c_m_n_dev_result( + ck_tile::host_tensor_descriptor(M, N, stride_C, this->is_row_major(CLayout{}))); + c_m_n_dev_buf.FromDevice(c_m_n_dev_result.mData.data()); + + // Calculate error tolerances + const float max_accumulated_value = + *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); + const auto rtol_atol = + this->template calculate_rtol_atol( + K, 1, max_accumulated_value); + + // Validate results + bool pass = ck_tile::check_err(c_m_n_dev_result, + c_m_n_host_ref, + "Error: Incorrect results!", + rtol_atol.at(ck_tile::number<0>{}), + rtol_atol.at(ck_tile::number<1>{})); + + EXPECT_TRUE(pass) << "TensorQuant validation failed with M=" << M << ", N=" << N + << ", K=" << K; + + if(!pass) + { + std::cout << "TensorQuant - Relative error threshold: " + << rtol_atol.at(ck_tile::number<0>{}) + << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) + << std::endl; + } + } + + private: + // TensorQuant-specific pipeline implementation + template + void run_quant_gemm_impl(const ck_tile::QuantGemmHostArgs& args, + const ck_tile::stream_config& s) + { + using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase; + + using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3; + + const ck_tile::index_t K_split = (args.K + Base::K_Tile - 1) / Base::K_Tile * Base::K_Tile; + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); + + const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { + constexpr bool has_hot_loop_v = has_hot_loop_.value; + constexpr auto tail_number_v = tail_number_.value; + constexpr bool transpose_c = false; + + using PipelineProblem = ck_tile::GemmRowColTensorQuantPipelineProblem< + ADataType, + BDataType, + AccDataType, + AccDataType, + CodegenGemmShape, + CodegenGemmTraits, + transpose_c, + ComputeDataType, + ck_tile::GemmPipelineScheduler::Intrawave, + has_hot_loop_v, + tail_number_v>; + + using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3; + using GemmEpilogue = ck_tile::CShuffleEpilogue< + ck_tile::CShuffleEpilogueProblem, + AccDataType, + CDataType, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough, + TilePartitioner::MPerBlock, + TilePartitioner::NPerBlock, + Base::M_Warp, + Base::N_Warp, + Base::M_Warp_Tile, + Base::N_Warp_Tile, + Base::K_Warp_Tile, + transpose_c, + ck_tile::memory_operation_enum::set>>; + + using Kernel = ck_tile::QuantGemmKernel; + + auto kargs = Kernel::MakeKernelArgs(args); + const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = Kernel::BlockSize(); + + if(!Kernel::IsSupportedArgument(kargs)) + { + throw std::runtime_error("Arguments not supported for TensorQuant kernel"); + } + + ck_tile::launch_kernel(s, + ck_tile::make_kernel( + Kernel{}, grids, blocks, 0, kargs)); + }; + + return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); + } +}; diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_typed.cpp b/test/ck_tile/gemm_block_scale/test_gemm_quant_typed.cpp new file mode 100644 index 0000000000..1926b7cd0f --- /dev/null +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_typed.cpp @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "ck_tile/host.hpp" +#include "ck_tile/ops/gemm.hpp" + +#include +#include + +#include "test_gemm_quant_fixtures.hpp" + +// Type aliases for readability +using RowMajor = ck_tile::tensor_layout::gemm::RowMajor; +using ColumnMajor = ck_tile::tensor_layout::gemm::ColumnMajor; +using FP8 = ck_tile::fp8_t; +using BF8 = ck_tile::bf8_t; +using Half = ck_tile::half_t; +using PkInt4 = ck_tile::pk_int4_t; +using AQuantGrouped = std::integral_constant; +using BQuantGrouped = std::integral_constant; +using RowColQuant = std::integral_constant; +using TensorQuant = std::integral_constant; +using GroupSize = std::integral_constant; + +// Type combinations for each quantization type +// clang-format off +using AQuantTypes = ::testing::Types< + std::tuple, + std::tuple, + std::tuple, + std::tuple +>; +// clang-format on + +// clang-format off +using BQuantTypes = ::testing::Types< + std::tuple, + std::tuple, + std::tuple, + std::tuple +>; +// clang-format on + +// clang-format off +using RowColQuantTypes = ::testing::Types< + std::tuple, + std::tuple +>; +// clang-format on + +// clang-format off +using TensorQuantTypes = ::testing::Types< + std::tuple, + std::tuple +>; +// clang-format on + +// Test suites for each quantization type +TYPED_TEST_SUITE(TestCkTileGemmAQuant, AQuantTypes); +TYPED_TEST_SUITE(TestCkTileGemmBQuant, BQuantTypes); +TYPED_TEST_SUITE(TestCkTileGemmRowColQuant, RowColQuantTypes); +TYPED_TEST_SUITE(TestCkTileGemmTensorQuant, TensorQuantTypes); + +#include "test_gemm_quant_ut_cases.inc" diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_ut_cases.inc b/test/ck_tile/gemm_block_scale/test_gemm_quant_ut_cases.inc new file mode 100644 index 0000000000..9b07afa2b3 --- /dev/null +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_ut_cases.inc @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +// AQuant tests +TYPED_TEST(TestCkTileGemmAQuant, AQuantGroupedTest) +{ + this->run_test_with_validation(1024, 1024, 1024); +} + +// BQuant tests +TYPED_TEST(TestCkTileGemmBQuant, BQuantGroupedTest) +{ + this->run_test_with_validation(1024, 1024, 1024); +} + +// RowColQuant tests +TYPED_TEST(TestCkTileGemmRowColQuant, RowColQuantTest) +{ + this->run_test_with_validation(1024, 1024, 1024); +} + +// TensorQuant tests +TYPED_TEST(TestCkTileGemmTensorQuant, TensorQuantTest) +{ + this->run_test_with_validation(1024, 1024, 1024); +} diff --git a/test/ck_tile/gemm_block_scale/test_run_gemm_aquant_example.inc b/test/ck_tile/gemm_block_scale/test_run_gemm_aquant_example.inc deleted file mode 100644 index dbe652ac62..0000000000 --- a/test/ck_tile/gemm_block_scale/test_run_gemm_aquant_example.inc +++ /dev/null @@ -1,616 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include -#include -#include -#include -#include -#include -#include - -#include "ck_tile/core/config.hpp" -#include "ck_tile/host.hpp" -#include "test_gemm_aquant_utils.hpp" -#include "ck_tile/host/permute_pk_int4.hpp" - -template -float gemm_calc_aquant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::stream_config& s) -{ - constexpr bool kPadM = false; - constexpr bool kPadN = false; - constexpr bool kPadK = false; - - constexpr int kBlockPerCu = 1; - - static_assert(std::is_same_v); - - constexpr ck_tile::index_t M_Tile = GemmConfig::M_Tile; - constexpr ck_tile::index_t N_Tile = GemmConfig::N_Tile; - constexpr ck_tile::index_t K_Tile = GemmConfig::K_Tile; - - constexpr ck_tile::index_t M_Warp = GemmConfig::M_Warp; - constexpr ck_tile::index_t N_Warp = GemmConfig::N_Warp; - constexpr ck_tile::index_t K_Warp = GemmConfig::K_Warp; - - constexpr ck_tile::index_t M_Warp_Tile = GemmConfig::M_Warp_Tile; - constexpr ck_tile::index_t N_Warp_Tile = GemmConfig::N_Warp_Tile; - constexpr ck_tile::index_t K_Warp_Tile = GemmConfig::K_Warp_Tile; - - using CodegenGemmShape = - ck_tile::TileGemmShape, - ck_tile::sequence, - ck_tile::sequence>; - - using TilePartitioner = ck_tile::GemmTile1DPartitioner; - - using CodegenGemmTraits = ck_tile::TileGemmQuantTraits; - - using GemmPipelineProblem = ck_tile::GemmPipelineProblemBase; - - using BaseGemmPipeline = ck_tile::BaseAQuantGemmPipelineAgBgCrCompV3; - - const ck_tile::index_t K_split = (args.K + K_Tile - 1) / K_Tile * K_Tile; - const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); - const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); - const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); - constexpr bool transposed_warp_gemm = false; - - const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { - constexpr bool has_hot_loop_v = has_hot_loop_.value; - constexpr auto tail_number_v = tail_number_.value; - - using CodegenPipelineProblem = - ck_tile::GemmAQuantPipelineProblem; - using CodegenGemmPipeline = ck_tile::AQuantGemmPipelineAgBgCrCompV3; - using GemmEpilogue = ck_tile::CShuffleEpilogue< - ck_tile::CShuffleEpilogueProblem, - AccDataType, - CDataType, - ck_tile::tuple<>, - CLayout, - ck_tile::element_wise::PassThrough, - TilePartitioner::MPerBlock, - TilePartitioner::NPerBlock, - M_Warp, - N_Warp, - M_Warp_Tile, - N_Warp_Tile, - K_Warp_Tile, - transposed_warp_gemm, - ck_tile::memory_operation_enum::set>>; - using Kernel = ck_tile::QuantGemmKernel; - - auto kargs = Kernel::MakeKernelArgs(args); - - const dim3 grids = Kernel::GridSize(args.M, args.N, args.k_batch); - const dim3 blocks = Kernel::BlockSize(); - - if(args.k_batch != 1) - { - throw std::runtime_error("split-k is not supported yet!"); - } - - if(!Kernel::IsSupportedArgument(kargs)) - { - throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n"); - } - - if(s.log_level_ > 0) - { - std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n' - << "shape: " << CodegenGemmShape::GetName() << '\n' - << "problem: " << CodegenPipelineProblem::GetName() << '\n' - << "pipeline: " << CodegenGemmPipeline::GetName() << '\n' - << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}" - << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}" - << std::endl; - } - - float ave_time = ck_tile::launch_kernel( - s, ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); - - return ave_time; - }; - return BaseGemmPipeline::TailHandler(Run, has_hot_loop, tail_num); -} - -template -static constexpr inline auto is_row_major(Layout layout_) -{ - return ck_tile::bool_constant, - ck_tile::tensor_layout::gemm::RowMajor>>{}; -} - -template -float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, - ck_tile::DeviceMem& aq_m_aqk_dev_buf, - ck_tile::DeviceMem& b_k_n_dev_buf, - ck_tile::DeviceMem& c_m_n_dev_buf, - ck_tile::index_t M, - ck_tile::index_t N, - ck_tile::index_t K, - ck_tile::index_t AQK, - ck_tile::index_t stride_A, - ck_tile::index_t stride_AQ, - ck_tile::index_t stride_B, - ck_tile::index_t stride_C, - ck_tile::index_t kbatch, - int n_warmup, - int n_repeat) -{ - ck_tile::QuantGemmHostArgs args; - args.a_ptr = a_m_k_dev_buf.GetDeviceBuffer(); - args.aq_ptr = aq_m_aqk_dev_buf.GetDeviceBuffer(); - args.b_ptr = b_k_n_dev_buf.GetDeviceBuffer(); - args.c_ptr = c_m_n_dev_buf.GetDeviceBuffer(); - args.k_batch = kbatch; - args.M = M; - args.N = N; - args.K = K; - args.QK_A = AQK; - args.stride_A = stride_A; - args.stride_B = stride_B; - args.stride_C = stride_C; - args.stride_AQ = stride_AQ; - - float ave_time = gemm_calc_aquant( - args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat}); - - std::size_t flop = std::size_t(2) * M * N * K; - std::size_t num_byte = sizeof(ADataType) * M * K + sizeof(AQDataType) * M * AQK + - sizeof(BDataType) * N * K + sizeof(CDataType) * M * N; - float tflops = static_cast(flop) / 1.E9 / ave_time; - float gb_per_sec = num_byte / 1.E6 / ave_time; - - std::cout << "Run Gemm kernel with M =" << M << " N =" << N << " K =" << K - << " StrideA =" << stride_A << " StrideAQ =" << stride_AQ << " StrideB =" << stride_B - << " StrideC =" << stride_C << " A_Layout =" << ALayout::name - << " B_Layout =" << BLayout::name << " C_Layout =" << CLayout::name - << " A_Type = " << DataTypeTraits::name - << " AQ_Type = " << DataTypeTraits::name - << " B_Type = " << DataTypeTraits::name - << " Acc_Type = " << DataTypeTraits::name - << " C_Type = " << DataTypeTraits::name << " : " << ave_time << " ms, " - << tflops << " TFlops, " << gb_per_sec << " GB/s, " << std::endl; - - return ave_time; -} - -template -bool run_gemm_test_with_layouts(int argc, - char* argv[], - const ALayout a_layout = ALayout{}, - const AQLayout aq_layout = AQLayout{}, - const BLayout b_layout = BLayout{}, - [[maybe_unused]] const CLayout c_layout = CLayout{}) -{ - auto [result, arg_parser] = create_args(argc, argv); - if(!result) - return false; - - using ADataType = typename TypeConfig::ADataType; - using AQDataType = typename TypeConfig::QDataType; - using BDataType = typename TypeConfig::BDataType; - using AccDataType = typename TypeConfig::AccDataType; - using CDataType = typename TypeConfig::CDataType; - - ck_tile::index_t M = arg_parser.get_int("m"); - ck_tile::index_t N = arg_parser.get_int("n"); - ck_tile::index_t K = arg_parser.get_int("k"); - - if(K % QuantGroupSize != 0) - { - throw std::runtime_error("K must be aligned with QuantGroupSize"); - } - - ck_tile::index_t AQK = K / QuantGroupSize; - - ck_tile::index_t stride_A = arg_parser.get_int("stride_a"); - ck_tile::index_t stride_AQ = arg_parser.get_int("stride_q"); - ck_tile::index_t stride_B = arg_parser.get_int("stride_b"); - ck_tile::index_t stride_C = arg_parser.get_int("stride_c"); - - ck_tile::index_t kbatch = arg_parser.get_int("split_k"); - int n_warmup = arg_parser.get_int("warmup"); - int n_repeat = arg_parser.get_int("repeat"); - ck_tile::index_t init_method = arg_parser.get_int("init"); - - stride_A = ck_tile::get_default_stride(M, K, stride_A, is_row_major(a_layout)); - stride_AQ = ck_tile::get_default_stride(M, AQK, stride_AQ, is_row_major(aq_layout)); - stride_B = ck_tile::get_default_stride(K, N, stride_B, is_row_major(b_layout)); - stride_C = ck_tile::get_default_stride(M, N, stride_C, is_row_major(CLayout{})); - - ck_tile::HostTensor a_m_k( - ck_tile::host_tensor_descriptor(M, K, stride_A, is_row_major(a_layout))); - ck_tile::HostTensor aq_m_aqk( - ck_tile::host_tensor_descriptor(M, AQK, stride_AQ, is_row_major(aq_layout))); - ck_tile::HostTensor b_k_n( - ck_tile::host_tensor_descriptor(K, N, stride_B, is_row_major(b_layout))); - ck_tile::HostTensor c_m_n_dev_result( - ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{}))); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution fill_seed(0, 500); - - if(init_method == 0) - { - if constexpr(std::is_same_v) - { - ck_tile::FillUniformDistribution{-5.0f, 5.0f, fill_seed(gen)}( - a_m_k); - } - else - { - ck_tile::FillUniformDistribution{-2.0f, 3.0f, fill_seed(gen)}(a_m_k); - } - ck_tile::FillUniformDistribution{-2.0f, 2.0f, fill_seed(gen)}(aq_m_aqk); - ck_tile::FillUniformDistribution{-5.0f, 5.0f, fill_seed(gen)}(b_k_n); - } - else if(init_method == 1) - { - std::cout << "Monotonic initialization is not supported." << std::endl; - return true; - } - else if(init_method == 2) - { - ck_tile::FillConstant{static_cast(0x22)}(a_m_k); - ck_tile::FillConstant{static_cast(0.5f)}(aq_m_aqk); - ck_tile::FillConstant{static_cast(0x38)}(b_k_n); - } - else - { - a_m_k.SetZero(); - aq_m_aqk.SetZero(); - b_k_n.SetZero(); - } - - ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes()); - ck_tile::DeviceMem aq_m_aqk_dev_buf(aq_m_aqk.get_element_space_size_in_bytes()); - ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes()); - ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes()); - - if constexpr(std::is_same_v) - { - // Permute vector pk_i4x4 data for device implementation - ck_tile::HostTensor a_m_k_dev = a_m_k; - ck_tile::permute_vectors_i4x4_b(a_m_k_dev); - a_m_k_dev_buf.ToDevice(a_m_k_dev.data()); - } - else - { - a_m_k_dev_buf.ToDevice(a_m_k.data()); - } - aq_m_aqk_dev_buf.ToDevice(aq_m_aqk.data()); - b_k_n_dev_buf.ToDevice(b_k_n.data()); - c_m_n_dev_buf.SetZero(); - c_m_n_dev_result.SetZero(); - - invoke_gemm(a_m_k_dev_buf, - aq_m_aqk_dev_buf, - b_k_n_dev_buf, - c_m_n_dev_buf, - M, - N, - K, - AQK, - stride_A, - stride_AQ, - stride_B, - stride_C, - kbatch, - n_warmup, - n_repeat); - - c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); - bool pass = true; - - if(arg_parser.get_int("v") == 1) - { - ck_tile::HostTensor c_m_n_host_ref( - ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{}))); - c_m_n_host_ref.SetZero(); - - ck_tile::reference_gemm_quant(a_m_k, aq_m_aqk, b_k_n, c_m_n_host_ref); - const float max_accumulated_value = - *std::max_element(c_m_n_host_ref.mData.begin(), c_m_n_host_ref.mData.end()); - const auto rtol_atol = calculate_rtol_atol( - K, kbatch, max_accumulated_value); - pass = ck_tile::check_err(c_m_n_dev_result, - c_m_n_host_ref, - "Error: Incorrect results!", - rtol_atol.at(ck_tile::number<0>{}), - rtol_atol.at(ck_tile::number<1>{})); - - if(!pass) - { - std::cout << "Relative error threshold: " << rtol_atol.at(ck_tile::number<0>{}) - << " Absolute error threshold: " << rtol_atol.at(ck_tile::number<1>{}) - << std::endl; - } - std::cout << "CPU verification " << (pass ? "Passed!" : "Failed ...") << std::endl; - } - else if(arg_parser.get_int("v") == 2) - { - std::cout << "GPU verification is not implemented yet. Re-run with -v=1" << std::endl; - return false; - } - - return pass; -} - -template -bool run_gemm_test_prec_type(std::string a_layout, std::string b_layout, int argc, char* argv[]) -{ - using Row = ck_tile::tensor_layout::gemm::RowMajor; - using Col = ck_tile::tensor_layout::gemm::ColumnMajor; - - if constexpr(std::is_same_v || - std::is_same_v || - std::is_same_v) - { - if(a_layout == "R" && b_layout == "C") - { - return run_gemm_test_with_layouts( - argc, argv, Row{}, Row{}, Col{}, Row{}); - } - else - { - throw std::runtime_error("Unsupported memory layout for the input matrices!"); - } - } - else - { - throw std::runtime_error("Unsupported data type for A."); - } - - return true; -} - -template