Skip to content

Commit

Permalink
Update the interface of shared tile iterator.
Browse files Browse the repository at this point in the history
  • Loading branch information
@lcy-seso
lcy-seso committed Feb 1, 2025
1 parent cb8b4e8 commit 372e1ca
Show file tree
Hide file tree
Showing 5 changed files with 140 additions and 56 deletions.
18 changes: 11 additions & 7 deletions include/cell/copy/global_to_shared.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -447,16 +447,18 @@ struct GlobalToSharedLoader {
"The shape of SharedTile must be divisible by the shape of "
"WarpLayout.");

using WarpShape = TileShape<Shared::kRows / WarpLayout::kRows,
Shared::kCols / WarpLayout::kCols>;
static const WarpReuse kMode = WarpReuse::kCont; // warp reuse mode
using WarpShape = TileShape<
warp::warp_tile_rows<Shared::kRows, WarpLayout::kRows, kMode>(),
warp::warp_tile_rows<Shared::kCols, WarpLayout::kCols, kMode>()>;

using BaseShape = warp::WarpBaseTileShape<DType, WarpShape, Shared::kType>;

static_assert(Shared::kRows % BaseShape ::kRows == 0,
"Shared::kRows must be divisible by BaseShape::kRows.");
static_assert(Shared::kCols % BaseShape::kCols == 0,
"Shared::kCols must be divisible by BaseShape::kCols.");

static const WarpReuse kMode = WarpReuse::kCont; // warp reuse mode
using ExecCounter = warp::ExecCounter<BaseShape, Shared, WarpLayout, kMode>;
using GlobalOffset = warp::GlobalOffsetHelper<WarpLayout, kMode>;
using SharedOffset =
Expand Down Expand Up @@ -501,17 +503,19 @@ struct SharedToGlobalStorer {
using DType = Shared::DType;
using WarpLayout = WarpLayout_;

using WarpShape = TileShape<Shared::kRows / WarpLayout::kRows,
Shared::kCols / WarpLayout::kCols>;
static const WarpReuse kMode = WarpReuse::kCont; // warp reuse mode

using WarpShape = TileShape<
warp::warp_tile_rows<Shared::kRows, WarpLayout::kRows, kMode>(),
warp::warp_tile_rows<Shared::kCols, WarpLayout::kCols, kMode>()>;

using BaseShape = warp::WarpBaseTileShape<DType, WarpShape, Shared::kType>;

static_assert(Shared::kRows % BaseShape::kRows == 0,
"Shared::kRows must be divisible by BaseShape::kRows.");
static_assert(Shared::kCols % BaseShape::kCols == 0,
"Shared::kCols must be divisible by BaseShape::kCols.");

static const WarpReuse kMode = WarpReuse::kCont; // warp reuse mode

using SharedOffset =
warp::SharedOffsetHelper<WarpLayout, BaseShape, Shared, kMode>;
using GlobalOffset = warp::GlobalOffsetHelper<WarpLayout, kMode>;
Expand Down
24 changes: 24 additions & 0 deletions include/cell/copy/warp.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,30 @@ DEVICE int warp_col_id() {
}
}

template <const int kSharedRows, const int kWarpRows, const WarpReuse kMode>
HOST_DEVICE constexpr int warp_tile_rows() {
if constexpr (kMode == WarpReuse::kCont) {
return kSharedRows / kWarpRows;
} else if constexpr (kMode == WarpReuse::kRowReuseCont) {
return kSharedRows / kWarpRows;
} else if constexpr (kMode == WarpReuse::kColReuseCont) {
return kSharedRows;
}
return -1;
}

template <const int kSharedCols, const int kWarpCols, const WarpReuse kMode>
HOST_DEVICE constexpr int warp_tile_cols() {
if constexpr (kMode == WarpReuse::kCont) {
return kSharedCols / kWarpCols;
} else if constexpr (kMode == WarpReuse::kRowReuseCont) {
return kSharedCols;
} else if constexpr (kMode == WarpReuse::kColReuseCont) {
return kSharedCols / kWarpCols;
}
return -1;
}

