mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-04-19 22:39:03 +00:00
Jimniu/ ck tile gemm stride validation (#2710)
* Add stride validation for gemm_basic * change default stride statement * Fix build error * Fix pre-commit failure * Addressed PR comments * clear the redundant code * clang format --------- Co-authored-by: mkumar16-amd <mkumar16@amd.com> Co-authored-by: ThomasNing <thomas.ning@amd.com>
This commit is contained in:
@@ -5,16 +5,32 @@
|
||||
#include "run_gemm_example.inc"
|
||||
#include "run_gemm_example_common.hpp"
|
||||
#include "gemm_basic_invoker.hpp"
|
||||
#include "ck_tile/core/utility/gemm_validation.hpp"
|
||||
|
||||
int run_gemm_example(ck_tile::ArgParser& arg_parser)
|
||||
{
|
||||
std::string data_type = arg_parser.get_str("prec");
|
||||
std::string a_layout = arg_parser.get_str("a_layout");
|
||||
std::string b_layout = arg_parser.get_str("b_layout");
|
||||
std::string c_layout = arg_parser.get_str("c_layout");
|
||||
|
||||
std::tuple<ck_tile::index_t, ck_tile::index_t, ck_tile::index_t> gemm_sizes =
|
||||
parse_gemm_size(arg_parser);
|
||||
|
||||
int m = std::get<0>(gemm_sizes);
|
||||
int n = std::get<1>(gemm_sizes);
|
||||
int k = std::get<2>(gemm_sizes);
|
||||
|
||||
int stride_a = arg_parser.get_int("stride_a");
|
||||
int stride_b = arg_parser.get_int("stride_b");
|
||||
int stride_c = arg_parser.get_int("stride_c");
|
||||
|
||||
using GemmConfig = GemmConfigBase;
|
||||
using Invoker = BasicInvoker;
|
||||
|
||||
ck_tile::validate_gemm_stride(
|
||||
a_layout, b_layout, c_layout, m, n, k, stride_a, stride_b, stride_c);
|
||||
|
||||
if(data_type == "fp16")
|
||||
{
|
||||
return run_gemm_example_prec_type<GemmConfig, Invoker, ck_tile::half_t>(
|
||||
|
||||
@@ -254,6 +254,15 @@ bool do_verify(const ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
|
||||
return pass;
|
||||
}
|
||||
|
||||
std::tuple<ck_tile::index_t, ck_tile::index_t, ck_tile::index_t>
|
||||
parse_gemm_size(ck_tile::ArgParser& arg_parser)
|
||||
{
|
||||
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");
|
||||
return std::make_tuple(M, N, K);
|
||||
}
|
||||
|
||||
template <typename GemmConfig,
|
||||
typename Invoker,
|
||||
typename ADataType,
|
||||
|
||||
@@ -71,6 +71,7 @@
|
||||
#include "ck_tile/core/utility/env.hpp"
|
||||
#include "ck_tile/core/utility/functional.hpp"
|
||||
#include "ck_tile/core/utility/functional_with_tuple.hpp"
|
||||
#include "ck_tile/core/utility/gemm_validation.hpp"
|
||||
#include "ck_tile/core/utility/ignore.hpp"
|
||||
#include "ck_tile/core/utility/literals.hpp"
|
||||
#include "ck_tile/core/utility/magic_div.hpp"
|
||||
|
||||
51
include/ck_tile/core/utility/gemm_validation.hpp
Normal file
51
include/ck_tile/core/utility/gemm_validation.hpp
Normal file
@@ -0,0 +1,51 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <stdexcept>
|
||||
#include "ck_tile/core/config.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
inline void
|
||||
validate_stride(std::string Layout, int M, int N, int stride, const std::string& stride_name)
|
||||
{
|
||||
if(Layout == "C" && stride < M)
|
||||
{
|
||||
throw std::runtime_error("For ColumnMajor layout, " + stride_name + "(" +
|
||||
std::to_string(stride) + ") must be greater or equal to dim " +
|
||||
std::to_string(M));
|
||||
}
|
||||
if(Layout == "R" && stride < N)
|
||||
{
|
||||
throw std::runtime_error("For RowMajor layout, " + stride_name + "(" +
|
||||
std::to_string(stride) + ") must be greater or equal to dim " +
|
||||
std::to_string(N));
|
||||
}
|
||||
}
|
||||
|
||||
inline void validate_gemm_stride(std::string a_layout,
|
||||
std::string b_layout,
|
||||
std::string c_layout,
|
||||
int M,
|
||||
int N,
|
||||
int K,
|
||||
int Stride_A,
|
||||
int Stride_B,
|
||||
int Stride_C)
|
||||
{
|
||||
// set default stride
|
||||
if(Stride_A <= 0)
|
||||
Stride_A = (a_layout == "R") ? K : M;
|
||||
if(Stride_B <= 0)
|
||||
Stride_B = (b_layout == "R") ? N : K;
|
||||
if(Stride_C <= 0)
|
||||
Stride_C = (c_layout == "R") ? N : M;
|
||||
|
||||
validate_stride(a_layout, M, K, Stride_A, "Stride_A");
|
||||
validate_stride(b_layout, K, N, Stride_B, "Stride_B");
|
||||
validate_stride(c_layout, M, N, Stride_C, "Stride_C");
|
||||
}
|
||||
} // namespace ck_tile
|
||||
Reference in New Issue
Block a user