Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(cell): Support different WarpLayout in GEMMs. #55

Merged
merged 3 commits into from
Feb 18, 2025
Merged

fix(cell): Support different WarpLayout in GEMMs. #55

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8b89491
fix(bugs): Support different WarpLayout in GEMMs.
KuangjuX Feb 16, 2025
25e31ec
chore: remove swizzle asserts.
KuangjuX Feb 18, 2025
bbb78c8
chore: Add kSwizzled flag in test_gemm.
KuangjuX Feb 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion 3rd-party/cutlass
Open in desktop
Submodule cutlass updated 2232 files
20 changes: 10 additions & 10 deletions include/cell/copy/shared_to_register.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -255,9 +255,9 @@ struct RegToSharedStorerImpl<Reg_, Shared_, kRowExec_, kColExec_,

DEVICE void operator()(const Reg& src, DType* dst, int start_offset) {
#pragma unroll
for (int i = 0; i < kRowExec; ++i) {
for (int j = 0; j < kColExec; ++j) {
#pragma unroll
for (int j = 0; j < kColExec; ++j) {
for (int i = 0; i < kRowExec; ++i) {
int lane_row = this->lane_row_id();
int lane_col = this->lane_col_id();

Expand Down Expand Up @@ -372,30 +372,30 @@ struct RegToSharedStorerImpl<Reg_, Shared_, kRowExec_, kColExec_,

DEVICE void operator()(const Reg& src, DType* dst, int start_offset) {
#pragma unroll
for (int i = 0; i < kColExec; ++i) {
for (int j = 0; j < kColExec; ++j) {
#pragma unroll
for (int j = 0; j < kRowExec; ++j) {
for (int i = 0; i < kRowExec; ++i) {
int tile_offset =
start_offset + i * kColStride + j * kRowStride;
start_offset + j * kColStride + i * kRowStride;
int lane_row = this->lane_row_id();
int lane_col = this->lane_col_id();

int row = 0, col = 0;
#pragma unroll
for (int m = 0; m < StoreMat::kSegRows; ++m) {
row = StoreMat::kElemPerSeg *
(lane_row + i * StoreMat::kThreadRows);
(lane_row + m * StoreMat::kThreadRows);
#pragma unroll
for (int n = 0; n < StoreMat::kSegCols; ++n) {
col = lane_col + j * StoreMat::kThreadCols;
col = lane_col + n * StoreMat::kThreadCols;

int in_tile_offset = col * Shared::kColStride + row;
int offset = tile_offset + in_tile_offset;
int swizzled_offset = get_swizzle_offset(offset);

const PackedType* src_ptr =
reinterpret_cast<const PackedType*>(
src(j, i).data());
src(i, j).data());
PackedType* dst_ptr =
reinterpret_cast<PackedType*>(dst);
dst_ptr[swizzled_offset / StoreMat::kElemPerSeg] =
Expand Down Expand Up @@ -463,9 +463,9 @@ struct RegToSharedStorerImpl<Reg_, Shared_, kRowExec_, kColExec_,
auto swizzled_tile_id = get_swizzled_tile_id(offset);
auto in_swizzled_tile_id = get_in_swizzle_tile_id(offset);
int swizzled_tile_offset =
src_tile_(swizzled_tile_id.x, swizzled_tile_id.y);
dst_tile_(swizzled_tile_id.x, swizzled_tile_id.y);
int in_swizzled_tile_offset =
in_src_tile_(in_swizzled_tile_id.x, in_swizzled_tile_id.y);
in_dst_tile_(in_swizzled_tile_id.x, in_swizzled_tile_id.y);

int offset_ = swizzled_tile_offset + in_swizzled_tile_offset;

Expand Down
33 changes: 26 additions & 7 deletions include/cell/copy/warp.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,10 +283,19 @@ template <typename WarpLayout_, typename BaseShape_, typename Shared_,
struct SharedOffsetHelper<WarpLayout_, BaseShape_, Shared_, kMode_,
tl::Layout::kRowMajor, false> {
DEVICE int get_warp_offset() {
// TODO(KuangjuX): hotfix this.
return warp_row_id<WarpLayout>() * kRowStride * BaseShape::kRows *
Shared::kCols +
warp_col_id<WarpLayout>() * kColStride * BaseShape::kCols;
switch (kMode) {
case WarpReuse::kCont:
return warp_row_id<WarpLayout>() * kRowStride *
BaseShape::kRows * Shared::kCols +
warp_col_id<WarpLayout>() * kColStride *
BaseShape::kCols;
case WarpReuse::kRowReuseCont:
return warp_row_id<WarpLayout>() * kRowStride *
BaseShape::kRows * Shared::kCols;
default:
assert(false && "Not implemented yet.");
return -1;
}
}

private:
Expand All @@ -308,9 +317,19 @@ template <typename WarpLayout_, typename BaseShape_, typename Shared_,
struct SharedOffsetHelper<WarpLayout_, BaseShape_, Shared_, kMode_,
tl::Layout::kColMajor, false> {
DEVICE int get_warp_offset() {
return warp_row_id<WarpLayout>() * kRowStride * BaseShape::kRows +
warp_col_id<WarpLayout>() * kColStride * BaseShape::kCols *
Shared::kRows;
switch (kMode) {
case WarpReuse::kCont:
return warp_row_id<WarpLayout>() * kRowStride *
BaseShape::kRows +
warp_col_id<WarpLayout>() * kColStride *
BaseShape::kCols * Shared::kRows;
case WarpReuse::kColReuseCont:
return warp_col_id<WarpLayout>() * kColStride *
BaseShape::kCols * Shared::kRows;
default:
assert(false && "Not implemented yet.");
return -1;
}
}

private:
Expand Down
8 changes: 6 additions & 2 deletions include/types/swizzle.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,9 @@ struct SwizzledLayout<Layout_, kB, kM, kS, tl::Layout::kRowMajor> {
*/
HOST_DEVICE auto operator()(int x, int y) const {
int idx = (x << (Mbits + Sbits)) | y;
assert(idx < (1 << (Bbits + Mbits + Sbits)));

// KuangjuX: This assert may affect the performance.
// assert(idx < (1 << (Bbits + Mbits + Sbits)));

int swizzled_idx = swizzle_(idx);
int swizzled_x = swizzled_idx >> (Mbits + Sbits);
Expand Down Expand Up @@ -114,7 +116,9 @@ struct SwizzledLayout<Layout_, kB, kM, kS, tl::Layout::kColMajor> {
*/
HOST_DEVICE auto operator()(int x, int y) const {
int idx = (y << (Bbits + Mbits)) | x;
assert(idx < (1 << (Bbits + Mbits + Sbits)));

// KuanjuX: This assert may affect the performance.
// assert(idx < (1 << (Bbits + Mbits + Sbits)));

int swizzled_idx = swizzle_(idx);
int swizzled_y = swizzled_idx >> (Mbits + Sbits);
Expand Down
4 changes: 4 additions & 0 deletions tests/cpp/cell/test_g2s_load.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,10 @@ TEST(GlobalToSharedLoad, test_row_major_load) {
run_test_row_major<float, tl::RowMajor<4, 1>, 64, 32, true>();
run_test_row_major<float, tl::RowMajor<2, 2>, 32, 64, true>();
run_test_row_major<float, tl::RowMajor<2, 4>, 32, 128, true>();

// To check correctness for next tests.
run_test_row_major<__half, tl::RowMajor<2, 1>, 64, 64, false>();
run_test_row_major<__half, tl::RowMajor<2, 1>, 64, 64, true>();
}

TEST(GlobalToSharedLoad, test_col_major_load) {
Expand Down
40 changes: 26 additions & 14 deletions tests/cpp/cell/test_gemm.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,8 @@ void cublas_hgemm(int m, int n, int k, const __half* A, const __half* B,
// @param strided_k: chunk size to partition the k dimension of the shared
// memory tile.
template <typename Element, typename ElementAcc, const int kM, const int kN,
const int kK, typename WarpLayout_, const int kChunkK>
const int kK, typename WarpLayout_, const int kChunkK,
const bool kSwizzled>
struct TestTraits {
using BaseShape = traits::BaseTileShape<Element>;

Expand All @@ -118,7 +119,6 @@ struct TestTraits {

/// == 2. configure tile transfer between global and shared using CuTe ==
using GlobalA = GlobalTile<Element, tl::RowMajor<kM, kK>>;
static const bool kSwizzled = true;
using SharedA = SharedTile<Element, tl::RowMajor<kM, kK>, kSwizzled>;
using LoadSharedA = GlobalToSharedLoader<SharedA, WarpLayout>;

Expand Down Expand Up @@ -216,7 +216,7 @@ __global__ void test_gemm(const Element* ga, const Element* gb,
} // namespace

template <const int kM, const int kN, const int kK, typename WarpLayout,
const int kChunkK>
const int kChunkK, const bool kSwizzled>
void run_test() {
// unittest for register-level gemm by calling into wmma PTX
using Element = __half;
Expand Down Expand Up @@ -249,8 +249,8 @@ void run_test() {
thrust::device_vector<ElementAcc> d_c = h_c;

// define the configuration of the test
using config =
TestTraits<Element, ElementAcc, kM, kN, kK, WarpLayout, kChunkK>;
using config = TestTraits<Element, ElementAcc, kM, kN, kK, WarpLayout,
kChunkK, kSwizzled>;

LOG(INFO) << "[" << kM << ", " << kN << ", " << kK << "], warps: ["
<< config::kWarpPerRow << ", " << config::kWarpPerCol
Expand Down Expand Up @@ -318,17 +318,29 @@ TEST(TestGemm, test) {
// as this will cause a shared memory overflow.

// 1 warp
run_test<16, 16, 64, tl::RowMajor<1, 1>, 64>(); // minimal shape
run_test<32, 16, 64, tl::RowMajor<1, 1>, 64>();
run_test<16, 32, 64, tl::RowMajor<1, 1>, 64>();
run_test<32, 32, 64, tl::RowMajor<1, 1>, 64>();
run_test<64, 64, 64, tl::RowMajor<1, 1>, 64>();
run_test<128, 64, 64, tl::RowMajor<1, 1>, 64>();
run_test<16, 16, 64, tl::RowMajor<1, 1>, 64, true>(); // minimal shape
run_test<32, 16, 64, tl::RowMajor<1, 1>, 64, true>();
run_test<16, 32, 64, tl::RowMajor<1, 1>, 64, true>();
run_test<32, 32, 64, tl::RowMajor<1, 1>, 64, true>();
run_test<64, 64, 64, tl::RowMajor<1, 1>, 64, true>();
run_test<128, 64, 64, tl::RowMajor<1, 1>, 64, true>();

// 2 x 1 warps
// TODO(KuangjuX): fix different warp layout.
// run_test<128, 64, 128, tl::RowMajor<2, 1>, 64>();
// run_test<128, 128, 128, tl::RowMajor<2, 1>, 64>();
run_test<32, 64, 128, tl::RowMajor<2, 1>, 128, true>();
run_test<64, 64, 128, tl::RowMajor<2, 1>, 128, true>();
run_test<32, 128, 128, tl::RowMajor<2, 1>, 128, true>();

// 1 x 2 warps
run_test<32, 128, 128, tl::RowMajor<1, 2>, 128, true>();
// TODO(KuangjuX): CUDA free failed: cudaErrorMisalignedAddress: misaligned
// address.
// run_test<64, 128, 128, tl::RowMajor<1, 2>, 128, true>();

// 2 x 2 warps
run_test<64, 64, 128, tl::RowMajor<2, 2>, 128, true>();

// 4 x 1 warps
run_test<64, 16, 256, tl::RowMajor<4, 1>, 256, true>();
}

} // namespace tilefusion::testing
89 changes: 62 additions & 27 deletions tests/cpp/cell/test_single_wmma.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,15 @@ __device__ void naive_gemm(int kM, int kN, int kK, //
}

__device__ void check_results(const float* hc1, const float* hc2, int numel) {
for (int i = 0; i < numel; ++i) assert(abs(hc1[i] - hc2[i]) < 1e-3);
for (int i = 0; i < numel; ++i) {
if (fabs(hc1[i] - hc2[i]) > 1e-3) {
printf("error: %d, %.4f, %.4f\n", i, hc1[i], hc2[i]);
printf("test failed!\n");
return;
}
}

printf("test passed!\n");

#if defined(DEBUG)
if (thread0()) {
Expand All @@ -59,10 +67,10 @@ __device__ void init_values(Element* a, Element* b, ElementAcc* c, int M, int N,
if (!thread0()) return;

for (int i = 0; i < M * K; ++i)
a[i] = static_cast<Element>(i % 2048 / 1000.);
a[i] = static_cast<Element>(i % 2048 / 1000);

for (int i = 0; i < K * N; ++i)
b[i] = static_cast<Element>(i % 2048 / 1000.);
b[i] = static_cast<Element>(i % 2048 / 1000);

for (int i = 0; i < M * N; ++i) c[i] = 0.;
}
Expand All @@ -84,6 +92,8 @@ __global__ void test_wmma(LoadRegA& load_rA, LoadRegB& load_rB,

init_values<Element, ElementAcc>(shared_a, shared_b, shared_c, M, N, K);

__syncthreads();

SharedC sC(shared_c);

TileIteratorA sAs(shared_a);
Expand All @@ -99,10 +109,13 @@ __global__ void test_wmma(LoadRegA& load_rA, LoadRegB& load_rB,

load_rA(sA, rA);
load_rB(sB, rB);
__syncthreads();

compute::gemm(rA, rB, acc);
}

__syncthreads();

store_rC(acc, sC);
__syncthreads();

Expand All @@ -119,29 +132,36 @@ __global__ void test_wmma(LoadRegA& load_rA, LoadRegB& load_rB,
} // namespace

template <typename Element, typename ElementAcc, const int kM, const int kN,
const int kK>
const int kK, typename WarpLayout>
struct TestTraits {
using WarpLayout = tl::RowMajor<1, 1>;
static const int kThreads = tl::get_numel<WarpLayout> * 32;

static constexpr int kWarpPerRow = tl::num_rows<WarpLayout>;
static constexpr int kWarpPerCol = tl::num_cols<WarpLayout>;

// ============= 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 constexpr int kAMs = kM / kWarpPerRow / BaseShape::kTileSize;
static constexpr int kAKs = kK / BaseShape::kTileSize;

static constexpr int kBKs = kK / BaseShape::kTileSize;
static constexpr int kBNs = kN / kWarpPerCol / BaseShape::kTileSize;

static constexpr int kCMs = kM / kWarpPerRow / BaseShape::kTileSize;
static constexpr int kCNs = kN / kWarpPerCol / BaseShape::kTileSize;

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 RegA = RegTile<BaseTileRowMajor<Element>, tl::RowMajor<kAMs, kAKs>>;
using LoadRegA =
SharedToRegLoader<RegA, WarpLayout, WarpReuse::kRowReuseCont>;

using SharedB = SharedTile<Element, tl::ColMajor<kK, kN>>;

using RegB = RegTile<BaseTileColMajor<Element>, tl::ColMajor<kKs, kNs>>;
using RegB = RegTile<BaseTileColMajor<Element>, tl::ColMajor<kBKs, kBNs>>;
using TileIteratorB = STileIterator<SharedB, TileShape<kK, kN>>;
using LoadRegB =
SharedToRegLoader<RegB, WarpLayout, WarpReuse::kColReuseCont>;
Expand All @@ -151,23 +171,21 @@ struct TestTraits {

// ============= register to shared storer =================
using SharedC = SharedTile<ElementAcc, tl::RowMajor<kM, kN>>;
using RegC = RegTile<BaseTileRowMajor<ElementAcc>, tl::RowMajor<kMs, kNs>>;
using RegC =
RegTile<BaseTileRowMajor<ElementAcc>, tl::RowMajor<kCMs, kCNs>>;
using StoreRegC = RegToSharedStorer<RegC, WarpLayout>;
};

template <const int kM, const int kN, const int kK>
template <const int kM, const int kN, const int kK, typename WarpLayout>
void run_test() {
using Element = cutlass::half_t;
using ElementAcc = float;

using config = TestTraits<Element, ElementAcc, kM, kN, kK>;
using config = TestTraits<Element, ElementAcc, kM, kN, kK, WarpLayout>;

dim3 dim_grid(1, 1, 1);
dim3 dim_block(config::kThreads, 1, 1);

int shm_size =
(kM + kN) * kK * sizeof(Element) + kM * kN * sizeof(ElementAcc);

typename config::LoadRegA load_rA;
typename config::LoadRegB load_rB;
typename config::StoreRegC store_rC;
Expand All @@ -181,23 +199,40 @@ void run_test() {
<< "RegB: " << RegB{} << std::endl
<< "RegC: " << RegC{} << std::endl;

test_wmma<Element, ElementAcc, kM, kN, kK, typename config::TileIteratorA,
RegA, typename config::LoadRegA, typename config::TileIteratorB,
RegB, typename config::LoadRegB, typename config::SharedC, RegC,
typename config::StoreRegC>
<<<dim_grid, dim_block, shm_size>>>(load_rA, load_rB, store_rC);
auto kernel =
test_wmma<Element, ElementAcc, kM, kN, kK,
typename config::TileIteratorA, RegA,
typename config::LoadRegA, typename config::TileIteratorB,
RegB, typename config::LoadRegB, typename config::SharedC,
RegC, typename config::StoreRegC>;

int shm_size =
(kM + kN) * kK * sizeof(Element) + 2 * kM * kN * sizeof(ElementAcc);

if (shm_size > 48 * 1024) {
cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size);
}

kernel<<<dim_grid, dim_block, shm_size>>>(load_rA, load_rB, store_rC);

cudaDeviceSynchronize();

LOG(INFO) << "[" << kM << ", " << kN << ", " << kK << "]. Test passed!"
<< std::endl;
}

TEST(TestWmma, test_m16n16k16_f) {
run_test<16, 16, 64>();
run_test<32, 32, 64>();
run_test<64, 64, 64>();
run_test<64, 128, 128>();
run_test<128, 128, 128>();
run_test<16, 16, 64, tl::RowMajor<1, 1>>();
run_test<32, 32, 64, tl::RowMajor<1, 1>>();
run_test<64, 64, 64, tl::RowMajor<1, 1>>();
// TODO(KuangjuX): It doesn't seem to be executed.
// run_test<64, 128, 128, tl::RowMajor<1, 1>>();

run_test<32, 64, 64, tl::RowMajor<2, 1>>();
run_test<32, 64, 64, tl::RowMajor<2, 1>>();
run_test<64, 64, 64, tl::RowMajor<2, 1>>();
run_test<128, 64, 64, tl::RowMajor<2, 1>>();
}

} // namespace tilefusion::testing
Loading