ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-03 13:11:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

[CK-tile] add more tests for batched transpose testing the rectangular block tile sizes (#2634)

Browse Source 
* add failing tests

* swap out and reference

* add constraint assert to transpose input distribution

* test both pipelines with rectangular block tile

* print mismatched indices

* add a smaller failing test for old pipeline

* print grid and block

* fill output before operating on it

* swap m/n tile sizes and make one test pass

* add device syncs

* add one more flipped test case

* flip block tile at host arg init

* fix tiles for lds pipeline

* clang-format

* rename tests

* roll back error check

* remove device syncs

* reduce large test case's size
This commit is contained in:
Max Podkorytov
2025-08-07 16:51:53 -07:00
committed by GitHub
parent 3c9400471d
commit ab26026835
5 changed files with 127 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
include/ck_tile/ops/batched_transpose/pipeline/batched_transpose_common_policy.hpp
View File

@@ -15,15 +15,15 @@ struct BatchedTransposeCommonPolicy
template <typename Problem>
CK_TILE_DEVICE static constexpr auto MakeInputDistribution()
{
constexpr index_t BlockSize = Problem::kBlockSize;
constexpr index_t LeadDimPerBlock = Problem::kMPerBlock;
constexpr index_t SecondDimPerBlock = Problem::kNPerBlock;
constexpr index_t kBlockSize = Problem::kBlockSize;
constexpr index_t kLeadDimPerBlock = Problem::kNPerBlock;
constexpr index_t kSecondDimPerBlock = Problem::kMPerBlock;
constexpr index_t kVectorSize = Problem::VectorSizeOutput;
using TileEncodingPattern = TileDistributionEncodingPattern2D<BlockSize,
SecondDimPerBlock,
LeadDimPerBlock,
constexpr index_t kVectorSize = Problem::VectorSizeInput;
static_assert((kLeadDimPerBlock * kVectorSize) % kBlockSize == 0, "");
using TileEncodingPattern = TileDistributionEncodingPattern2D<kBlockSize,
kSecondDimPerBlock,
kLeadDimPerBlock,
kVectorSize,
TileAccessPattern>;
return TileEncodingPattern::Make2DStaticTileDistribution();

20
include/ck_tile/ops/batched_transpose/pipeline/batched_transpose_lds_problem.hpp
View File

@@ -18,19 +18,19 @@ struct BatchedTransposeLdsProblem
{
using DataType = remove_cvref_t<DataType_>;
static constexpr index_t kRowWarps_ = NumWarps::at(number<1>{});
static constexpr index_t kColWarps_ = NumWarps::at(number<0>{});
static constexpr index_t kRowWarps_ = NumWarps::at(number<0>{});
static constexpr index_t kColWarps_ = NumWarps::at(number<1>{});
static constexpr index_t kBlockSize_ = get_warp_size() * kRowWarps_ * kColWarps_;
static constexpr index_t kRowPerBlock_ = BlockTile::at(number<1>{});
static constexpr index_t kColPerBlock_ = BlockTile::at(number<0>{});
static constexpr index_t kRowPerBlock_ = BlockTile::at(number<0>{});
static constexpr index_t kColPerBlock_ = BlockTile::at(number<1>{});
static constexpr index_t kBlockSize = kBlockSize_;
// warps per block
static constexpr index_t kLeadNumWarps = kRowWarps_;
static constexpr index_t kSecondNumWarps = kColWarps_;
static constexpr index_t kLeadNumWarps = kColWarps_;
static constexpr index_t kSecondNumWarps = kRowWarps_;
static constexpr index_t kLeadSizePerBlock = kRowPerBlock_;
static constexpr index_t kSecondSizePerBlock = kColPerBlock_;
static constexpr index_t kLeadSizePerBlock = kColPerBlock_;
static constexpr index_t kSecondSizePerBlock = kRowPerBlock_;
static constexpr index_t kQuadrantLeadDim = LaneGroupTransposeTraits<DataType>::kleadDim;
static constexpr index_t kQuadrantSecondDim = LaneGroupTransposeTraits<DataType>::ksecondDim;
@@ -60,8 +60,8 @@ struct BatchedTransposeLdsProblem
static constexpr bool kPadM = kPadM_;
static constexpr bool kPadN = kPadN_;
static constexpr auto kMPerBlock = kLeadSizePerBlock;
static constexpr auto kNPerBlock = kSecondSizePerBlock;
static constexpr auto kMPerBlock = kSecondSizePerBlock;
static constexpr auto kNPerBlock = kLeadSizePerBlock;
// 128-bit is the max single-instruction bandwidth for load/store
static constexpr index_t MaxLoadStoreSize = 16;

2
include/ck_tile/ops/batched_transpose/pipeline/batched_transpose_policy.hpp
View File

@@ -19,8 +19,8 @@ struct BatchedTransposePolicy : public BatchedTransposeCommonPolicy
constexpr index_t VecLoadSize = Problem::VectorSizeOutput;
using TileEncodingPattern = TileDistributionEncodingPattern2D<BlockSize,
NPerBlock,
MPerBlock,
NPerBlock,
VecLoadSize,
TileAccessPattern>;
return TileEncodingPattern::MakeShuffled2DStaticTileDistribution();