Commit f0831350 authored by Jun Liu's avatar Jun Liu
Browse files

Merge branch 'amd-develop' into amd-master

parents f0fd0263 6e01019b
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
#include <iostream>
#include <numeric>
#include <initializer_list>
#include <cstdlib>
#include "profiler/profile_gemm_streamk_impl.hpp"
#include "profiler_operation_registry.hpp"
enum struct GemmMatrixLayout
{
MK_KN_MN, // 0
MK_NK_MN, // 1
KM_KN_MN, // 2
KM_NK_MN, // 3
};
enum struct GemmDataType
{
F32_F32_F32, // 0
F16_F16_F16, // 1
BF16_BF16_BF16, // 2
INT8_INT8_INT8, // 3
};
#define OP_NAME "gemm_streamk"
#define OP_DESC "StreamK GEMM"
int profile_gemm_streamk(int argc, char* argv[])
{
if(argc < 14)
{
printf("arg1: tensor operation (" OP_NAME ": " OP_DESC ")\n");
printf("arg2: data type (0: fp32; 1: fp16; 2: bf16; 3: int8)\n");
printf("arg3: matrix layout (0: A[m, k] * B[k, n] = C[m, n];\n");
printf(" 1: A[m, k] * B[n, k] = C[m, n];\n");
printf(" 2: A[k, m] * B[k, n] = C[m, n];\n");
printf(" 3: A[k, m] * B[n, k] = C[m, n])\n");
printf("arg4: verification (0: no; 1: yes)\n");
printf("arg5: initialization (0: no init; 1: integer value; 2: decimal value)\n");
printf("arg6: print tensor value (0: no; 1: yes)\n");
printf("arg7: time kernel (0=no, 1=yes)\n");
printf("arg8 to 13: M, N, K, StrideA, StrideB, StrideC\n");
printf("arg14: num_sk_blocks (optional)\n");
exit(1);
}
const auto data_type = static_cast<GemmDataType>(std::stoi(argv[2]));
const auto layout = static_cast<GemmMatrixLayout>(std::stoi(argv[3]));
const bool do_verification = std::stoi(argv[4]);
const int init_method = std::stoi(argv[5]);
const bool do_log = std::stoi(argv[6]);
const bool time_kernel = std::stoi(argv[7]);
const int M = std::stoi(argv[8]);
const int N = std::stoi(argv[9]);
const int K = std::stoi(argv[10]);
const int StrideA = std::stoi(argv[11]);
const int StrideB = std::stoi(argv[12]);
const int StrideC = std::stoi(argv[13]);
const uint32_t NumSKBlocks =
argc >= 15 ? static_cast<uint32_t>(std::stoul(std::string(argv[14]))) : 0xffffffff;
using F32 = float;
using F16 = ck::half_t;
using Row = ck::tensor_layout::gemm::RowMajor;
using Col = ck::tensor_layout::gemm::ColumnMajor;
auto profile = [&](auto a_type,
auto b_type,
auto acc_type,
auto c_type,
auto a_layout,
auto b_layout,
auto c_layout) {
using ADataType = decltype(a_type);
using BDataType = decltype(b_type);
using AccDataType = decltype(acc_type);
using CDataType = decltype(c_type);
using ALayout = decltype(a_layout);
using BLayout = decltype(b_layout);
using CLayout = decltype(c_layout);
const int DefaultStrideA = ck::is_same_v<ALayout, Row> ? K : M;
const int DefaultStrideB = ck::is_same_v<BLayout, Row> ? N : K;
const int DefaultStrideC = ck::is_same_v<CLayout, Row> ? N : M;
bool pass = ck::profiler::profile_gemm_streamk_impl<ADataType,
BDataType,
AccDataType,
CDataType,
ALayout,
BLayout,
CLayout>(
do_verification,
init_method,
do_log,
time_kernel,
M,
N,
K,
(StrideA <= 0) ? DefaultStrideA : StrideA,
(StrideB <= 0) ? DefaultStrideB : StrideB,
(StrideC <= 0) ? DefaultStrideC : StrideC,
NumSKBlocks);
return pass ? 0 : 1;
};
if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::MK_KN_MN)
{
return profile(F32{}, F32{}, F32{}, F32{}, Row{}, Row{}, Row{});
}
else if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::MK_NK_MN)
{
return profile(F32{}, F32{}, F32{}, F32{}, Row{}, Col{}, Row{});
}
else if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::KM_KN_MN)
{
return profile(F32{}, F32{}, F32{}, F32{}, Col{}, Row{}, Row{});
}
else if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::KM_NK_MN)
{
return profile(F32{}, F32{}, F32{}, F32{}, Col{}, Col{}, Row{});
}
else if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::MK_KN_MN)
{
return profile(F16{}, F16{}, F32{}, F16{}, Row{}, Row{}, Row{});
}
else if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::MK_NK_MN)
{
return profile(F16{}, F16{}, F32{}, F16{}, Row{}, Col{}, Row{});
}
else if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::KM_KN_MN)
{
return profile(F16{}, F16{}, F32{}, F16{}, Col{}, Row{}, Row{});
}
else if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::KM_NK_MN)
{
return profile(F16{}, F16{}, F32{}, F16{}, Col{}, Col{}, Row{});
}
else
{
std::cout << "this data_type & layout is not implemented" << std::endl;
return 1;
}
}
REGISTER_PROFILER_OPERATION(OP_NAME, OP_DESC, profile_gemm_streamk);
...@@ -3,13 +3,6 @@ ...@@ -3,13 +3,6 @@
## GPU visibility ## GPU visibility
export HIP_VISIBLE_DEVICES=0 export HIP_VISIBLE_DEVICES=0
DRIVER="../build/bin/ckProfiler" DRIVER="../build/bin/ckProfiler"
OP=$1
DATATYPE=$2
LAYOUT=$3
VERIFY=$4
INIT=$5
LOG=$6
TIME=$7
OP=$1 OP=$1
DATATYPE=$2 DATATYPE=$2
......
# TODO: Enable for gfx90a after complier fix # TODO: Enable for gfx90a after complier fix
if(NOT GPU_TARGETS MATCHES "gfx90a") if(DL_KERNELS)
add_gtest_executable(test_batched_gemm_multi_d test_batched_gemm_multi_d.cpp) add_gtest_executable(test_batched_gemm_multi_d test_batched_gemm_multi_d.cpp)
target_link_libraries(test_batched_gemm_multi_d PRIVATE utility device_batched_gemm_multi_d_instance) target_link_libraries(test_batched_gemm_multi_d PRIVATE utility device_batched_gemm_multi_d_instance)
endif() endif()
#include <stdio.h>
#include <string>
#include <algorithm>
#include <vector>
#include <limits>
#include "simple_args.h"
simple_args_t create_arg(int argc, char** argv)
{
simple_args_t args;
args.insert("m", "1024", "matrix m")
.insert("n", "1024", "matrix n")
.insert("k", "1024", "matrix k")
.insert("m_per_block", "128", "m_per_block")
.insert("n_per_block", "128", "n_per_block")
.insert("k_per_block", "32", "k_per_block")
.insert("num_cu", "104", "num cu")
.insert("occupancy", "2", "occupancy")
.parse(argc, argv);
return args;
}
namespace impl {
template <typename T>
T integer_divide_ceil(T n, T d)
{
return (n + d - 1) / d;
}
template <typename T>
T min(T a, T b)
{
return a > b ? b : a;
}
template <typename T>
T max(T a, T b)
{
return a > b ? a : b;
}
} // namespace impl
struct block_dispatcher_t
{
public:
uint32_t m_per_block;
uint32_t n_per_block;
uint32_t k_per_block;
uint32_t num_cu;
uint32_t occupancy;
uint32_t m;
uint32_t n;
uint32_t k;
//--------------------------------------
uint32_t sk_num_blocks;
uint32_t sk_num_big_blocks;
uint32_t sk_total_iters;
// uint32_t sk_num_blocks_per_tile; // how many
uint32_t dp_start_block_idx;
uint32_t dp_iters_per_block;
uint32_t dp_num_blocks;
uint32_t k_iters_per_tile;
uint32_t k_iters_per_big_block;
//--------------------------------------
static constexpr uint32_t min_k_iters_per_sk_block = 1;
void dump()
{
printf("%dx%dx%d(%dx%dx%d), cu:%d, occ:%d, grids:%d, sk_num_big_blocks:%d, "
"sk_num_blocks:%d, sk_total_iters:%d, dp_start_block_idx:%d, dp_iters_per_block:%d, "
"dp_num_blocks:%d, k_iters_per_tile:%d, k_iters_per_big_block:%d\n",
m,
n,
k,
m_per_block,
n_per_block,
k_per_block,
num_cu,
occupancy,
get_grid_dims_x(),
sk_num_big_blocks,
sk_num_blocks,
sk_total_iters,
dp_start_block_idx,
dp_iters_per_block,
dp_num_blocks,
k_iters_per_tile,
k_iters_per_big_block);
}
block_dispatcher_t(uint32_t m_per_block_,
uint32_t n_per_block_,
uint32_t k_per_block_,
uint32_t num_cu_,
uint32_t occupancy_,
uint32_t m_,
uint32_t n_,
uint32_t k_)
: m_per_block(m_per_block_),
n_per_block(n_per_block_),
k_per_block(k_per_block_),
num_cu(num_cu_),
occupancy(occupancy_),
m(m_),
n(n_),
k(k_)
{
init();
}
uint32_t get_grid_dims_x() { return dp_start_block_idx + dp_num_blocks; }
uint32_t get_block_idx(uint32_t bid)
{
// block id is linearily allocated along sk blocks (dp blocks are fine)
// this function will compute blockIdx.x and the linear sk block mapping
// uint32_t block_idx = 0;
// if(bid < sk_num_big_blocks) {
// uint32_t current_k_iter = bid * k_iters_per_big_block;
// tile_idx = current_k_iter / k_iters_per_tile;
// }
return bid;
}
uint32_t get_current_itr(uint32_t block_idx)
{
uint32_t current_itr = 0;
if(block_idx < sk_num_big_blocks)
{
current_itr = block_idx * k_iters_per_big_block;
}
else if(block_idx < sk_num_blocks)
{
current_itr = (sk_num_big_blocks * k_iters_per_big_block) +
(block_idx - sk_num_big_blocks) * (k_iters_per_big_block - 1);
}
else if(block_idx >= dp_start_block_idx)
{
current_itr = sk_total_iters + (block_idx - dp_start_block_idx) * dp_iters_per_block;
}
return current_itr;
}
void get_block_itr(uint32_t block_idx, uint32_t& iter_start, uint32_t& iter_end)
{
if(block_idx < sk_num_big_blocks)
{
iter_start = block_idx * k_iters_per_big_block;
iter_end = iter_start + k_iters_per_big_block;
}
else if(block_idx < sk_num_blocks)
{
iter_start = (sk_num_big_blocks * k_iters_per_big_block) +
(block_idx - sk_num_big_blocks) * (k_iters_per_big_block - 1);
iter_end = iter_start + (k_iters_per_big_block - 1);
}
else if(block_idx >= dp_start_block_idx)
{
iter_start = sk_total_iters + (block_idx - dp_start_block_idx) * dp_iters_per_block;
iter_end = iter_start + dp_iters_per_block;
}
}
private:
void init()
{
uint32_t num_tiles =
impl::integer_divide_ceil(m, m_per_block) * impl::integer_divide_ceil(n, n_per_block);
k_iters_per_tile = impl::integer_divide_ceil(k, k_per_block);
// one cu can hold one wg at one time, from the whole chip's point of view
// if number of wg is same as num_cu, we call it 1 dispatch
// if number of wg is 2x num_cu, we call it 2 dispatches.
// one dispatch can deliever wg same as num_cu (full dispatch), or less than num_cu (partial
// dispatch)
//
uint32_t full_dispatches = num_tiles / num_cu;
uint32_t full_dispatch_tiles = full_dispatches * num_cu;
uint32_t partial_dispatche_tiles = num_tiles - full_dispatch_tiles;
uint32_t sk_occupancy = occupancy;
uint32_t dp_tiles = full_dispatch_tiles;
uint32_t sk_tiles = partial_dispatche_tiles;
if(full_dispatches < occupancy)
{
// in this case, we allocate all blocks as sk blocks
// sk_occupancy = occupancy - full_dispatches;
sk_occupancy = 1; // TODO: single occ seems better
dp_tiles = full_dispatch_tiles;
sk_tiles = partial_dispatche_tiles;
}
else if((occupancy > 1) && (full_dispatches % occupancy == occupancy - 1))
{
// e.g. occupancy = 2, full_dispatches = 3, 5, 7 ...
// occupancy = 3, full_dispatches = 5, 8, 11 ...
// occupancy = 4, full_dispatches = 7, 11 ...
sk_occupancy = 1; // left 1 slot for sk occupancy
dp_tiles = full_dispatch_tiles;
sk_tiles = partial_dispatche_tiles;
}
else
{
// others, we reduce 1 dispatch from dp, together with partial dispatch,
// to construct sk dispatch
sk_occupancy = occupancy - ((full_dispatches - 1) % occupancy);
dp_tiles = full_dispatch_tiles - num_cu;
sk_tiles = partial_dispatche_tiles + num_cu;
}
// dp_num_blocks = dp_tiles;
// dp_start_block_idx = num_cu * sk_occupancy;
dp_iters_per_block = k_iters_per_tile;
sk_total_iters = k_iters_per_tile * sk_tiles;
// printf("num_tiles:%d, full_dispatches:%d, full_dispatch_tiles:%d,
// partial_dispatche_tiles:%d\n",
// num_tiles, full_dispatches, full_dispatch_tiles, partial_dispatche_tiles);
{
uint32_t min_sk_tiles = (sk_tiles >= num_cu) ? num_cu : (sk_tiles + 1);
uint32_t max_sk_tiles =
(sk_tiles >= num_cu) ? num_cu * sk_occupancy
: impl::min(num_cu, sk_total_iters / min_k_iters_per_sk_block);
// if use dp for sk-block, how many iters do we need
uint32_t dp_for_sk_iters = k_iters_per_tile;
uint32_t best_sk_score =
std::numeric_limits<int>::max(); // we need to find the smallest sk iters
for(uint32_t tentative_sk_blocks = min_sk_tiles; tentative_sk_blocks < max_sk_tiles;
tentative_sk_blocks++)
{
uint32_t tentative_sk_iters_per_block =
(sk_total_iters + tentative_sk_blocks - 1) / tentative_sk_blocks;
uint32_t tentative_sk_iters = tentative_sk_iters_per_block;
uint32_t sk_blocks_per_tile = (tentative_sk_blocks + sk_tiles - 1) / sk_tiles;
// TODO: carefully adjust this parameter
// the more sk_blocks_per_tile, the worse the overhead
uint32_t cross_sk_blocks_overhead = sk_blocks_per_tile;
if(tentative_sk_blocks % sk_tiles != 0)
{
// penalty for uneven divide
cross_sk_blocks_overhead +=
sk_blocks_per_tile * tentative_sk_iters_per_block / 50;
}
uint32_t tentative_sk_score = tentative_sk_iters + cross_sk_blocks_overhead;
if(tentative_sk_score < best_sk_score)
{
best_sk_score = tentative_sk_score;
sk_num_blocks = tentative_sk_blocks;
}
}
if(best_sk_score >= dp_for_sk_iters)
{
sk_num_blocks = 0;
}
if(sk_num_blocks == 0)
{
sk_num_big_blocks = 0;
k_iters_per_big_block = 0;
dp_num_blocks = num_tiles; // all tile to be dp block
dp_start_block_idx = 0;
sk_total_iters = 0; // clear this tiles
}
else
{
uint32_t k_iters_per_sk_block = sk_total_iters / sk_num_blocks;
sk_num_big_blocks = sk_total_iters - k_iters_per_sk_block * sk_num_blocks;
k_iters_per_big_block = k_iters_per_sk_block + 1;
dp_num_blocks = dp_tiles;
dp_start_block_idx = (sk_num_blocks + num_cu - 1) / num_cu * num_cu;
}
}
}
};
struct tile_work_t
{
uint32_t tile_idx;
uint32_t iter_begin;
uint32_t k_begin;
uint32_t k_end;
uint32_t k_iters_remaining;
};
int main(int argc, char** argv)
{
simple_args_t arg = create_arg(argc, argv);
block_dispatcher_t block_dispatcher{arg.get_uint32("m_per_block"),
arg.get_uint32("n_per_block"),
arg.get_uint32("k_per_block"),
arg.get_uint32("num_cu"),
arg.get_uint32("occupancy"),
arg.get_uint32("m"),
arg.get_uint32("n"),
arg.get_uint32("k")};
block_dispatcher.dump();
// simulate actual kernel launch
uint32_t dim_x = block_dispatcher.get_grid_dims_x();
uint32_t total_k_iters =
impl::integer_divide_ceil(arg.get_uint32("k"), arg.get_uint32("k_per_block"));
uint32_t num_tiles =
impl::integer_divide_ceil(arg.get_uint32("m"), arg.get_uint32("m_per_block")) *
impl::integer_divide_ceil(arg.get_uint32("n"), arg.get_uint32("n_per_block"));
std::vector<int> valid_tile_record(num_tiles * total_k_iters);
for(uint32_t bid = 0; bid < dim_x; bid++)
{
uint32_t block_idx = block_dispatcher.get_block_idx(bid);
bool is_sk_block = block_idx < (block_dispatcher.sk_num_blocks);
bool is_dp_block = block_idx >= block_dispatcher.dp_start_block_idx;
uint32_t iter_start, iter_end;
block_dispatcher.get_block_itr(block_idx, iter_start, iter_end);
uint32_t total_iter_length = iter_end - iter_start;
while(true)
{
uint32_t iter_length_mod = iter_end % block_dispatcher.k_iters_per_tile;
uint32_t current_iter_length =
impl::min(iter_length_mod == 0 ? (iter_end - iter_start) : iter_length_mod,
total_iter_length);
uint32_t tile_idx = (iter_end - 1) / block_dispatcher.k_iters_per_tile;
uint32_t tile_iter_start =
((iter_end - 1) % block_dispatcher.k_iters_per_tile) - current_iter_length + 1;
if(is_sk_block)
{
printf("[sk_block] bid:%3d, block_idx:%3d, tile_idx:%3d, iter_start:%d(%d | %d), "
"iter_end:%d (len:%d)\n",
bid,
block_idx,
tile_idx,
iter_end - current_iter_length,
tile_iter_start,
iter_start,
iter_end,
current_iter_length);
}
else if(is_dp_block)
{
printf("[dp_block] bid:%3d, block_idx:%3d, tile_idx:%3d, iter_start:%d(%d | %d), "
"iter_end:%d (len:%d)\n",
bid,
block_idx,
tile_idx,
iter_end - current_iter_length,
tile_iter_start,
iter_start,
iter_end,
current_iter_length);
}
else
{
printf("[other ] bid:%3d, block_idx:%3d\n", bid, block_idx);
}
// some validation check
for(auto i = iter_end - current_iter_length; i < iter_end; i++)
{
if(i >= valid_tile_record.size())
{
printf("unexpected, current iter:%d larger than max:%d\n",
i,
valid_tile_record.size());
return -1;
}
valid_tile_record[i] = 1;
}
iter_end -= current_iter_length;
if(iter_end <= iter_start)
break;
}
}
int untouched = 0;
for(auto i = 0; i < valid_tile_record.size(); i++)
{
if(valid_tile_record[i] != 1)
{
printf("untouched at %d (%d)\n", i, valid_tile_record.size());
untouched++;
}
}
printf("untouched %d/%d, %s\n",
untouched,
valid_tile_record.size(),
untouched == 0 ? "valid" : "fail");
}
CC=g++
$CC -Wall -std=c++17 -Iinclude -O3 block_swizzle_test.cpp -o block_swizzle_test.exe
\ No newline at end of file
#pragma once
#include <iomanip>
#include <iostream>
#include <stdlib.h>
#include <string>
#include <unordered_map>
#include <vector>
#include <assert.h>
struct arg_content_t
{
std::string name; // key
std::string value;
std::string help_text;
};
class simple_args_t
{
public:
simple_args_t() {}
simple_args_t& insert(const std::string& name_,
const std::string& default_value_,
const std::string& help_text_)
{
arg_content_t arg{name_, default_value_, help_text_};
if(arg_map.count(arg.name) != 0)
{
std::cout << "arg:" << arg.name << "already exist" << std::endl;
}
else
{
arg_map[arg.name] = arg;
}
return *this;
}
void usage()
{
for(auto& content : arg_map)
{
std::vector<std::string> help_text_lines;
size_t pos = 0;
for(size_t next_pos = content.second.help_text.find('\n', pos);
next_pos != std::string::npos;)
{
help_text_lines.push_back(
std::string(content.second.help_text.begin() + pos,
content.second.help_text.begin() + next_pos++));
pos = next_pos;
next_pos = content.second.help_text.find('\n', pos);
}
help_text_lines.push_back(std::string(content.second.help_text.begin() + pos,
content.second.help_text.end()));
int arg_name_width = 16 - content.second.name.length();
arg_name_width = arg_name_width > 0 ? arg_name_width : 2;
std::cout << std::setw(4) << "-" << content.second.name << std::setw(arg_name_width)
<< " " << help_text_lines[0] << std::endl;
for(auto help_next_line = std::next(help_text_lines.begin());
help_next_line != help_text_lines.end();
++help_next_line)
{
std::cout << std::setw(28) << " " << *help_next_line << std::endl;
}
}
}
bool parse(int argc, char* argv[], int start_index = 1)
{
if(argc <= start_index)
{
// std::cout << "not enough args (" << argc << ") with starting index " << start_index
// << std::endl;
return true;
}
for(int i = start_index; i < argc; i++)
{
std::string cur_arg = std::string(argv[i]);
if(cur_arg[0] != '-')
{
std::cout << "illegal input" << std::endl;
usage();
return false;
}
else if(cur_arg[0] == '-' && cur_arg[1] == '?')
{
usage();
return false;
}
else
{
size_t found_equal = cur_arg.find('=');
if(found_equal == std::string::npos || found_equal == (cur_arg.length() - 1))
{
std::cout << "failed while parsing \"" << cur_arg << "\", "
<< "arg must be in the form \"-name=value\"" << std::endl;
return false;
}
std::string arg_name = cur_arg.substr(1, found_equal - 1);
std::string arg_value = cur_arg.substr(found_equal + 1);
if(arg_map.count(arg_name) == 0)
{
std::cout << "no such arg \"" << arg_name << "\" registered" << std::endl;
return false;
}
arg_map[arg_name].value = arg_value;
}
}
return true;
}
std::string get(const std::string& name) const { return get_str(name); }
std::string get_str(const std::string& name) const
{
assert(arg_map.count(name) != 0);
std::string value = arg_map.at(name).value;
return value;
}
int get_int(const std::string& name) const
{
assert(arg_map.count(name) != 0);
int value = atoi(arg_map.at(name).value.c_str());
return value;
}
uint32_t get_uint32(const std::string& name) const
{
assert(arg_map.count(name) != 0);
uint32_t value = strtoul(arg_map.at(name).value.c_str(), nullptr, 10);
return value;
}
uint64_t get_uint64(const std::string& name) const
{
assert(arg_map.count(name) != 0);
uint64_t value = strtoull(arg_map.at(name).value.c_str(), nullptr, 10);
return value;
}
double get_double(const std::string& name) const
{
assert(arg_map.count(name) != 0);
double value = atof(arg_map.at(name).value.c_str());
return value;
}
float get_float(const std::string& name) const
{
assert(arg_map.count(name) != 0);
float value = atof(arg_map.at(name).value.c_str());
return value;
}
private:
std::unordered_map<std::string, arg_content_t> arg_map;
};
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment