ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-12 02:05:50 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
lwpck-4181
composable_kernel/include/ck_tile/host/tensor_shuffle_utils.hpp
khuagarw cece6c0c2c debugging permuteN
2026-01-15 22:46:53 +00:00

274 lines
10 KiB
C++
Raw Permalink Blame History

// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include <stdexcept>
namespace ck_tile {
template <typename T>
auto shuffle_aq(const ck_tile::HostTensor<T>* t, int block_aq_k)
{
if(t->get_lengths().size() != 2)
{
throw std::runtime_error("Host tensor is not rank 2 tensor.");
}
int m_ = t->get_lengths()[0];
int aqk_ = t->get_lengths()[1];
if(aqk_ % block_aq_k != 0)
{
throw std::runtime_error("shuffle_aq needs a aqk of multiple times of block_aq_k.");
}
ck_tile::HostTensor<T> t_view({m_, aqk_ / block_aq_k, block_aq_k});
std::copy(t->begin(), t->end(), t_view.begin());
return ck_tile::reference_permute(t_view, {1, 0, 2});
}
template <typename T>
auto shuffle_bq(const ck_tile::HostTensor<T>* t, int block_bq_k)
{
const auto& lengths = t->get_lengths();
const size_t rank = lengths.size();
// Validate block_bq_k divisibility based on rank
int bqk_dim = (rank == 5) ? lengths[4] : (rank == 2) ? lengths[0] : -1;
if(bqk_dim < 0)
{
throw std::runtime_error("shuffle_bq expects either rank-2 or rank-5 tensor, got rank " +
std::to_string(rank));
}
if(bqk_dim % block_bq_k != 0)
{
throw std::runtime_error("shuffle_bq needs bqk dimension to be a multiple of block_bq_k.");
}
// For TilePermuteN
if(rank == 5)
{
// Handle 5D tensor: [n, nrepeat, nwarp, n_warp_tile, bqk]
ck_tile::HostTensor<T> t_view({static_cast<int>(lengths[0]),
static_cast<int>(lengths[1]),
static_cast<int>(lengths[2]),
static_cast<int>(lengths[3]),
bqk_dim / block_bq_k,
block_bq_k});
std::copy(t->begin(), t->end(), t_view.begin());
return ck_tile::reference_permute(t_view, {4, 0, 1, 2, 3, 5});
}
else // rank == 2
{
// Handle 2D tensor: [bqk, n]
int n_ = lengths[1];
ck_tile::HostTensor<T> t_view({n_, bqk_dim / block_bq_k, block_bq_k});
std::copy(t->begin(), t->end(), t_view.begin());
// printf("I am inside shuffle_bq\n");
// printf("t_view.get_lengths(): %lu, %lu, %lu\n",
// t_view.get_lengths()[0],
// t_view.get_lengths()[1],
// t_view.get_lengths()[2]);
// for(int i = 0; i < static_cast<int>(t_view.get_lengths()[0]); i++)
// {
// for(int j = 0; j < static_cast<int>(t_view.get_lengths()[1]); j++)
// {
// for(int k = 0; k < static_cast<int>(t_view.get_lengths()[2]); k++)
// {
// printf("t_view[%d][%d][%d]: %f\n", i, j, k, t_view(i, j, k));
// }
// }
// }
// printf("I am inside shuffle_bq\n");
return ck_tile::reference_permute(t_view, {1, 0, 2});
}
}
template <typename GemmConfig, typename T>
auto shuffle_b(const ck_tile::HostTensor<T>& t, const GemmConfig& gemmConfig)
{
assert(t.get_lengths().size() == 2);
int n_ = t.get_lengths()[1];
int k_ = t.get_lengths()[0];
if(ck_tile::is_gfx12_supported())
{
constexpr int divisor = 2;
constexpr int kABK1PerLane = 8;
int kABK0PerLane = gemmConfig.K_Warp_Tile / divisor / kABK1PerLane;
ck_tile::HostTensor<T> t_view({n_ / gemmConfig.N_Warp_Tile,
gemmConfig.N_Warp_Tile,
k_ / gemmConfig.K_Warp_Tile,
kABK0PerLane,
divisor,
kABK1PerLane});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 2, 4, 1, 3, 5});
}
else
{
int divisor = 1;
if(ck_tile::is_gfx11_supported())
{
divisor = 1;
}
else
{
assert(is_wave32() == false);
divisor = gemmConfig.N_Warp_Tile == 32 ? 2 : 4;
}
ck_tile::HostTensor<T> t_view({n_ / gemmConfig.N_Warp_Tile,
gemmConfig.N_Warp_Tile,
k_ / gemmConfig.K_Warp_Tile,
divisor,
gemmConfig.K_Warp_Tile / divisor});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4});
}
}
template <typename GemmConfig, typename T>
auto shuffle_b(const ck_tile::HostTensor<T>& t)
{
return shuffle_b(t, GemmConfig{});
}
template <typename GemmConfig, typename T>
auto bq_permuteN(const ck_tile::HostTensor<T>& t, index_t group_n)
{
assert(t.get_lengths().size() == 2);
int n_ = t.get_lengths()[1]; // N/grp_size_N
int bqk_ = t.get_lengths()[0];
constexpr int NRepeat = GemmConfig::N_Tile / GemmConfig::N_Warp_Tile / GemmConfig::N_Warp;
int dim = n_ / (GemmConfig::N_Warp * (GemmConfig::N_Warp_Tile / group_n) * NRepeat);
ck_tile::HostTensor<T> t_view(
{dim, GemmConfig::N_Warp, GemmConfig::N_Warp_Tile / group_n, NRepeat, bqk_});
std::copy(t.begin(), t.end(), t_view.begin());
printf("I am inside shuffle_bq PermuteN\n");
printf("t_view.get_lengths(): %lu, %lu, %lu, %lu, %lu\n",
t_view.get_lengths()[0],
t_view.get_lengths()[1],
t_view.get_lengths()[2],
t_view.get_lengths()[3],
t_view.get_lengths()[4]);
for(int i = 0; i < static_cast<int>(t_view.get_lengths()[0]); i++)
{
for(int j = 0; j < static_cast<int>(t_view.get_lengths()[1]); j++)
{
for(int k = 0; k < static_cast<int>(t_view.get_lengths()[2]); k++)
{
for(int l = 0; l < static_cast<int>(t_view.get_lengths()[3]); l++)
{
for(int m = 0; m < static_cast<int>(t_view.get_lengths()[4]); m++)
{
printf("t_view[%d][%d][%d][%d][%d]: %f\n",
i,
j,
k,
l,
m,
static_cast<float>(t_view(i, j, k, l, m)));
}
}
}
}
}
return ck_tile::reference_permute(t_view, {0, 3, 1, 2, 4});
}
template <typename GemmConfig, typename T>
auto shuffle_b_permuteN(const ck_tile::HostTensor<T>& t, const GemmConfig& gemmConfig)
{
assert(t.get_lengths().size() == 2);
int n_ = t.get_lengths()[1];
int k_ = t.get_lengths()[0];
printf("KxN: %dx%d\n", k_, n_);
int NRepeat = gemmConfig.N_Tile / gemmConfig.N_Warp_Tile / gemmConfig.N_Warp;
printf("gemmConfig.N_tile: %d, gemmConfig.N_warp_tile: %d, gemmConfig.N_Warp: %d, "
"gemmConfig.K_warp_tile: %d\n",
gemmConfig.N_Tile,
gemmConfig.N_Warp_Tile,
gemmConfig.N_Warp,
gemmConfig.K_Warp_Tile);
if(ck_tile::is_gfx12_supported())
{
constexpr int divisor = 2;
constexpr int kABK1PerLane = 8;
int kABK0PerLane = gemmConfig.K_Warp_Tile / divisor / kABK1PerLane;
ck_tile::HostTensor<T> t_view({n_ / gemmConfig.N_Tile,
gemmConfig.N_Warp,
gemmConfig.N_Warp_Tile,
NRepeat,
k_ / gemmConfig.K_Warp_Tile,
kABK0PerLane,
divisor,
kABK1PerLane});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 3, 1, 4, 6, 5, 2, 7});
}
else
{
int divisor = 1;
if(ck_tile::is_gfx11_supported())
{
divisor = 1;
}
else
{
assert(is_wave32() == false);
divisor = gemmConfig.N_Warp_Tile == 32 ? 2 : 4;
}
ck_tile::HostTensor<T> t_view({n_ / gemmConfig.N_Tile,
gemmConfig.N_Warp,
gemmConfig.N_Warp_Tile,
NRepeat,
k_ / gemmConfig.K_Warp_Tile,
divisor,
gemmConfig.K_Warp_Tile / divisor});
std::copy(t.begin(), t.end(), t_view.begin());
printf("I am inside shuffle_b PermuteN\n");
printf("t_view.get_lengths(): %lu, %lu, %lu, %lu, %lu, %lu, %lu\n",
t_view.get_lengths()[0],
t_view.get_lengths()[1],
t_view.get_lengths()[2],
t_view.get_lengths()[3],
t_view.get_lengths()[4],
t_view.get_lengths()[5],
t_view.get_lengths()[6]);
// for(int i = 0; i < static_cast<int>(t_view.get_lengths()[0]); i++)
// {
// for(int j = 0; j < static_cast<int>(t_view.get_lengths()[1]); j++)
// {
// for(int k = 0; k < static_cast<int>(t_view.get_lengths()[2]); k++)
// {
// for(int l = 0; l < static_cast<int>(t_view.get_lengths()[3]); l++)
// {
// for(int m = 0; m < static_cast<int>(t_view.get_lengths()[4]); m++)
// {
// for(int n = 0; n < static_cast<int>(t_view.get_lengths()[5]); n++)
// {
// for(int p = 0; p < static_cast<int>(t_view.get_lengths()[6]);
// p++)
// {
// printf("t_view[%d][%d][%d][%d][%d][%d][%d]: %d\n", i, j, k,
// l, m, n, p, static_cast<int>(t_view(i, j, k,l, m, n, p)));
// }
// }
// }
// }
// }
// }
// }
return ck_tile::reference_permute(t_view, {0, 3, 1, 4, 5, 2, 6});
}
}
template <typename GemmConfig, typename T>
auto shuffle_b_permuteN(const ck_tile::HostTensor<T>& t)
{
return shuffle_b_permuteN(t, GemmConfig{});
}
} // namespace ck_tile
Reference in New Issue View Git Blame Copy Permalink