Skip to content

GitLab

Unverified Commit 6491acda authored by carlushuang's avatar carlushuang Committed by GitHub
Browse files

add tensor slicing API (#7) 



* add tensor slicing API

* remove redundant ck namespace

* better gemm_gemm interface

* modify gemm_gemm

* add slice_tile api

* fix merge bug

* update to 3d padding, since we no longer need that much LDS size

* clean

* cleang

* clean

* clean

* clean

* clean

* clean

* clean

* clean

* clean

* clean

* clean

* clean

* clean

---------
Co-authored-by: default avatarChao Liu <chao.liu2@amd.com>
parent 1cf54e86
Hide whitespace changes
Inline Side-by-side
Showing with 583 additions and 55 deletions
example/91_tile_program/gemm_gemm.cpp
View file @ 6491acda
...@@ -84,6 +84,7 @@ int main(int argc, char* argv[]) ...@@ -84,6 +84,7 @@ int main(int argc, char* argv[])
constexpr ck::index_t kN0PerBlock = 128; constexpr ck::index_t kN0PerBlock = 128;
constexpr ck::index_t kK0PerBlock = 32; constexpr ck::index_t kK0PerBlock = 32;
constexpr ck::index_t kN1PerBlock = 128; constexpr ck::index_t kN1PerBlock = 128;
constexpr ck::index_t kK1PerBlock = 32;
constexpr ck::index_t kBlockSize = 256; constexpr ck::index_t kBlockSize = 256;
ck::index_t kGridSize = (M0 / kM0PerBlock) * (N1 / kN1PerBlock); ck::index_t kGridSize = (M0 / kM0PerBlock) * (N1 / kN1PerBlock);
...@@ -107,7 +108,8 @@ int main(int argc, char* argv[]) ...@@ -107,7 +108,8 @@ int main(int argc, char* argv[])
kM0PerBlock, kM0PerBlock,
kN0PerBlock, kN0PerBlock,
kK0PerBlock, kK0PerBlock,
kN1PerBlock>{}, kN1PerBlock,
kK1PerBlock>{},
kGridSize, kGridSize,
kBlockSize, kBlockSize,
0, 0,
......
example/91_tile_program/gemm_gemm.hpp
View file @ 6491acda
...@@ -11,6 +11,7 @@ ...@@ -11,6 +11,7 @@
#include "ck/tile_program/tile/tile_distribution.hpp" #include "ck/tile_program/tile/tile_distribution.hpp"
#include "ck/tile_program/tile/tile_elementwise.hpp" #include "ck/tile_program/tile/tile_elementwise.hpp"
#include "ck/tile_program/tile/tile_gemm_shape.hpp" #include "ck/tile_program/tile/tile_gemm_shape.hpp"
#include "ck/tile_program/tile/slice_tile.hpp"
#include "ck/tile_program/warp_tile/warp_gemm.hpp" #include "ck/tile_program/warp_tile/warp_gemm.hpp"
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp" #include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp"
#include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_problem.hpp" #include "ck/tile_program/block_tile_pipeline/block_gemm_pipeline_problem.hpp"
...@@ -29,9 +30,18 @@ template <typename A0DataType, ...@@ -29,9 +30,18 @@ template <typename A0DataType,
ck::index_t kM0PerBlock, ck::index_t kM0PerBlock,
ck::index_t kN0PerBlock, ck::index_t kN0PerBlock,
ck::index_t kK0PerBlock, ck::index_t kK0PerBlock,
ck::index_t kN1PerBlock> ck::index_t kN1PerBlock,
ck::index_t kK1PerBlock>
struct GemmGemm struct GemmGemm
{ {
static constexpr auto I0 = ck::Number<0>{};
static constexpr auto BlockSize = ck::Number<kBlockSize>{};
static constexpr auto M0PerBlock = ck::Number<kM0PerBlock>{};
static constexpr auto N0PerBlock = ck::Number<kN0PerBlock>{};
static constexpr auto K0PerBlock = ck::Number<kK0PerBlock>{};
static constexpr auto N1PerBlock = ck::Number<kN1PerBlock>{};
static constexpr auto K1PerBlock = ck::Number<kK1PerBlock>{};
// block gemm0 pipeline // block gemm0 pipeline
using BlockGemm0Pipeline = ck::tile_program::block::BlockGemmPipelineAGmemBGmemCRegV2< using BlockGemm0Pipeline = ck::tile_program::block::BlockGemmPipelineAGmemBGmemCRegV2<
ck::tile_program::block::BlockGemmPipelineProblem< ck::tile_program::block::BlockGemmPipelineProblem<
...@@ -49,7 +59,7 @@ struct GemmGemm ...@@ -49,7 +59,7 @@ struct GemmGemm
B1DataType, B1DataType,
Acc1DataType, Acc1DataType,
kBlockSize, kBlockSize,
ck::tile_program::TileGemmShape<kM0PerBlock, kN1PerBlock, kN0PerBlock>>, ck::tile_program::TileGemmShape<kM0PerBlock, kN1PerBlock, kK1PerBlock>>,
ck::tile_program::block::BlockGemmARegBSmemCRegV1DefaultPolicy>; ck::tile_program::block::BlockGemmARegBSmemCRegV1DefaultPolicy>;
#if 0 #if 0
...@@ -59,23 +69,51 @@ struct GemmGemm ...@@ -59,23 +69,51 @@ struct GemmGemm
using namespace ck; using namespace ck;
constexpr index_t kNPerBlock = kN1PerBlock; constexpr index_t kNPerBlock = kN1PerBlock;
constexpr index_t kKPerBlock = kN0PerBlock; constexpr index_t kKPerBlock = kK1PerBlock;
constexpr auto b_lds_block_desc = constexpr auto b_lds_block_desc =
make_naive_tensor_descriptor_packed(make_tuple(kNPerBlock, kKPerBlock), Number<32>{}); make_naive_tensor_descriptor_packed(make_tuple(kNPerBlock, kKPerBlock), Number<32>{});
return b_lds_block_desc;
}
#elif 1
// 3d, with padding
__device__ static constexpr auto MakeB1LdsBlockDescriptor()
{
using namespace ck;
// using BDataType = B1DataType;
constexpr index_t kNPerBlock = kN1PerBlock;
constexpr index_t kKPerBlock = kK1PerBlock;
constexpr index_t kPad = 1;
constexpr index_t kK1 = 8;
constexpr auto b_lds_block_desc_0 = make_naive_tensor_descriptor(
make_tuple(Number<kKPerBlock / kK1>{}, Number<kNPerBlock>{}, Number<kK1>{}),
make_tuple(Number<(kNPerBlock + kPad) * kK1>{}, Number<kK1>{}, Number<1>{}),
Number<kK1>{},
Number<1>{});
constexpr auto b_lds_block_desc = transform_tensor_descriptor(
b_lds_block_desc_0,
make_tuple(make_pass_through_transform(kNPerBlock),
make_merge_transform(make_tuple(Number<kKPerBlock / kK1>{}, Number<kK1>{}))),
make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
make_tuple(Sequence<0>{}, Sequence<1>{}));
return b_lds_block_desc; return b_lds_block_desc;
} }
#else #else
// fake XOR // fake XOR
__device__ static constexpr auto MakeB1LdsBlockDescriptor() __host__ __device__ static constexpr auto MakeB1LdsBlockDescriptor()
{ {
using namespace ck; using namespace ck;
using BDataType = B1DataType; using BDataType = B1DataType;
constexpr index_t kNPerBlock = kN1PerBlock; constexpr index_t kNPerBlock = kN1PerBlock;
constexpr index_t kKPerBlock = kN0PerBlock; constexpr index_t kKPerBlock = kK1PerBlock;
constexpr auto b_lds_block_desc_d1_d2_d3 = make_naive_tensor_descriptor_packed( constexpr auto b_lds_block_desc_d1_d2_d3 = make_naive_tensor_descriptor_packed(
make_tuple(kNPerBlock / 2, 2, kKPerBlock), Number<kKPerBlock>{}); make_tuple(kNPerBlock / 2, 2, kKPerBlock), Number<kKPerBlock>{});
...@@ -108,7 +146,7 @@ struct GemmGemm ...@@ -108,7 +146,7 @@ struct GemmGemm
using BDataType = B1DataType; using BDataType = B1DataType;
constexpr index_t kNPerBlock = kN1PerBlock; constexpr index_t kNPerBlock = kN1PerBlock;
constexpr index_t kKPerBlock = kN0PerBlock; constexpr index_t kKPerBlock = kK1PerBlock;
constexpr index_t K1 = 16 / sizeof(BDataType); constexpr index_t K1 = 16 / sizeof(BDataType);
constexpr index_t K0 = kKPerBlock / K1; constexpr index_t K0 = kKPerBlock / K1;
...@@ -138,14 +176,14 @@ struct GemmGemm ...@@ -138,14 +176,14 @@ struct GemmGemm
const B0DataType* p_b0, const B0DataType* p_b0,
const B1DataType* p_b1, const B1DataType* p_b1,
C1DataType* p_c1, C1DataType* p_c1,
ck::index_t M0, const ck::index_t M0,
ck::index_t N0, const ck::index_t N0,
ck::index_t K0, const ck::index_t K0,
ck::index_t N1, const ck::index_t N1,
ck::index_t Lda0, const ck::index_t Lda0,
ck::index_t Ldb0, const ck::index_t Ldb0,
ck::index_t Ldb1, const ck::index_t Ldb1,
ck::index_t Ldc1) const ck::index_t Ldc1)
{ {
using namespace ck; using namespace ck;
using namespace ck::tile_program; using namespace ck::tile_program;
...@@ -177,38 +215,40 @@ struct GemmGemm ...@@ -177,38 +215,40 @@ struct GemmGemm
__shared__ char p_smem_char[GetStaticLdsSize()]; __shared__ char p_smem_char[GetStaticLdsSize()];
// A0 DRAM block window // A0 DRAM block window
auto a0_dram_block_window = make_tile_window( auto a0_dram_block_window =
a0_dram_grid, make_tuple(Number<kM0PerBlock>{}, Number<kK0PerBlock>{}), {iM0, 0}); make_tile_window(a0_dram_grid, make_tuple(M0PerBlock, K0PerBlock), {iM0, 0});
// B0 DRAM block window // B0 DRAM block window
auto b0_dram_block_window = make_tile_window( auto b0_dram_block_window =
b0_dram_grid, make_tuple(Number<kN0PerBlock>{}, Number<kK0PerBlock>{}), {0, 0}); make_tile_window(b0_dram_grid, make_tuple(N0PerBlock, K0PerBlock), {0, 0});
// Block GEMM0 pipeline // Block GEMM0 pipeline
constexpr auto block_gemm0_pipeline = BlockGemm0Pipeline{}; constexpr auto block_gemm0_pipeline = BlockGemm0Pipeline{};
// B1 DRAM window // B1 DRAM window
auto b1_dram_block_window = auto b1_dram_block_window = make_tile_window(b1_dram_grid,
make_tile_window(b1_dram_grid, make_tuple(N1PerBlock, K1PerBlock),
make_tuple(Number<kN1PerBlock>{}, Number<kN0PerBlock>{}), {iN1, 0},
{iN1, 0}, MakeB1DramTileDistribution());
MakeB1DramTileDistribution());
// B1 LDS tensor view: occupies the same LDS allocation as block_gemm0_pipeline // B1 LDS tensor view: occupies the same LDS allocation as block_gemm0_pipeline
auto b1_lds_block = make_tensor_view<AddressSpaceEnum::Lds>( auto b1_lds_block = make_tensor_view<AddressSpaceEnum::Lds>(
reinterpret_cast<B1DataType*>(p_smem_char), MakeB1LdsBlockDescriptor()); reinterpret_cast<B1DataType*>(p_smem_char), MakeB1LdsBlockDescriptor());
auto b1_lds_block_window = make_tile_window( auto b1_lds_block_window =
b1_lds_block, make_tuple(Number<kN1PerBlock>{}, Number<kN0PerBlock>{}), {0, 0}); make_tile_window(b1_lds_block, make_tuple(N1PerBlock, K1PerBlock), {0, 0});
// Bock GEMM1 // Bock GEMM1
constexpr auto block_gemm1 = BlockGemm1{}; constexpr auto block_gemm1 = BlockGemm1{};
// Acc1 tile // Acc1 tile
auto acc1_block_tile = decltype(block_gemm1( auto acc1_block_tile = decltype(block_gemm1(
tile_elementwise_in( get_slice_tile(
type_convert<C0DataType, Acc0DataType>, tile_elementwise_in(
block_gemm0_pipeline(a0_dram_block_window, b0_dram_block_window, 0, nullptr)), type_convert<C0DataType, Acc0DataType>,
block_gemm0_pipeline(a0_dram_block_window, b0_dram_block_window, 0, nullptr)),
Sequence<0, 0>{},
Sequence<kM0PerBlock, kK1PerBlock>{}),
b1_dram_block_window)){}; b1_dram_block_window)){};
// init Acc1 // init Acc1
...@@ -226,30 +266,44 @@ struct GemmGemm ...@@ -226,30 +266,44 @@ struct GemmGemm
const auto c0_block_tile = const auto c0_block_tile =
tile_elementwise_in(type_convert<C0DataType, Acc0DataType>, acc0_block_tile); tile_elementwise_in(type_convert<C0DataType, Acc0DataType>, acc0_block_tile);
// Block GEMM1: acc1 += c0 * b1 // prefetch load b1
{ const auto b1_block_tile = load_tile(b1_dram_block_window);
// load b1 move_tile_window(b1_dram_block_window, {0, kK1PerBlock});
const auto b1_block_tile = load_tile(b1_dram_block_window);
// wait for block gemm0 pipeline to finish
block_sync_lds();
store_tile(b1_lds_block_window, b1_block_tile); block_sync_lds();
// wait for store_tile to finish store_tile(b1_lds_block_window, b1_block_tile);
block_sync_lds();
// acc1 += c0 * b1 constexpr index_t k1_loops = kN0PerBlock / kK1PerBlock;
block_gemm1(acc1_block_tile, c0_block_tile, b1_lds_block_window);
// wait for block gemm1 to finish if constexpr(k1_loops > 1)
{
static_for<0, k1_loops - 1, 1>{}([&](auto i) {
// acc1 += c0 * b1
const auto b1_block_tile_1 = load_tile(b1_dram_block_window);
block_sync_lds();
block_gemm1(acc1_block_tile,
get_slice_tile(c0_block_tile,
Sequence<0, i * kK1PerBlock>{},
Sequence<kM0PerBlock, (i + 1) * kK1PerBlock>{}),
b1_lds_block_window);
block_sync_lds();
move_tile_window(b1_dram_block_window, {0, kK1PerBlock});
store_tile(b1_lds_block_window, b1_block_tile_1);
});
}
// tail
{
block_sync_lds(); block_sync_lds();
block_gemm1(acc1_block_tile,
get_slice_tile(c0_block_tile,
Sequence<0, (k1_loops - 1) * kK1PerBlock>{},
Sequence<kM0PerBlock, kN0PerBlock>{}),
b1_lds_block_window);
} }
// move tile windows
move_tile_window(b0_dram_block_window, {kN0PerBlock, 0}); move_tile_window(b0_dram_block_window, {kN0PerBlock, 0});
move_tile_window(b1_dram_block_window, {0, kN0PerBlock}); block_sync_lds();
iN0 += kN0PerBlock; iN0 += kN0PerBlock;
} while(iN0 < N0); } while(iN0 < N0);
...@@ -262,11 +316,10 @@ struct GemmGemm ...@@ -262,11 +316,10 @@ struct GemmGemm
auto c1_dram_grid = make_naive_tensor_view<AddressSpaceEnum::Global>( auto c1_dram_grid = make_naive_tensor_view<AddressSpaceEnum::Global>(
p_c1, make_tuple(M0, N1), make_tuple(Ldc1, 1), Number<32>{}, Number<1>{}); p_c1, make_tuple(M0, N1), make_tuple(Ldc1, 1), Number<32>{}, Number<1>{});
auto c1_dram_window = auto c1_dram_window = make_tile_window(c1_dram_grid,
make_tile_window(c1_dram_grid, make_tuple(M0PerBlock, N1PerBlock),
make_tuple(Number<kM0PerBlock>{}, Number<kN1PerBlock>{}), {iM0, iN1},
{iM0, iN1}, c1_block_tile.GetTileDistribution());
c1_block_tile.GetTileDistribution());
store_tile(c1_dram_window, c1_block_tile); store_tile(c1_dram_window, c1_block_tile);
} }
......
include/ck/tile_program/tile/slice_tile.hpp 0 → 100644
View file @ 6491acda
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck/utility/common_header.hpp"
#include "ck/tensor_description/tensor_descriptor.hpp"
#include "ck/tensor_description/tensor_descriptor_helper.hpp"
#include "ck/tensor_description/tensor_adaptor.hpp"
#include "ck/tensor_description/tensor_space_filling_curve.hpp"
#include "ck/tile_program/tile/tile_distribution.hpp"
#include "ck/tile_program/tile/tile_window.hpp"
#include "ck/tile_program/tile/static_distributed_tensor.hpp"
namespace ck {
namespace tile_program {
namespace detail {
template <typename, typename, typename, index_t>
struct reverse_slice_sequence_impl;
template <index_t x,
index_t... xs,
index_t m,
index_t... ms,
index_t id,
index_t... ids,
index_t SliceSize>
struct reverse_slice_sequence_impl<Sequence<x, xs...>,
Sequence<m, ms...>,
Sequence<id, ids...>,
SliceSize>
{
using old_scan =
reverse_slice_sequence_impl<Sequence<xs...>, Sequence<ms...>, Sequence<ids...>, SliceSize>;
static constexpr auto slice_size = old_scan::remaining_slice_sizes::Front().value;
static constexpr auto slice_length =
std::conditional_t<m, Number<math::gcd(x, slice_size)>, Number<x>>::value;
using dim_lengths =
typename sequence_merge<Sequence<slice_length>, typename old_scan::dim_lengths>::type;
using dim_slices =
typename sequence_merge<Sequence<x / slice_length>, typename old_scan::dim_slices>::type;
using remaining_slice_sizes = typename sequence_merge<
std::conditional_t<m, Sequence<slice_size / slice_length>, Sequence<slice_size>>,
typename old_scan::remaining_slice_sizes>::type;
// the first idx that sliced length not equal to original length
static constexpr index_t _flag =
slice_length != x && remaining_slice_sizes{}.Front().value == 1;
static constexpr index_t _split_flag = std::conditional_t<m, Number<_flag>, Number<0>>::value;
static constexpr index_t _split_idx =
std::conditional_t<_split_flag, Number<id>, Number<0>>::value;
static constexpr index_t split_flag = _split_flag || old_scan::split_flag;
static constexpr index_t split_idx = std::
conditional_t<old_scan::split_flag, Number<old_scan::split_idx>, Number<_split_idx>>::value;
};
template <index_t x, index_t m, index_t id, index_t SliceSize>
struct reverse_slice_sequence_impl<Sequence<x>, Sequence<m>, Sequence<id>, SliceSize>
{
static constexpr auto slice_size = SliceSize;
static constexpr auto slice_length =
std::conditional_t<m, Number<math::gcd(x, slice_size)>, Number<x>>::value;
using dim_lengths = Sequence<slice_length>;
using dim_slices = Sequence<x / slice_length>;
using remaining_slice_sizes =
std::conditional_t<m, Sequence<slice_size / slice_length>, Sequence<slice_size>>;
// the first idx that sliced length not equal to original length
static constexpr index_t _flag =
slice_length != x && remaining_slice_sizes{}.Front().value == 1;
static constexpr index_t split_flag = std::conditional_t<m, Number<_flag>, Number<0>>::value;
static constexpr index_t split_idx =
std::conditional_t<split_flag, Number<id>, Number<0>>::value;
};
// clang-format off
// input a sequence(with optional mask), and the SliceSize : size per slice
// output the sequence each slice, and Number of slices
//
// e.g. <2, 1, 4, 2>, 8 -> lengths:<1, 1, 4, 2> , nums: <2, 1, 1, 1> : 2 slices , slice_idx: 0
// <4, 2, 4, 1, 2>, 4 -> lengths:<1, 1, 2, 1, 2> , nums: <4, 2, 2, 1, 1> : 16 slices , slice_idx: 2
// <4, 2, 4, 1, 6>, 4 -> lengths:<1, 1, 2, 1, 2> , nums: <4, 2, 2, 1, 3> : 48 slices , slice_idx: 2
// <4, 2, 5, 1, 2>, 10 -> lengths:<1, 1, 5, 1, 2> , nums: <4, 2, 1, 1, 1> : 8 slices , slice_idx: 1
//
// <4, 2, 8>, 64 -> lengths:<4, 2, 8> , nums: <1, 1, 1> : 1 slices , slice_idx: 0
// <4, 2, 8>, 32 -> lengths:<2, 2, 8> , nums: <2, 1, 1> : 2 slices , slice_idx: 0
// <4, 2, 8>, 16 -> lengths:<1, 2, 8> , nums: <4, 1, 1> : 4 slices , slice_idx: 0
// <4, 2, 8>, 8 -> lengths:<1, 1, 8> , nums: <4, 2, 1> : 8 slices , slice_idx: 1
// <4, 2, 8>, 4 -> lengths:<1, 1, 4> , nums: <4, 2, 2> : 16 slices , slice_idx: 2
// <4, 2, 8>, 2 -> lengths:<1, 1, 2> , nums: <4, 2, 4> : 32 slices , slice_idx: 2
// <4, 2, 8>, 1 -> lengths:<1, 1, 1> , nums: <4, 2, 8> : 64 slices , slice_idx: 2
//
// <4, 2, 1, 4, 2> / 4 ->
// mask:<1, 1, 1, 0, 1>, -> lengths:<1, 2, 1, 4, 2> , nums: <4, 1, 1, 1, 1> : 8 slices , slice_idx: 0
//
// return Tuple<slice_lengths, slice_nums, slice_index>, slice_index is at which index will start
// have split slices (right -> left)
// or the first index that sliced length is different from the original length
// clang-format on
template <typename Seq,
index_t SliceSize,
typename Mask = typename uniform_sequence_gen<Seq::Size(), 1>::type>
constexpr auto reverse_slice_sequence(Seq,
Number<SliceSize>,
Mask = typename uniform_sequence_gen<Seq::Size(), 1>::type{})
{
static_assert(Seq::Size() == Mask::Size());
using sliced_type =
reverse_slice_sequence_impl<Seq,
Mask,
typename arithmetic_sequence_gen<0, Seq::Size(), 1>::type,
SliceSize>;
static_assert(sliced_type::remaining_slice_sizes::Front().value == 1,
"can not evenly divide this sequence, please check");
return make_tuple(typename sliced_type::dim_lengths{},
typename sliced_type::dim_slices{},
Number<sliced_type::split_idx>{});
}
//
// slice tensor from x_dim, result in split in y_dim, not p_dim.
// We don't support slice cross p_dim (aka, slice different threads)
// also, sliced along y_dim need be the first dim of current dim.
// Multiply Y dim before sliced dim does not make sense
//
// e.g
// X0 X1
// <1, 4, 32> - <4, 1, 4, 2, 4> | slice origin:<0, 0>, len:<0, 32>, (0 means all length)
// Y P P Y P Y P Y
// => <1, 4, 32> - <1, 1, 4, 2, 4> -> OK
// |--> slice along this Y dim, is the first dim of X1, totally 4 slices
//
// X0 X1
// <1, 4, 32> - <4, 1, 4, 2, 4> | slice origin:<0, 0>, len:<0, 8>, (0 means all length)
// Y P P Y P Y P Y
// => <1, 4, 32> - <1, 1, 1, 2, 4> -> OK
// |--> slice along this Y dim, the P dim is 1 in the left, so is OK
// totally 16 slices
//
// X0 X1
// <1, 4, 32> - <4, 1, 4, 2, 4> | slice origin:<0, 0>, len:<0, 4>, (0 means all length)
// Y P P Y P Y P Y
// => <1, 4, 32> - <1, 1, 1, 1, 4> -> Fail
// |--> slice along this P dim, will split threads, not supported
//
// X0 X1
// <1, 4, 32> - <4, 1, 4, 2, 4> | slice origin:<0, 0>, len:<0, 16>, (0 means all length)
// Y P P Y P Y P Y
// => <1, 4, 32> - <1, 1, 2, 2, 4> -> OK
// |--> slice along this Y dim, but this Y sim need to split into 2
// subdime
// the P dim in the left is 1, means actually not crossing P
//
template <typename Distribution, index_t... XSliceBegins, index_t... XSliceEnds>
__host__ __device__ constexpr auto slice_distribution_from_x(
Distribution, Sequence<XSliceBegins...> x_slice_begins, Sequence<XSliceEnds...> x_slice_ends)
{
// NOTE: this function need to be called under constexpr context,
// due to https://wg21.link/p2280r0 we have to use non-reference type for distribution
using Encoding = decltype(Distribution::GetStaticTileDistributionEncoding());
static_assert(sizeof...(XSliceBegins) == sizeof...(XSliceEnds));
constexpr auto x_slice_lengths = x_slice_ends - x_slice_begins;
constexpr auto src_h_prefix_sum = Encoding::Detail::GetHDimLengthsPrefixSum();
constexpr auto src_y_info = Encoding::Detail::GetSortedYInfo();
constexpr auto src_y_dims = src_y_info[Number<0>{}];
constexpr auto src_y_maps = src_y_info[Number<1>{}];
constexpr auto src_y_prefix_sum = src_y_info[Number<2>{}];
constexpr auto sliced_hlen_yidx_ylen = [&]() constexpr
{
auto y_slice_sorted_origins = make_zero_multi_index<Distribution::NDimY>();
auto y_slice_lengths =
to_array<index_t, Distribution::NDimY>(Distribution{}.GetYs2DDescriptor().GetLengths());
// This lambda will modify some value outside, so c++ will not treat return value as
// constexpr
// TODO: ugly
auto new_h_lengths = transform_tuples(
[&](auto h_len, auto id) {
constexpr auto sliced_h =
reverse_slice_sequence(h_len, Number<x_slice_lengths[id]>{});
constexpr auto sliced_h_lens = sliced_h[Number<0>{}];
constexpr auto sliced_h_index = sliced_h[Number<2>{}];
// update y_slice_lengths
constexpr auto uniformed_h_index = sliced_h_index + Number<src_h_prefix_sum[id]>{};
constexpr auto found_y_index = container_find(src_y_dims, uniformed_h_index);
static_assert(found_y_index >= 0 && found_y_index < src_y_dims.Size(),
"not sliced at y dim, please check");
static_for<0, sliced_h_index + 1, 1>{}([&](auto i) {
y_slice_lengths(src_y_maps[found_y_index - i]) =
sliced_h_lens[sliced_h_index - i];
});
// TODO: add validations not across p dim
// NOTE: this y_origin is for all dims, not only current dim
// will later use pick to select target dim
constexpr auto y_origin = [&]() {
constexpr auto h_trans = make_merge_transform_v3_division_mod(h_len);
auto h_origin_ = make_zero_multi_index<h_trans.NDimLow>();
h_trans.CalculateLowerIndex(h_origin_, Sequence<x_slice_begins[id].value>{});
auto y_origin_ = make_zero_multi_index<Distribution::NDimY>();
static_for<0, sliced_h_index + 1, 1>{}([&](auto i) {
y_origin_(found_y_index - i) = h_origin_[sliced_h_index - i];
});
return y_origin_;
}();
constexpr auto y_picks = typename arithmetic_sequence_gen<src_y_prefix_sum[id],
src_y_prefix_sum[id + 1],
1>::type{};
set_container_subset(
y_slice_sorted_origins, y_picks, get_container_subset(y_origin, y_picks));
return sliced_h_lens;
},
typename Encoding::HsLengthss{},
typename arithmetic_sequence_gen<0, Encoding::HsLengthss::Size(), 1>::type{});
auto y_slice_origins = container_reorder_given_old2new(y_slice_sorted_origins, src_y_maps);
return make_tuple(new_h_lengths, y_slice_origins, y_slice_lengths);
}
();
return sliced_hlen_yidx_ylen;
}
} // namespace detail
template <typename StaticDistributedTensor_, index_t... SliceBegins, index_t... SliceEnds>
__host__ __device__ constexpr auto get_slice_tile(const StaticDistributedTensor_& tile,
Sequence<SliceBegins...> slice_begins,
Sequence<SliceEnds...> slice_ends)
{
using Distribution = decltype(StaticDistributedTensor_::GetTileDistribution());
using Encoding = decltype(Distribution::GetStaticTileDistributionEncoding());
using DataType = typename StaticDistributedTensor_::DataType;
constexpr auto sliced_hlen_yidx_ylen =
detail::slice_distribution_from_x(Distribution{}, slice_begins, slice_ends);
constexpr auto sliced_h_lengths = sliced_hlen_yidx_ylen[Number<0>{}];
constexpr auto sliced_y_origins_array = sliced_hlen_yidx_ylen[Number<1>{}];
constexpr auto sliced_y_origins_size = sliced_y_origins_array.Size();
constexpr auto sliced_y_lengths_array = sliced_hlen_yidx_ylen[Number<2>{}];
constexpr auto sliced_y_lengths_size = sliced_y_lengths_array.Size();
constexpr auto sliced_y_origins = TO_SEQUENCE(sliced_y_origins_array, sliced_y_origins_size);
constexpr auto sliced_y_lengths = TO_SEQUENCE(sliced_y_lengths_array, sliced_y_lengths_size);
using SlicedEnc =
StaticTileDistributionEncoding<typename Encoding::RsLengths,
decltype(sliced_h_lengths), // only need to change the
// h_lengths type
typename Encoding::Ps2RHssMajor,
typename Encoding::Ps2RHssMinor,
typename Encoding::Ys2RHsMajor,
typename Encoding::Ys2RHsMinor>;
auto sliced_tensor =
make_static_distributed_tensor<DataType>(make_static_tile_distribution(SlicedEnc{}));
sliced_tensor.GetThreadBuffer() = tile.GetSlicedThreadData(sliced_y_origins, sliced_y_lengths);
return sliced_tensor;
}
template <typename DstStaticDistributedTensor_,
typename SrcStaticDistributedTensor_,
index_t... SliceBegins,
index_t... SliceEnds>
__host__ __device__ constexpr auto set_slice_tile(DstStaticDistributedTensor_& dst_tile,
const SrcStaticDistributedTensor_& src_tile,
Sequence<SliceBegins...> slice_begins,
Sequence<SliceEnds...> slice_ends)
{
using DstDistribution = decltype(DstStaticDistributedTensor_::GetTileDistribution());
constexpr auto sliced_hlen_yidx_ylen =
detail::slice_distribution_from_x(DstDistribution{}, slice_begins, slice_ends);
constexpr auto sliced_h_lengths = sliced_hlen_yidx_ylen[Number<0>{}];
constexpr auto sliced_y_origins_array = sliced_hlen_yidx_ylen[Number<1>{}];
constexpr auto sliced_y_origins_size = sliced_y_origins_array.Size();
constexpr auto sliced_y_lengths_array = sliced_hlen_yidx_ylen[Number<2>{}];
constexpr auto sliced_y_lengths_size = sliced_y_lengths_array.Size();
constexpr auto sliced_y_origins = TO_SEQUENCE(sliced_y_origins_array, sliced_y_origins_size);
constexpr auto sliced_y_lengths = TO_SEQUENCE(sliced_y_lengths_array, sliced_y_lengths_size);
dst_tile.SetSlicedThreadData(sliced_y_origins, sliced_y_lengths, src_tile.GetThreadBuffer());
}
} // namespace tile_program
} // namespace ck
include/ck/tile_program/tile/static_tile_distribution_encoding.hpp
View file @ 6491acda
...@@ -41,6 +41,7 @@ struct StaticTileDistributionEncoding ...@@ -41,6 +41,7 @@ struct StaticTileDistributionEncoding
static constexpr auto ys_to_rhs_minor_ = Ys2RHsMinor{}; static constexpr auto ys_to_rhs_minor_ = Ys2RHsMinor{};
// redundant but useful info // redundant but useful info
// TODO: really bad code, should be over-hauled
struct Detail struct Detail
{ {
// ndim_rh_major_, ndim_span_mainor_ // ndim_rh_major_, ndim_span_mainor_
...@@ -232,6 +233,62 @@ struct StaticTileDistributionEncoding ...@@ -232,6 +233,62 @@ struct StaticTileDistributionEncoding
} }
}(); }();
// e.g. tuple<seq<1, 4, 32>, seq<4, 1, 4, 2, 4>> --> seq<3, 5> --> seq<0, 3, 8>
__host__ __device__ static constexpr auto GetHDimLengthsPrefixSum()
{
// <len_d0, len_d1, ...>
// e.g. tuple<seq<1, 4, 32>, seq<4, 1, 4, 2, 4>> --> seq<3, 5>
constexpr auto uniformed_h_dim_lengths = generate_sequence_v2(
[&](auto i) {
constexpr index_t size = HsLengthss{}[i].Size();
return Number<size>{};
},
Number<NDimX>{});
// <0, len_d0, len_d0+len_d1, ...>
// e.g. seq<3, 5> --> seq<0, 3, 8>
constexpr auto h_dim_prefix_sum = prefix_sum_sequence(uniformed_h_dim_lengths);
return h_dim_prefix_sum;
}
__host__ __device__ static constexpr auto GetUniformedIdxY2H()
{
constexpr auto all_ys_2_rhss = transform_sequences(
[](auto major, auto minor) constexpr {
// <0, 0, len_d0, len_d0+len_d1, ...>
constexpr auto x_dim_prefix_sum =
merge_sequences(Sequence<0>{} /*for R dims*/, GetHDimLengthsPrefixSum());
return x_dim_prefix_sum.At(major) + minor;
},
Ys2RHsMajor{},
Ys2RHsMinor{});
return all_ys_2_rhss;
}
// return tuple<sorted_dims, sorted_maps, sorted_prefix_sum>
template <typename IdxSeq, typename PrefixSumSeq>
__host__ __device__ static constexpr auto GetSortedInfo(IdxSeq, PrefixSumSeq)
{
using sorted_idx =
sequence_unique_sort<IdxSeq, math::less<index_t>, math::equal<index_t>>;
constexpr auto sorted_dims = typename sorted_idx::type{};
constexpr auto sorted_maps = typename sorted_idx::sorted2unsorted_map{};
constexpr auto sorted_histogram =
histogram_sorted_sequence(sorted_dims, PrefixSumSeq{});
constexpr auto sorted_prefix_sum = prefix_sum_sequence(sorted_histogram);
return make_tuple(sorted_dims, sorted_maps, sorted_prefix_sum);
}
__host__ __device__ static constexpr auto GetSortedYInfo()
{
return GetSortedInfo(GetUniformedIdxY2H(), GetHDimLengthsPrefixSum());
}
__host__ __device__ void Print() const __host__ __device__ void Print() const
{ {
printf("StaticTileDistributionEncoding::Detail{"); printf("StaticTileDistributionEncoding::Detail{");
......
include/ck/utility/container_helper.hpp
View file @ 6491acda
...@@ -447,6 +447,20 @@ __host__ __device__ constexpr void set_container_subset(Y& y, Sequence<Is...> pi ...@@ -447,6 +447,20 @@ __host__ __device__ constexpr void set_container_subset(Y& y, Sequence<Is...> pi
} }
} }
// return the index of first occurance in the sequence.
// return seq.Size(), if not found
template <index_t... Is>
constexpr index_t container_find(Sequence<Is...> seq, index_t value)
{
for(auto i = 0; i < seq.Size(); i++)
{
if(seq[i] == value)
return i;
}
return seq.Size();
}
template <index_t... Is> template <index_t... Is>
__host__ __device__ constexpr auto sequence_to_tuple_of_number(Sequence<Is...>) __host__ __device__ constexpr auto sequence_to_tuple_of_number(Sequence<Is...>)
{ {
......
include/ck/utility/number.hpp
View file @ 6491acda
// SPDX-License-Identifier: MIT // SPDX-License-Identifier: MIT
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
#ifndef CK_NUMBER_HPP #pragma once
#define CK_NUMBER_HPP
#include "integral_constant.hpp" #include "ck/utility/integral_constant.hpp"
namespace ck { namespace ck {
...@@ -15,4 +14,3 @@ template <index_t N> ...@@ -15,4 +14,3 @@ template <index_t N>
using LongNumber = integral_constant<long_index_t, N>; using LongNumber = integral_constant<long_index_t, N>;
} // namespace ck } // namespace ck
#endif
include/ck/utility/sequence.hpp
View file @ 6491acda
...@@ -809,6 +809,46 @@ __host__ __device__ constexpr auto inclusive_scan_sequence(Seq, Reduce, Number<I ...@@ -809,6 +809,46 @@ __host__ __device__ constexpr auto inclusive_scan_sequence(Seq, Reduce, Number<I
return reverse_inclusive_scan_sequence(Seq{}.Reverse(), Reduce{}, Number<Init>{}).Reverse(); return reverse_inclusive_scan_sequence(Seq{}.Reverse(), Reduce{}, Number<Init>{}).Reverse();
} }
// e.g. Seq<2, 3, 4> --> Seq<0, 2, 5>, Init=0, Reduce=Add
// ResultSeq TargetSeq Reduce
template <typename, typename, typename>
struct sequence_exclusive_scan;
template <index_t... Xs, index_t Y, index_t... Ys, typename Reduce>
struct sequence_exclusive_scan<Sequence<Xs...>, Sequence<Y, Ys...>, Reduce>
{
using old_scan = typename sequence_merge<Sequence<Xs...>,
Sequence<Reduce{}(Y, Sequence<Xs...>{}.Back())>>::type;
using type = typename sequence_exclusive_scan<old_scan, Sequence<Ys...>, Reduce>::type;
};
template <index_t... Xs, index_t Y, typename Reduce>
struct sequence_exclusive_scan<Sequence<Xs...>, Sequence<Y>, Reduce>
{
using type = Sequence<Xs...>;
};
template <index_t... Xs, typename Reduce>
struct sequence_exclusive_scan<Sequence<Xs...>, Sequence<>, Reduce>
{
using type = Sequence<Xs...>;
};
template <typename Seq, typename Reduce, index_t Init>
constexpr auto exclusive_scan_sequence(Seq, Reduce, Number<Init>)
{
// TODO: c++20 and later can pass in Reduce with a lambda expression
return typename sequence_exclusive_scan<Sequence<Init>, Seq, Reduce>::type{};
}
template <typename Seq>
constexpr auto prefix_sum_sequence(Seq)
{
return typename sequence_exclusive_scan<Sequence<0>,
typename sequence_merge<Seq, Sequence<0>>::type,
math::plus<index_t>>::type{};
}
template <typename Seq, index_t... Is> template <typename Seq, index_t... Is>
__host__ __device__ constexpr auto pick_sequence_elements_by_ids(Seq, Sequence<Is...> /* ids */) __host__ __device__ constexpr auto pick_sequence_elements_by_ids(Seq, Sequence<Is...> /* ids */)
{ {
......
include/ck/utility/sequence_helper.hpp
View file @ 6491acda
...@@ -3,7 +3,10 @@ ...@@ -3,7 +3,10 @@
#pragma once #pragma once
#include "ck/utility/sequence.hpp"
#include "ck/utility/array.hpp"
#include "ck/utility/tuple.hpp" #include "ck/utility/tuple.hpp"
#include "ck/utility/macro_func_array_to_sequence.hpp"
namespace ck { namespace ck {
...@@ -14,7 +17,7 @@ __host__ __device__ constexpr auto make_sequence(Number<Is>...) ...@@ -14,7 +17,7 @@ __host__ __device__ constexpr auto make_sequence(Number<Is>...)
} }
// F() returns index_t // F() returns index_t
// F use default constructor // F use default constructor, so F cannot be lambda function
template <typename F, index_t N> template <typename F, index_t N>
__host__ __device__ constexpr auto generate_sequence(F, Number<N>) __host__ __device__ constexpr auto generate_sequence(F, Number<N>)
{ {
...@@ -22,6 +25,7 @@ __host__ __device__ constexpr auto generate_sequence(F, Number<N>) ...@@ -22,6 +25,7 @@ __host__ __device__ constexpr auto generate_sequence(F, Number<N>)
} }
// F() returns Number<> // F() returns Number<>
// F could be lambda function
template <typename F, index_t N> template <typename F, index_t N>
__host__ __device__ constexpr auto generate_sequence_v2(F&& f, Number<N>) __host__ __device__ constexpr auto generate_sequence_v2(F&& f, Number<N>)
{ {
...@@ -35,4 +39,55 @@ __host__ __device__ constexpr auto to_sequence(Tuple<Number<Is>...>) ...@@ -35,4 +39,55 @@ __host__ __device__ constexpr auto to_sequence(Tuple<Number<Is>...>)
return Sequence<Is...>{}; return Sequence<Is...>{};
} }
namespace detail {
template <index_t h_idx, typename SeqSortedSamples, typename SeqRange>
struct sorted_sequence_histogram;
template <index_t h_idx, index_t x, index_t... xs, index_t r, index_t... rs>
struct sorted_sequence_histogram<h_idx, Sequence<x, xs...>, Sequence<r, rs...>>
{
template <typename Histogram>
constexpr auto operator()(Histogram& h)
{
if constexpr(x < r)
{
h.template At<h_idx>() += 1;
sorted_sequence_histogram<h_idx, Sequence<xs...>, Sequence<r, rs...>>{}(h);
}
else
{
h.template At<h_idx + 1>() = 1;
sorted_sequence_histogram<h_idx + 1, Sequence<xs...>, Sequence<rs...>>{}(h);
}
}
};
template <index_t h_idx, index_t x, index_t r, index_t... rs>
struct sorted_sequence_histogram<h_idx, Sequence<x>, Sequence<r, rs...>>
{
template <typename Histogram>
constexpr auto operator()(Histogram& h)
{
if constexpr(x < r)
{
h.template At<h_idx>() += 1;
}
}
};
} // namespace detail
// SeqSortedSamples: <0, 2, 3, 5, 7>, SeqRange: <0, 3, 6, 9> -> SeqHistogram : <2, 2, 1>
template <typename SeqSortedSamples, index_t r, index_t... rs>
constexpr auto histogram_sorted_sequence(SeqSortedSamples, Sequence<r, rs...>)
{
constexpr auto bins = sizeof...(rs); // or categories
constexpr auto histogram = [&]() {
Array<index_t, bins> h{0}; // make sure this can clear all element to zero
detail::sorted_sequence_histogram<0, SeqSortedSamples, Sequence<rs...>>{}(h);
return h;
}();
return TO_SEQUENCE(histogram, bins);
}
} // namespace ck } // namespace ck
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