mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-04-19 22:39:03 +00:00
transpose load api development (#2177)
* add transpose load; no real logic * fix some compile errors * fix some issues * update transpose load logic * add some fixes * fix a distribution issue * update some codes * add some fix * can pass; but no logic * transpose load enable * update tile transpose * miss output tile distribution mapping * hack for transpose 16x16 * update output tensor distribution * delete unused variables * fix transpose related codes * update transpose load example * exchange the iteration order * fix 16x16 related dimension transpose * fix a transpose index issue * fix a transpose index issue * fix clang format check * update load tile transpose related codes * fix compile errors and pass 16x16 tests * fix a typo * update logic * check other data types * add transpose load api * update transpose load api * fix clang format check * change file name * refactor codes * update code name * delete some unused codes * delete the unused oob flag for transpose load * update tensor view api for transpose load * update for testing * fix a typo error * move transpose ops to example directory * update transpose api * update include file * fix for pr review * fix compile errors * add transpose load; no real logic * fix some compile errors * fix some issues * update transpose load logic * add some fixes * fix a distribution issue * update some codes * add some fix * can pass; but no logic * transpose load enable * update tile transpose * miss output tile distribution mapping * hack for transpose 16x16 * update output tensor distribution * delete unused variables * fix transpose related codes * update transpose load example * exchange the iteration order * fix 16x16 related dimension transpose * fix a transpose index issue * fix a transpose index issue * fix clang format check * update load tile transpose related codes * fix compile errors and pass 16x16 tests * fix a typo * update logic * check other data types * add transpose load api * update transpose load api * fix clang format check * change file name * refactor codes * update code name * delete some unused codes * delete the unused oob flag for transpose load * update tensor view api for transpose load * update for testing * fix a typo error * move transpose ops to example directory * update transpose api * update include file * fix for pr review * fix compile errors * change directory name * delete the duplicated directory * update cmakelists file * delete the unused codes * update function names * update transpose policy * update code after remod.py * update codes * add some comment * Polish the instr infrastructure * build up the fixed instr * redesign the transpose api, currently it has numerical error * add the bf16 transpose * fix some issues * add some comments * update document * Finished the refactor of API and pass through the verification * fix the merging issue --------- Co-authored-by: ThomasNing <thomas.ning@amd.com>
This commit is contained in:
joyeamd
9
example/ck_tile/37_transpose/CMakeLists.txt
Normal file
9
example/ck_tile/37_transpose/CMakeLists.txt
Normal file
@@ -0,0 +1,9 @@
|
||||
set(TARGET_NAME tile_example_transpose)
|
||||
add_executable(${TARGET_NAME} EXCLUDE_FROM_ALL transpose_example.cpp transpose_api.cpp)
|
||||
target_include_directories(${TARGET_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/)
|
||||
|
||||
# NOTE: we turn off undefined-func-template to let source compile without explicit declare function specializations
|
||||
list(APPEND EXAMPLE_BATCHED_TRANSPOSE_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal)
|
||||
# list(APPEND EXAMPLE_BATCHED_TRANSPOSE_COMPILE_OPTIONS -v --save-temps -Wno-gnu-line-marker)
|
||||
target_compile_options(tile_example_transpose PRIVATE ${EXAMPLE_BATCHED_TRANSPOSE_COMPILE_OPTIONS})
|
||||
|
||||
27
example/ck_tile/37_transpose/README.md
Normal file
27
example/ck_tile/37_transpose/README.md
Normal file
@@ -0,0 +1,27 @@
|
||||
# Batched Transpose
|
||||
This folder contains example for transpose load for architecture gfx950. This transpose load has some constraints in input tile distribution.
|
||||
|
||||
## 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 ../ <arch>
|
||||
# Make the transpose executable
|
||||
make tile_example_transpose -j
|
||||
```
|
||||
This will result in an executable `build/bin/tile_example_transpose`
|
||||
|
||||
## example
|
||||
```
|
||||
args:
|
||||
-N input batch size (default:2)
|
||||
-C input channel size. (default:64)
|
||||
-H input height size. (default:1)
|
||||
-W input width size. (default:64)
|
||||
-v whether do CPU validation or not (default: 1)
|
||||
-layout_in input tensor data layout - NCHW by default
|
||||
-layout_out output tensor data layout - NHWC by default
|
||||
-seed seed to be used, -1 means random every time (default:-1)
|
||||
-k_name t to 1 will print kernel name (default:0)
|
||||
```
|
||||
120
example/ck_tile/37_transpose/batched_transpose_kernel.hpp
Normal file
120
example/ck_tile/37_transpose/batched_transpose_kernel.hpp
Normal file
@@ -0,0 +1,120 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/ops/common.hpp"
|
||||
#include "ck_tile/ops/elementwise.hpp"
|
||||
#include "ck_tile/host/hip_check_error.hpp"
|
||||
#include <string>
|
||||
#include <type_traits>
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
struct BatchedTransposeHostArgs
|
||||
{
|
||||
const void* p_input;
|
||||
void* p_output;
|
||||
index_t batch;
|
||||
index_t height;
|
||||
index_t width;
|
||||
// index_t dim_blocks;
|
||||
index_t dim_stride;
|
||||
index_t dim_block_h;
|
||||
index_t dim_block_w;
|
||||
};
|
||||
|
||||
template <typename Pipeline_>
|
||||
struct BatchedTransposeKernel
|
||||
{
|
||||
using Pipeline = remove_cvref_t<Pipeline_>;
|
||||
using Problem = remove_cvref_t<typename Pipeline::Problem>;
|
||||
|
||||
using Type = typename Problem::DataType;
|
||||
|
||||
struct BatchedTransposeKargs
|
||||
{
|
||||
const void* p_input;
|
||||
void* p_output;
|
||||
index_t batch;
|
||||
index_t height;
|
||||
index_t width;
|
||||
index_t dim_stride;
|
||||
};
|
||||
|
||||
using Kargs = BatchedTransposeKargs;
|
||||
using Hargs = BatchedTransposeHostArgs;
|
||||
|
||||
CK_TILE_HOST static constexpr auto GridSize(const Hargs& h)
|
||||
{
|
||||
size_t grid_size_x = h.dim_block_w;
|
||||
size_t grid_size_y = h.dim_block_h;
|
||||
size_t grid_size_z = h.batch;
|
||||
return dim3(grid_size_x, grid_size_y, grid_size_z);
|
||||
}
|
||||
|
||||
CK_TILE_HOST static constexpr auto MakeKargs(const Hargs& h)
|
||||
{
|
||||
Kargs k;
|
||||
k.p_input = h.p_input;
|
||||
k.p_output = h.p_output;
|
||||
k.batch = h.batch;
|
||||
k.height = h.height;
|
||||
k.width = h.width;
|
||||
k.dim_stride = h.dim_stride;
|
||||
return k;
|
||||
}
|
||||
|
||||
CK_TILE_HOST_DEVICE static constexpr auto BlockSize() { return Problem::kBlockSize; }
|
||||
|
||||
CK_TILE_DEVICE void operator()(Kargs kargs) const
|
||||
{
|
||||
__shared__ char smem[Pipeline::GetSmemSize()];
|
||||
static constexpr ck_tile::index_t kMPerBlock = Problem::kSecondSizePerBlock;
|
||||
static constexpr ck_tile::index_t kNPerBlock = Problem::kLeadSizePerBlock;
|
||||
|
||||
const auto iDim = blockIdx.z;
|
||||
const auto x_m_n = [&]() {
|
||||
const auto x_dram_naive = make_naive_tensor_view<address_space_enum::global>(
|
||||
static_cast<const Type*>(kargs.p_input) + iDim * kargs.dim_stride,
|
||||
make_tuple(kargs.height, kargs.width),
|
||||
make_tuple(kargs.width, 1),
|
||||
number<Pipeline::GetVectorSize()>{},
|
||||
number<1>{});
|
||||
|
||||
return pad_tensor_view(x_dram_naive,
|
||||
make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}),
|
||||
sequence<false, false>{});
|
||||
}();
|
||||
|
||||
const auto iM = __builtin_amdgcn_readfirstlane(blockIdx.y * kMPerBlock);
|
||||
const auto iN = __builtin_amdgcn_readfirstlane(blockIdx.x * kNPerBlock);
|
||||
|
||||
const auto y_n_m = [&]() {
|
||||
const auto y_dram_naive = make_naive_tensor_view<address_space_enum::global>(
|
||||
static_cast<Type*>(kargs.p_output) + iDim * kargs.dim_stride,
|
||||
make_tuple(kargs.width, kargs.height),
|
||||
make_tuple(kargs.height, 1),
|
||||
number<Pipeline::GetVectorSize()>{},
|
||||
number<1>{});
|
||||
|
||||
return pad_tensor_view(y_dram_naive,
|
||||
make_tuple(number<kNPerBlock>{}, number<kMPerBlock>{}),
|
||||
sequence<false, false>{});
|
||||
}();
|
||||
|
||||
auto x_block_window = make_tile_window(
|
||||
x_m_n,
|
||||
make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}),
|
||||
{static_cast<ck_tile::index_t>(iM), static_cast<ck_tile::index_t>(iN)});
|
||||
|
||||
auto y_block_window = make_tile_window(
|
||||
y_n_m,
|
||||
make_tuple(number<kNPerBlock>{}, number<kMPerBlock>{}),
|
||||
{static_cast<ck_tile::index_t>(iN), static_cast<ck_tile::index_t>(iM)});
|
||||
|
||||
Pipeline{}(x_block_window, y_block_window, smem);
|
||||
}
|
||||
};
|
||||
} // namespace ck_tile
|
||||
149
example/ck_tile/37_transpose/block_transpose.hpp
Normal file
149
example/ck_tile/37_transpose/block_transpose.hpp
Normal file
@@ -0,0 +1,149 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "transpose_policy.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
template <typename Layout_, index_t kRow, index_t kCol>
|
||||
struct TransposeTraits
|
||||
{
|
||||
static constexpr index_t kLeadDim = kCol;
|
||||
static constexpr index_t kSecondDim = kRow;
|
||||
};
|
||||
|
||||
template <index_t kRow, index_t kCol>
|
||||
struct TransposeTraits<tensor_layout::gemm::ColumnMajor, kRow, kCol>
|
||||
{
|
||||
static constexpr index_t kLeadDim = kRow;
|
||||
static constexpr index_t kSecondDim = kCol;
|
||||
};
|
||||
|
||||
// supports 2D transpose which will store to lds, then use ds_read_b*_tr_b* instruction to get the
|
||||
// transposed data; Layout in TransposePipelineProblem is the original layout of the data in the
|
||||
// global memory
|
||||
template <typename DataType_,
|
||||
typename Layout_,
|
||||
index_t kBlockSize_,
|
||||
index_t kRowWarps_, // how many warps in row direction
|
||||
index_t kColWarps_, // how many warps in col direction
|
||||
index_t kRowPerBlock_, // row number per block
|
||||
index_t kColPerBlock_, // col number per block
|
||||
index_t kRowPerXdl_, // row number per xdl ops
|
||||
index_t kColPerXdl_> // col number per xdl ops
|
||||
struct TransposePipelineProblem
|
||||
{
|
||||
static_assert(kRowWarps_ * kColWarps_ * get_warp_size() == kBlockSize_,
|
||||
"the block size is not correct!");
|
||||
using DataType = remove_cvref_t<DataType_>;
|
||||
using Layout = remove_cvref_t<Layout_>;
|
||||
static constexpr index_t kBlockSize = kBlockSize_;
|
||||
static constexpr index_t kLeadNumWarps =
|
||||
TransposeTraits<Layout, kRowWarps_, kColWarps_>::kLeadDim;
|
||||
static constexpr index_t kSecondNumWarps =
|
||||
TransposeTraits<Layout, kRowWarps_, kColWarps_>::kSecondDim;
|
||||
static constexpr index_t kLeadSizePerBlock =
|
||||
TransposeTraits<Layout, kRowPerBlock_, kColPerBlock_>::kLeadDim;
|
||||
static constexpr index_t kSecondSizePerBlock =
|
||||
TransposeTraits<Layout, kRowPerBlock_, kColPerBlock_>::kSecondDim;
|
||||
static constexpr index_t kLeadSizePerXdl =
|
||||
TransposeTraits<Layout, kRowPerXdl_, kColPerXdl_>::kLeadDim;
|
||||
static constexpr index_t kSecondSizePerXdl =
|
||||
TransposeTraits<Layout, kRowPerXdl_, kColPerXdl_>::kSecondDim;
|
||||
|
||||
static constexpr index_t kQuadrantLeadDim = LaneGroupTransposeTraits<DataType>::kleadDim;
|
||||
static constexpr index_t kQuadrantSecondDim = LaneGroupTransposeTraits<DataType>::ksecondDim;
|
||||
|
||||
static_assert(kLeadSizePerBlock % kLeadNumWarps == 0,
|
||||
"block dim should be divided by warp dim!");
|
||||
static_assert(kSecondSizePerBlock % kSecondNumWarps == 0,
|
||||
"block dim should be divided by warp dim!");
|
||||
// how many rows/cols implemented in one warp
|
||||
static constexpr index_t kLeadSizePerWarp = kLeadSizePerBlock / kLeadNumWarps;
|
||||
static constexpr index_t kSecondSizePerWarp = kSecondSizePerBlock / kSecondNumWarps;
|
||||
|
||||
static_assert(kLeadSizePerWarp % kLeadSizePerXdl == 0,
|
||||
"warp dim should be divided by xdl dim!");
|
||||
static_assert(kSecondSizePerWarp % kSecondSizePerXdl == 0,
|
||||
"warp dim should be divided by xdl dim!");
|
||||
|
||||
// warp rows/cols is divided into xdl.
|
||||
static constexpr index_t kLeadXdlNumPerWarp = kLeadSizePerWarp / kLeadSizePerXdl;
|
||||
static constexpr index_t kSecondXdlNumPerWarp = kSecondSizePerWarp / kSecondSizePerXdl;
|
||||
|
||||
static_assert(kLeadSizePerXdl % kQuadrantLeadDim == 0,
|
||||
"xdl dim should be divided by quad dim!");
|
||||
static_assert(kSecondSizePerXdl % kQuadrantSecondDim == 0,
|
||||
"xdl dim should be divided by quad dim!");
|
||||
// xdl rows/cols is divided into quadrants.
|
||||
static constexpr index_t kQuadNumPerLeadDim = kLeadSizePerXdl / kQuadrantLeadDim;
|
||||
static constexpr index_t kQuadNumPerSecondDim = kSecondSizePerXdl / kQuadrantSecondDim;
|
||||
|
||||
static constexpr index_t kIterationsInSecondDim =
|
||||
kQuadNumPerLeadDim * kQuadNumPerSecondDim * 16 / get_warp_size();
|
||||
};
|
||||
|
||||
template <typename Problem_, typename Policy_ = TransposePolicy>
|
||||
struct BlockTranspose
|
||||
{
|
||||
using Problem = remove_cvref_t<Problem_>;
|
||||
using Policy = remove_cvref_t<Policy_>;
|
||||
|
||||
using DataType = remove_cvref_t<typename Problem::DataType>;
|
||||
using Layout = remove_cvref_t<typename Problem::Layout>;
|
||||
|
||||
static constexpr index_t kBlockSize = Problem::kBlockSize;
|
||||
static constexpr index_t kLeadSizePerBlock = Problem::kLeadSizePerBlock;
|
||||
static constexpr index_t kSecondSizePerBlock = Problem::kSecondSizePerBlock;
|
||||
|
||||
static constexpr index_t GetVectorSize() { return Policy::template GetVectorSize<Problem>(); }
|
||||
|
||||
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
|
||||
{
|
||||
return Policy::template GetSmemSize<Problem>();
|
||||
}
|
||||
|
||||
template <typename InputTileWindow, typename OutputTileWindow>
|
||||
CK_TILE_DEVICE void operator()(const InputTileWindow& input_window,
|
||||
OutputTileWindow& output_window,
|
||||
void* __restrict__ p_smem)
|
||||
{
|
||||
auto input_tile_window =
|
||||
make_tile_window(input_window, Policy::template MakeInputDistribution<Problem>());
|
||||
auto output_tile_window =
|
||||
make_tile_window(output_window, Policy::template MakeOutputDistribution<Problem>());
|
||||
|
||||
DataType* p_lds_ptr = static_cast<DataType*>(p_smem);
|
||||
constexpr auto in_lds_block_desc = Policy::template MakeLdsStoreBlockDescriptor<Problem>();
|
||||
auto input_lds_block =
|
||||
make_tensor_view<address_space_enum::lds>(p_lds_ptr, in_lds_block_desc);
|
||||
|
||||
constexpr auto out_lds_block_desc = Policy::template MakeLdsLoadBlockDescriptor<Problem>();
|
||||
auto output_lds_block =
|
||||
make_tensor_view<address_space_enum::lds>(p_lds_ptr, out_lds_block_desc);
|
||||
|
||||
auto copy_to_lds_window =
|
||||
make_tile_window(input_lds_block,
|
||||
make_tuple(number<kSecondSizePerBlock>{}, number<kLeadSizePerBlock>{}),
|
||||
{0, 0});
|
||||
auto load_from_lds_window =
|
||||
make_tile_window(output_lds_block,
|
||||
make_tuple(number<kSecondSizePerBlock>{}, number<kLeadSizePerBlock>{}),
|
||||
{0, 0},
|
||||
Policy::template MakeLdsLoadTileDistribution<Problem>());
|
||||
|
||||
auto x = load_tile(input_tile_window);
|
||||
|
||||
store_tile(copy_to_lds_window, x);
|
||||
block_sync_lds();
|
||||
|
||||
auto y = load_tile_transpose(load_from_lds_window);
|
||||
|
||||
store_tile(output_tile_window, y);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
59
example/ck_tile/37_transpose/transpose_api.cpp
Normal file
59
example/ck_tile/37_transpose/transpose_api.cpp
Normal file
@@ -0,0 +1,59 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
#include "transpose_example.hpp"
|
||||
#include <iostream>
|
||||
|
||||
template <typename ts_type,
|
||||
ck_tile::index_t block_x,
|
||||
ck_tile::index_t block_y,
|
||||
ck_tile::index_t warp_x,
|
||||
ck_tile::index_t warp_y>
|
||||
float batched_transpose_dispatch(batched_transpose_kargs& a, ck_tile::stream_config& s)
|
||||
{
|
||||
uint32_t dim_block_h = (a.height + block_y - 1) / block_y;
|
||||
uint32_t dim_block_w = (a.width + block_x - 1) / block_x;
|
||||
uint32_t dim_stride = a.height * a.width;
|
||||
|
||||
a.dim_stride = dim_stride;
|
||||
a.dim_block_h = dim_block_h;
|
||||
a.dim_block_w = dim_block_w;
|
||||
|
||||
using ts_problem = ck_tile::TransposePipelineProblem<ts_type,
|
||||
ck_tile::tensor_layout::gemm::RowMajor,
|
||||
64,
|
||||
1,
|
||||
1,
|
||||
block_y,
|
||||
block_x,
|
||||
warp_y,
|
||||
warp_x>;
|
||||
using ts_pipeline = ck_tile::BlockTranspose<ts_problem>;
|
||||
|
||||
using kernel = ck_tile::BatchedTransposeKernel<ts_pipeline>;
|
||||
|
||||
auto kargs = kernel::MakeKargs(a);
|
||||
|
||||
const dim3 grids = kernel::GridSize(a);
|
||||
constexpr dim3 blocks = kernel::BlockSize();
|
||||
|
||||
float ave_time = ck_tile::launch_kernel(
|
||||
s, ck_tile::make_kernel<blocks.x, 1>(kernel{}, grids, blocks, 0, kargs));
|
||||
|
||||
return ave_time;
|
||||
}
|
||||
|
||||
float batched_transpose(batched_transpose_trait t,
|
||||
batched_transpose_kargs a,
|
||||
ck_tile::stream_config s)
|
||||
{
|
||||
if(t.type == "fp16")
|
||||
{
|
||||
return batched_transpose_dispatch<ck_tile::fp16_t, 16, 32, 16, 32>(a, s);
|
||||
}
|
||||
else if(t.type == "fp8")
|
||||
{
|
||||
return batched_transpose_dispatch<ck_tile::fp8_t, 16, 64, 16, 64>(a, s);
|
||||
}
|
||||
|
||||
return -1;
|
||||
}
|
||||
257
example/ck_tile/37_transpose/transpose_example.cpp
Normal file
257
example/ck_tile/37_transpose/transpose_example.cpp
Normal file
@@ -0,0 +1,257 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#include <vector>
|
||||
#include <iostream>
|
||||
#include <numeric>
|
||||
#include <cassert>
|
||||
#include <cstdlib>
|
||||
#include <iostream>
|
||||
#include <time.h>
|
||||
#include <unordered_set>
|
||||
|
||||
#include "transpose_example.hpp"
|
||||
|
||||
#if 0
|
||||
template <typename T>
|
||||
void dump_host_tensor_4d(const ck_tile::HostTensor<T>& x)
|
||||
{
|
||||
auto len = x.get_lengths();
|
||||
assert(len.size() == 4);
|
||||
std::cout << "[";
|
||||
for(size_t i = 0; i < len[0]; i++)
|
||||
{
|
||||
std::cout << i << ": [";
|
||||
for(size_t j = 0; j < len[1]; j++)
|
||||
{
|
||||
std::cout << j << ": [";
|
||||
for(size_t k = 0; k < len[2]; k++)
|
||||
{
|
||||
std::cout << k << ": [";
|
||||
for(size_t v = 0; v < len[3]; v++)
|
||||
{
|
||||
if constexpr(std::is_same_v<T, ck_tile::fp16_t>)
|
||||
{
|
||||
auto m =
|
||||
ck_tile::type_convert<float>(x(std::vector<std::size_t>{i, j, k, v}));
|
||||
|
||||
std::cout << m;
|
||||
if(v != len[3] - 1)
|
||||
std::cout << ",";
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << x(std::vector<std::size_t>{i, j, k, v}) << " ";
|
||||
}
|
||||
}
|
||||
std::cout << "]" << std::endl;
|
||||
}
|
||||
std::cout << "]" << std::endl;
|
||||
}
|
||||
std::cout << std::endl;
|
||||
}
|
||||
std::cout << "--------------------" << std::endl;
|
||||
}
|
||||
#endif
|
||||
|
||||
// different threshold for different dtype
|
||||
template <typename DataType>
|
||||
auto get_elimit(std::string /*init_method*/)
|
||||
{
|
||||
double rtol = 1e-3;
|
||||
double atol = 1e-3;
|
||||
return ck_tile::make_tuple(rtol, atol);
|
||||
}
|
||||
|
||||
template <>
|
||||
auto get_elimit<ck_tile::bf16_t>(std::string /*init_method*/)
|
||||
{
|
||||
double rtol = 1e-2;
|
||||
double atol = 1e-2;
|
||||
return ck_tile::make_tuple(rtol, atol);
|
||||
}
|
||||
|
||||
template <>
|
||||
auto get_elimit<ck_tile::fp8_t>(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);
|
||||
}
|
||||
}
|
||||
|
||||
auto create_args(int argc, char* argv[])
|
||||
{
|
||||
ck_tile::ArgParser arg_parser;
|
||||
arg_parser.insert("v", "1", "whether do CPU validation or not")
|
||||
.insert("pr", "fp16", "input data type. fp16/fp32 (representing 8/16/32 bit data)")
|
||||
.insert("N", "2", "input batch size. ")
|
||||
.insert("C", "64", "input channel size.")
|
||||
.insert("H", "1", "input height size.")
|
||||
.insert("W", "64", "input width size. ")
|
||||
.insert("layout_in", "NCHW", "input tensor data layout - NCHW by default")
|
||||
.insert("layout_out", "NHWC", "output tensor data layout - NHWC by default ")
|
||||
.insert("seed", "-1", "seed to be used, -1 means random every time")
|
||||
.insert("kname", "0", "t to 1 will print kernel name");
|
||||
|
||||
bool result = arg_parser.parse(argc, argv);
|
||||
return std::make_tuple(result, arg_parser);
|
||||
}
|
||||
|
||||
template <typename Type>
|
||||
bool run_batched_transpose(ck_tile::ArgParser args)
|
||||
{
|
||||
int validate = args.get_int("v");
|
||||
std::string prec = args.get_str("pr");
|
||||
int N = args.get_int("N");
|
||||
int C = args.get_int("C");
|
||||
int H = args.get_int("H");
|
||||
int W = args.get_int("W");
|
||||
std::string layout_in = args.get_str("layout_in");
|
||||
std::string layout_out = args.get_str("layout_out");
|
||||
int seed = args.get_int("seed");
|
||||
|
||||
int dim_in[4], dim_out[4];
|
||||
int stride_dim_in[4], stride_dim_out[4];
|
||||
bool nchw2nhwc = layout_in == "NCHW" && layout_out == "NHWC";
|
||||
bool nhwc2nchw = layout_in == "NHWC" && layout_out == "NCHW";
|
||||
assert(nchw2nhwc != nhwc2nchw);
|
||||
(void)nhwc2nchw;
|
||||
|
||||
dim_in[0] = N;
|
||||
dim_in[1] = nchw2nhwc ? C : H;
|
||||
dim_in[2] = nchw2nhwc ? H : W;
|
||||
dim_in[3] = nchw2nhwc ? W : C;
|
||||
dim_out[0] = N;
|
||||
dim_out[1] = nchw2nhwc ? H : C;
|
||||
dim_out[2] = nchw2nhwc ? W : H;
|
||||
dim_out[3] = nchw2nhwc ? C : W;
|
||||
stride_dim_in[0] = C * H * W;
|
||||
stride_dim_in[1] = nchw2nhwc ? H * W : C * W;
|
||||
stride_dim_in[2] = nchw2nhwc ? W : C;
|
||||
stride_dim_in[3] = 1;
|
||||
stride_dim_out[0] = C * H * W;
|
||||
stride_dim_out[1] = nchw2nhwc ? C * W : H * W;
|
||||
stride_dim_out[2] = nchw2nhwc ? C : W;
|
||||
stride_dim_out[3] = 1;
|
||||
|
||||
if(seed < 0)
|
||||
{
|
||||
seed = std::time(nullptr);
|
||||
}
|
||||
|
||||
ck_tile::HostTensor<Type> x_host(
|
||||
{dim_in[0], dim_in[1], dim_in[2], dim_in[3]},
|
||||
{stride_dim_in[0], stride_dim_in[1], stride_dim_in[2], stride_dim_in[3]});
|
||||
ck_tile::HostTensor<Type> y_host(
|
||||
{dim_out[0], dim_out[1], dim_out[2], dim_out[3]},
|
||||
{stride_dim_out[0], stride_dim_out[1], stride_dim_out[2], stride_dim_out[3]});
|
||||
|
||||
ck_tile::FillUniformDistribution<Type>{-.5f, .5f}(x_host);
|
||||
|
||||
ck_tile::DeviceMem x_dev(x_host.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem y_dev(y_host.get_element_space_size_in_bytes());
|
||||
|
||||
x_dev.ToDevice(x_host.data());
|
||||
|
||||
auto trait = batched_transpose_trait{prec, layout_in};
|
||||
|
||||
uint32_t height = nchw2nhwc ? C : H * W;
|
||||
uint32_t width = nchw2nhwc ? H * W : C;
|
||||
|
||||
batched_transpose_kargs karg = [&]() {
|
||||
batched_transpose_kargs a_;
|
||||
a_.p_input = x_dev.GetDeviceBuffer();
|
||||
a_.p_output = y_dev.GetDeviceBuffer();
|
||||
a_.batch = N;
|
||||
a_.height = height;
|
||||
a_.width = width;
|
||||
return a_;
|
||||
}();
|
||||
|
||||
ck_tile::stream_config sc{nullptr, true};
|
||||
|
||||
auto ms = batched_transpose(trait, karg, sc);
|
||||
|
||||
std::size_t num_operations = N * C * H * (W - 1);
|
||||
std::size_t num_bytes = N * C * H * W * sizeof(Type);
|
||||
|
||||
float ave_time = ms * 1E-3;
|
||||
float gb_per_sec = num_bytes / ms * 1.E-6;
|
||||
float tflops = static_cast<float>(num_operations) / ms * 1.E-6;
|
||||
|
||||
std::cout << "Run Batched Transpose kernel with N=" << N << ", C=" << C << ", H=" << H
|
||||
<< ", W=" << W << ", layout_in=" << layout_in << ", layout_out=" << layout_out
|
||||
<< " : " << ms << " ms (" << ave_time << " ave_time), " << tflops << " TFlops"
|
||||
<< gb_per_sec << " GB/s, " << std::endl;
|
||||
|
||||
printf("[%s]N:%d, C:%d, H:%d, W:%d, layout_in:%s, %f\n",
|
||||
prec.c_str(),
|
||||
N,
|
||||
C,
|
||||
H,
|
||||
W,
|
||||
layout_in.c_str(),
|
||||
ms);
|
||||
if(ms < 0)
|
||||
printf("not supported\n");
|
||||
fflush(stdout);
|
||||
|
||||
if(ms < 0)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
y_dev.FromDevice(y_host.data());
|
||||
|
||||
bool rtn = true;
|
||||
if(validate)
|
||||
{
|
||||
// this host buffer will not copy to GPU, so no need use stride
|
||||
ck_tile::HostTensor<Type> y_ref(
|
||||
{dim_out[0], dim_out[1], dim_out[2], dim_out[3]},
|
||||
{stride_dim_out[0], stride_dim_out[1], stride_dim_out[2], stride_dim_out[3]});
|
||||
|
||||
ck_tile::reference_batched_transpose<Type>(x_host, y_ref, layout_in, layout_out);
|
||||
|
||||
auto [rtol, atol] = get_elimit<Type>("");
|
||||
|
||||
rtn &= ck_tile::check_err(
|
||||
y_host, y_ref, std::string("y Error: Incorrect results!"), rtol, atol);
|
||||
}
|
||||
printf("valid:%s\n", rtn ? "y" : "n");
|
||||
fflush(stdout);
|
||||
return rtn;
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
auto [result, args] = create_args(argc, argv);
|
||||
if(!result)
|
||||
return -1;
|
||||
std::string prec = args.get_str("pr");
|
||||
|
||||
bool r = true;
|
||||
if(prec.compare("fp16") == 0)
|
||||
{
|
||||
r &= run_batched_transpose<ck_tile::fp16_t>(args);
|
||||
}
|
||||
else if(prec.compare("fp8") == 0)
|
||||
{
|
||||
r &= run_batched_transpose<ck_tile::fp8_t>(args);
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cerr << "Unsupported data type: " << prec << std::endl;
|
||||
}
|
||||
|
||||
return r ? 0 : -1;
|
||||
}
|
||||
27
example/ck_tile/37_transpose/transpose_example.hpp
Normal file
27
example/ck_tile/37_transpose/transpose_example.hpp
Normal file
@@ -0,0 +1,27 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "ck_tile/ops/reduce.hpp"
|
||||
#include "batched_transpose_kernel.hpp"
|
||||
#include "block_transpose.hpp"
|
||||
#include "transpose_policy.hpp"
|
||||
|
||||
#include <vector>
|
||||
#include <string>
|
||||
|
||||
#pragma once
|
||||
|
||||
struct batched_transpose_trait
|
||||
{
|
||||
std::string type;
|
||||
std::string layout;
|
||||
};
|
||||
|
||||
struct batched_transpose_kargs : public ck_tile::BatchedTransposeHostArgs
|
||||
{
|
||||
};
|
||||
|
||||
float batched_transpose(batched_transpose_trait t,
|
||||
batched_transpose_kargs a,
|
||||
ck_tile::stream_config s);
|
||||
151
example/ck_tile/37_transpose/transpose_policy.hpp
Normal file
151
example/ck_tile/37_transpose/transpose_policy.hpp
Normal file
@@ -0,0 +1,151 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "ck_tile/core.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
struct TransposePolicy
|
||||
{
|
||||
static constexpr auto TileAccessPattern = tile_distribution_pattern::thread_raked;
|
||||
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto GetVectorSize()
|
||||
{
|
||||
return 16 / sizeof(typename Problem::DataType);
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
|
||||
{
|
||||
return integer_least_multiple(
|
||||
sizeof(typename Problem::DataType) *
|
||||
MakeLdsStoreBlockDescriptor<Problem>().get_element_space_size(),
|
||||
16);
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto MakeInputDistribution()
|
||||
{
|
||||
constexpr index_t BlockSize = Problem::kBlockSize;
|
||||
constexpr index_t LeadDimPerBlock = Problem::kLeadSizePerBlock;
|
||||
constexpr index_t SecondDimPerBlock = Problem::kSecondSizePerBlock;
|
||||
constexpr index_t VecLoadSize = 16 / sizeof(typename Problem::DataType);
|
||||
|
||||
using TileEncodingPattern = TileDistributionEncodingPattern2D<BlockSize,
|
||||
SecondDimPerBlock,
|
||||
LeadDimPerBlock,
|
||||
VecLoadSize,
|
||||
TileAccessPattern>;
|
||||
return TileEncodingPattern::Make2DStaticTileDistribution();
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto MakeOutputDistribution()
|
||||
{
|
||||
constexpr auto input_dstr = MakeLdsLoadTileDistribution<Problem>();
|
||||
|
||||
using OutTileDstrEncode =
|
||||
typename OutputTileDistributionTraits<remove_cvref_t<decltype(input_dstr)>,
|
||||
typename Problem::DataType>::OutDstrEncode;
|
||||
constexpr auto block_dstr = make_static_tile_distribution(OutTileDstrEncode{});
|
||||
|
||||
return block_dstr;
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto MakeLdsStoreBlockDescriptor()
|
||||
{
|
||||
constexpr index_t kLeadDimPerBlock = Problem::kLeadSizePerBlock;
|
||||
constexpr index_t kSecondDimPerBlock = Problem::kSecondSizePerBlock;
|
||||
constexpr index_t kVectorSize = 16 / sizeof(typename Problem::DataType);
|
||||
|
||||
constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(
|
||||
make_tuple(number<kSecondDimPerBlock>{},
|
||||
number<kLeadDimPerBlock / kVectorSize>{},
|
||||
number<kVectorSize>{}),
|
||||
make_tuple(number<kLeadDimPerBlock>{}, number<kVectorSize>{}, number<1>{}),
|
||||
number<kVectorSize>{},
|
||||
number<1>{});
|
||||
|
||||
constexpr auto lds_block_desc = transform_tensor_descriptor(
|
||||
lds_block_desc_0,
|
||||
make_tuple(make_pass_through_transform(number<kSecondDimPerBlock>{}),
|
||||
make_merge_transform(make_tuple(number<kLeadDimPerBlock / kVectorSize>{},
|
||||
number<kVectorSize>{}))),
|
||||
make_tuple(sequence<0>{}, sequence<1, 2>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}));
|
||||
|
||||
return lds_block_desc;
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto MakeLdsLoadBlockDescriptor()
|
||||
{
|
||||
constexpr index_t kLeadDimPerBlock = Problem::kLeadSizePerBlock;
|
||||
constexpr index_t kSecondDimPerBlock = Problem::kSecondSizePerBlock;
|
||||
|
||||
constexpr index_t kVectorSize = 8 / sizeof(typename Problem::DataType);
|
||||
|
||||
constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(
|
||||
make_tuple(number<kSecondDimPerBlock>{},
|
||||
number<kLeadDimPerBlock / kVectorSize>{},
|
||||
number<kVectorSize>{}),
|
||||
make_tuple(number<kLeadDimPerBlock>{}, number<kVectorSize>{}, number<1>{}),
|
||||
number<kVectorSize>{},
|
||||
number<1>{});
|
||||
|
||||
constexpr auto lds_block_desc = transform_tensor_descriptor(
|
||||
lds_block_desc_0,
|
||||
make_tuple(make_pass_through_transform(number<kSecondDimPerBlock>{}),
|
||||
make_merge_transform(make_tuple(number<kLeadDimPerBlock / kVectorSize>{},
|
||||
number<kVectorSize>{}))),
|
||||
make_tuple(sequence<0>{}, sequence<1, 2>{}),
|
||||
make_tuple(sequence<0>{}, sequence<1>{}));
|
||||
|
||||
return lds_block_desc;
|
||||
}
|
||||
|
||||
template <typename Problem>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto MakeLdsLoadTileDistribution()
|
||||
{
|
||||
using DataType = typename Problem::DataType;
|
||||
|
||||
// Extract base dimensions from the traits
|
||||
constexpr index_t kBaseLeadDim = LaneGroupTransposeTraits<DataType>::kleadDim;
|
||||
constexpr index_t kBaseSecondDim = LaneGroupTransposeTraits<DataType>::ksecondDim;
|
||||
|
||||
// Calculate block-level dimensions
|
||||
constexpr index_t kLead = Problem::kLeadSizePerXdl;
|
||||
constexpr index_t kSecond = Problem::kSecondSizePerXdl;
|
||||
constexpr index_t kLeadIterPerWarp = Problem::kLeadXdlNumPerWarp;
|
||||
constexpr index_t kSecondIterPerWarp = Problem::kSecondXdlNumPerWarp;
|
||||
constexpr index_t kLeadNumWarps = Problem::kLeadNumWarps;
|
||||
constexpr index_t kSecondNumWarps = Problem::kSecondNumWarps;
|
||||
|
||||
// Calculate repetitions of base pattern
|
||||
constexpr index_t kLeadRepetitions = kLead / kBaseLeadDim;
|
||||
constexpr index_t kSecondRepetitions = kSecond / kBaseSecondDim;
|
||||
constexpr index_t kSecondDimIterations = Problem::kIterationsInSecondDim;
|
||||
constexpr index_t kSecondDimStrSub = kSecondRepetitions / kSecondDimIterations;
|
||||
|
||||
constexpr auto xdllevel_dstr_encoding = make_transposed_distr_encode<DataType,
|
||||
kSecondDimStrSub,
|
||||
kSecondDimIterations,
|
||||
kLeadRepetitions,
|
||||
1>();
|
||||
|
||||
constexpr auto input_tile_encode =
|
||||
InputTileDistributionEncoding<decltype(xdllevel_dstr_encoding),
|
||||
kLeadIterPerWarp,
|
||||
kSecondIterPerWarp,
|
||||
kLeadNumWarps,
|
||||
kSecondNumWarps>();
|
||||
constexpr auto block_dstr = make_static_tile_distribution(input_tile_encode);
|
||||
return block_dstr;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
Reference in New Issue
Block a user