template <typename BaseTile_, typename Tile_, typename WarpLayout_,
const WarpReuse kMode_>
struct ExecCounter {
Expand Down
7 changes: 4 additions & 3 deletions include/types/shared_tile_iterator.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ struct STileIteratorPrettyPrinter {
/// @tparam Tile_: The type of the large tile to chunk.
/// @tparam ChunkShape_: The shape of the smaller tiles into which the large
/// tile is partitioned (chunk shape).
template <class Tile_, class ChunkShape_>
template <class Tile_, class ChunkShape_, class BaseShape_>
class STileIterator {
public:
using Tile = Tile_;
Expand Down Expand Up @@ -133,9 +133,10 @@ class STileIterator {

/// @brief Pretty printer for the static shape information of a TileIterator.
/// Note: This printer function works ONLY on the host.
template <typename TileShape, typename ChunkShape>
template <typename TileShape, typename ChunkShape, typename BaseShape>
static HOST std::ostream& operator<<(
std::ostream& out, const STileIterator<TileShape, ChunkShape>& itr) {
std::ostream& out,
const STileIterator<TileShape, ChunkShape, BaseShape>& itr) {
STileIteratorPrettyPrinter::print(out, itr);
return out;
}
Expand Down
54 changes: 34 additions & 20 deletions tests/cpp/cell/test_single_wmma.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -120,29 +120,43 @@ template <typename Element, typename ElementAcc, const int kM, const int kN,
const int kK>
struct TestTraits {
using WarpLayout = tl::RowMajor<1, 1>;
static const int kThreads = tl::get_numel<WarpLayout> * 32;

// ============= shared to register loader =================
// TODO: whether BaseTileShape should depend on Element type?
using BaseShape = traits::BaseTileShape<Element>;
// how many elements a BaseTile are executed along the m, n, k dimension
static constexpr int kMs = kM / BaseShape::kTileSize;
static constexpr int kNs = kN / BaseShape::kTileSize;
static constexpr int kKs = kK / BaseShape::kTileSize;
static const int kThreads = WarpLayout::kNumel * 32;

// ============= shared to register loader A =================
static constexpr WarpReuse kModeA = WarpReuse::kRowReuseCont;
static constexpr int kWarpRowsA =
warp::warp_tile_rows<kM, WarpLayout::kRows, kModeA>();
static constexpr int kWarpColsA =
warp::warp_tile_cols<kK, WarpLayout::kCols, kModeA>();
using BaseShapeA =
warp::WarpBaseTileShape<Element, TileShape<kWarpRowsA, kWarpColsA>,
tl::Layout::kRowMajor>;

using SharedA = SharedTile<Element, tl::RowMajor<kM, kK>>;
using TileIteratorA = STileIterator<SharedA, TileShape<kM, kK>>;

using RegA = RegTile<BaseTileRowMajor<Element>, tl::RowMajor<kMs, kKs>>;
using LoadRegA =
SharedToRegLoader<RegA, WarpLayout, WarpReuse::kRowReuseCont>;
using TileIteratorA = STileIterator<SharedA, TileShape<kM, kK>, BaseShapeA>;

static constexpr int kMs = kWarpRowsA / BaseShapeA::kRows;
static constexpr int kAKs = kWarpColsA / BaseShapeA::kCols;
using RegA = RegTile<BaseTileRowMajor<Element>, tl::RowMajor<kMs, kAKs>>;
using LoadRegA = SharedToRegLoader<RegA, WarpLayout, kModeA>;

// ============= shared to register loader B =================
static constexpr WarpReuse kModeB = WarpReuse::kColReuseCont;
static constexpr int kWarpRowsB =
warp::warp_tile_rows<kK, WarpLayout::kRows, kModeB>();
static constexpr int kWarpColsB =
warp::warp_tile_cols<kN, WarpLayout::kCols, kModeB>();
using BaseShapeB =
warp::WarpBaseTileShape<Element, TileShape<kWarpRowsB, kWarpColsB>,
tl::Layout::kColMajor>;

using SharedB = SharedTile<Element, tl::ColMajor<kK, kN>>;
using TileIteratorB = STileIterator<SharedB, TileShape<kK, kN>, BaseShapeB>;

using RegB = RegTile<BaseTileColMajor<Element>, tl::ColMajor<kKs, kNs>>;
using TileIteratorB = STileIterator<SharedB, TileShape<kK, kN>>;
using LoadRegB =
SharedToRegLoader<RegB, WarpLayout, WarpReuse::kColReuseCont>;
static constexpr int kBKs = kWarpRowsB / BaseShapeB::kRows;
static constexpr int kNs = kWarpColsB / BaseShapeB::kCols;
using RegB = RegTile<BaseTileColMajor<Element>, tl::ColMajor<kBKs, kNs>>;
using LoadRegB = SharedToRegLoader<RegB, WarpLayout, kModeB>;

static_assert(TileIteratorA::sc1 == TileIteratorB::sc0,
"dimension mismatch!");
Expand Down Expand Up @@ -193,8 +207,8 @@ void run_test() {
TEST(TestWmma, test_m16n16k16_f) {
run_test<16, 16, 16>(); // Test the `BaseTile` 16x16x16
run_test<16, 32, 16>();
run_test<96, 48, 80>();
run_test<64, 128, 128>();
// run_test<96, 48, 80>();
// run_test<64, 128, 128>();
}

} // namespace tilefusion::testing
93 changes: 67 additions & 26 deletions tests/cpp/cell/test_swizzled_copy.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,10 @@
#include <sstream>

namespace tilefusion::testing {

using namespace cell;
using namespace copy;

namespace tl = tile_layout;

namespace {
Expand Down Expand Up @@ -91,36 +93,55 @@ __global__ void swizzled_copy(const Element* data, G2S1& g2s,

/// @brief This unit test verifies the correctness of the swizzled row-major
/// format for loading operand A in GEMM.
/// NOTE: This test is specifically for the "kRowReuseCont" warp reuse
/// mode.
template <typename WarpLayout, const int kRows, const int kCols,
const int kShmRows, const int kShmCols, const int kChunkShm>
void run_test_rowmajor() {
static_assert(kShmRows == kRows, "kShmRows must be equal to kRows");

using Element = __half;
const int kThreads = WarpLayout::kNumel * 32;
static constexpr int kWarpPerRow = WarpLayout::kRows;
static constexpr WarpReuse kMode = WarpReuse::kRowReuseCont;

using Global = GlobalTile<Element, tl::RowMajor<kRows, kCols>>;
using GIterator = GTileIterator<Global, TileShape<kRows, kShmCols>>;

// for non-swizzled layout
using Shared1 =
SharedTile<Element, tl::RowMajor<kShmRows, kShmCols>, false>;
using SIterator1 = STileIterator<Shared1, TileShape<kShmRows, kChunkShm>>;

// for swizzled layout
using Shared2 = SharedTile<Element, tl::RowMajor<kShmRows, kShmCols>, true>;
using SIterator2 = STileIterator<Shared2, TileShape<kShmRows, kChunkShm>>;

using BaseShape = traits::BaseTileShape<Element>;
using WarpShape = // automatically infer the BaseTile shape
TileShape<warp::warp_tile_rows<kShmRows, WarpLayout::kRows, kMode>(),
warp::warp_tile_cols<kShmCols, WarpLayout::kCols, kMode>()>;
using BaseShape =
warp::WarpBaseTileShape<Element, WarpShape, Global::kType>;

// TODO(ying): The user is currently responsible for ensuring the correct
// coordination between `BaseShape` and `TileIterator`. However, since
// `BaseShape` is intended to be an internal concept, update this to make it
// more transparent for the user.
static_assert(
kShmRows % BaseShape::kRows == 0 && kShmCols % BaseShape::kCols == 0,
"kRows and kCols must be multiples of BaseShape::kRows and "
"BaseShape::kCols, respectively.");
using GIterator = GTileIterator<Global, TileShape<kRows, kShmCols>>;

const int kSc0 = kShmRows / kWarpPerRow / BaseShape::kRows;
const int kSc1 = kChunkShm / BaseShape::kCols;
using SIterator1 =
STileIterator<Shared1, TileShape<kShmRows, kChunkShm>, BaseShape>;
using SIterator2 =
STileIterator<Shared2, TileShape<kShmRows, kChunkShm>, BaseShape>;

const int kSc0 = kShmRows / WarpLayout::kRows / BaseShape::kRows;
const int kSc1 = kChunkShm / BaseShape::kCols;
using Reg = RegTile<BaseTileRowMajor<Element>, tl::RowMajor<kSc0, kSc1>>;

#ifdef DEBUG
LOG(INFO) << "GIterator: " << GIterator{} << std::endl
LOG(INFO) << std::endl
<< "WarpShape: (" << dim_size<0, WarpShape> << ", "
<< dim_size<1, WarpShape> << ")" << std::endl
<< "BaseShape: " << BaseShape{} << std::endl
<< "GIterator: " << GIterator{} << std::endl
<< "SIterator1: " << SIterator1{} << std::endl
<< "SIterator2: " << SIterator2{} << std::endl
<< "GlobalTile: " << Global{} << std::endl
Expand All @@ -130,7 +151,7 @@ void run_test_rowmajor() {

using G2S1 = GlobalToSharedLoader<Shared1, WarpLayout>;
using G2S2 = GlobalToSharedLoader<Shared2, WarpLayout>;
using S2R = SharedToRegLoader<Reg, WarpLayout, WarpReuse::kRowReuseCont>;
using S2R = SharedToRegLoader<Reg, WarpLayout, kMode>;

dim3 dim_grid(1, 1, 1);
dim3 dim_block(kThreads, 1, 1);
Expand Down Expand Up @@ -174,34 +195,52 @@ void run_test_rowmajor() {
template <typename WarpLayout, const int kRows /*K*/, const int kCols /*N*/,
const int kShmRows, const int kShmCols, const int kChunkShm>
void run_test_colmajor() {
using Element = __half;
const int kThreads = tl::get_numel<WarpLayout> * 32;
static constexpr int kWarpPerCol = tl::num_cols<WarpLayout>;

static_assert(kShmCols == kCols, "kShmCols must be equal to kCols.");

using Global = GlobalTile<Element, tl::ColMajor<kRows, kCols>>;
using GIterator = GTileIterator<Global, TileShape<kShmRows, kShmCols>>;
using Element = __half;
const int kThreads = WarpLayout::kNumel * 32;
static constexpr WarpReuse kMode = WarpReuse::kColReuseCont;

using Global = GlobalTile<Element, tl::ColMajor<kRows, kCols>>;
// for non-swizzled layout
using Shared1 = SharedTile<Element, tl::ColMajor<kShmRows, kShmCols>,
false /*disable swizzled layout on shared*/>;
using SIterator1 = STileIterator<Shared1, TileShape<kChunkShm, kShmCols>>;

// for swizzled layout
using Shared2 = SharedTile<Element, tl::ColMajor<kShmRows, kShmCols>,
true /*enable swizzled layout on shared*/>;
using SIterator2 = STileIterator<Shared2, TileShape<kChunkShm, kShmCols>>;

using BaseShape = traits::BaseTileShape<Element>;
using WarpShape = // automatically infer the BaseTile shape
TileShape<warp::warp_tile_rows<kShmRows, WarpLayout::kRows, kMode>(),
warp::warp_tile_cols<kShmCols, WarpLayout::kCols, kMode>()>;
using BaseShape =
warp::WarpBaseTileShape<Element, WarpShape, Global::kType>;

// TODO(ying): The user is currently responsible for ensuring the correct
// coordination between `BaseShape` and `TileIterator`. However, since
// `BaseShape` is intended to be an internal concept, update this to make it
// more transparent for the user.
static_assert(
kShmRows % BaseShape::kRows == 0 && kShmCols % BaseShape::kCols == 0,
"kRows and kCols must be multiples of BaseShape::kRows and "
"BaseShape::kCols, respectively.");
using GIterator = GTileIterator<Global, TileShape<kShmRows, kShmCols>>;

using SIterator1 =
STileIterator<Shared1, TileShape<kChunkShm, kShmCols>, BaseShape>;
using SIterator2 =
STileIterator<Shared2, TileShape<kChunkShm, kShmCols>, BaseShape>;

const int kSc0 = kChunkShm / BaseShape::kRows;
const int kSc1 = kShmCols / BaseShape::kCols / kWarpPerCol;
const int kSc1 = kShmCols / BaseShape::kCols / WarpLayout::kCols;

using Reg = RegTile<BaseTileColMajor<Element>, tl::ColMajor<kSc0, kSc1>>;

#ifdef DEBUG
LOG(INFO) << "GIterator: " << GIterator{} << std::endl
LOG(INFO) << std::endl
<< "WarpShape: (" << dim_size<0, WarpShape> << ", "
<< dim_size<1, WarpShape> << ")" << std::endl
<< "BaseShape: " << BaseShape{} << std::endl
<< "GIterator: " << GIterator{} << std::endl
<< "SIterator1: " << SIterator1{} << std::endl
<< "SIterator2: " << SIterator2{} << std::endl
<< "GlobalTile: " << Global{} << std::endl
Expand All @@ -211,7 +250,7 @@ void run_test_colmajor() {

using G2S1 = GlobalToSharedLoader<Shared1, WarpLayout>;
using G2S2 = GlobalToSharedLoader<Shared2, WarpLayout>;
using S2R = SharedToRegLoader<Reg, WarpLayout, WarpReuse::kColReuseCont>;
using S2R = SharedToRegLoader<Reg, WarpLayout, kMode>;

dim3 dim_grid(1, 1, 1);
dim3 dim_block(kThreads, 1, 1);
Expand Down Expand Up @@ -392,7 +431,8 @@ void test_col_major_store() {
} // namespace

TEST(TestSwizzledLoad, test_load_row_major) {
// run_test_rowmajor<tl::RowMajor<1, 1>, 16, 16, 16, 16, 16>();
run_test_rowmajor<tl::RowMajor<1, 1>, 16, 16, 16, 16, 16>();

// run_test_rowmajor<tl::RowMajor<1, 1>, 16, 32, 16, 32, 16>();
// run_test_rowmajor<tl::RowMajor<1, 1>, 16, 32, 16, 32, 32>();
// run_test_rowmajor<tl::RowMajor<1, 1>, 32, 32, 32, 32, 16>();
Expand Down Expand Up @@ -425,7 +465,8 @@ TEST(TestSwizzledLoad, test_load_row_major) {
}

TEST(TestSwizzledLoad, test_load_col_major) {
// run_test_colmajor<tl::RowMajor<1, 1>, 16, 16, 16, 16, 16>();
run_test_colmajor<tl::RowMajor<1, 1>, 16, 16, 16, 16, 16>();

// run_test_colmajor<tl::RowMajor<1, 1>, 32, 16, 32, 16, 16>();
// run_test_colmajor<tl::RowMajor<1, 1>, 32, 16, 32, 16, 32>();
// run_test_colmajor<tl::RowMajor<1, 1>, 32, 32, 16, 32, 16>();
Expand Down

0 comments on commit 372e1ca

Please sign in to comment.