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:
Jimniu
2025-10-29 22:45:09 -04:00
committed by GitHub
parent 1e77695fe8
commit 8c4cb4f9f4
4 changed files with 77 additions and 0 deletions

View File

@@ -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>(

View File

@@ -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,

View File

@@ -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"

View 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