Skip to content

GEMM using split-k #13

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

Open
Yzhonghua wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

GEMM using split-k #13

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
109 changes: 109 additions & 0 deletions example/35_splitK_gemm/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
### Project Overview

This project focuses on optimizing matrix multiplication operations on GPUs. The initial code uses a standard GEMM operation without any advanced memory optimizations or block matrix multiplication techniques. Our goal is to enhance the performance by introducing pinned memory and block matrix multiplication techniques.

### Original Code

The original code performs GEMM (General Matrix Multiplication) using standard memory allocations and operations. It does not leverage pinned memory or block matrix multiplication techniques. The core operation is performed using the following class and methods:

```
using DeviceGemmInstance = ck::tensor_operation::device::DeviceGemmXdlSplitKCShuffle
<ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout,
AElementOp, BElementOp, CElementOp, GemmDefault, 256, 256, 128, 4, 8, 32, 32, 4, 2,
S<1, 4, 64, 1>, S<0, 2, 1, 3>, S<0, 2, 1, 3>, 3, 8, 8, true,
S<1, 4, 64, 1>, S<0, 1, 3, 2>, S<0, 1, 3, 2>, 3, 8, 8, true, 1, 1, S<1, 32, 1, 8>, 8>;

```

### Changes and Enhancements

* Pinned Memory Optimization

**Purpose:** To reduce memory transfer time between the host and the GPU by using pinned (page-locked) memory.

```
// Use pinned memory
void* pinned_memory_a;
void* pinned_memory_b;
void* pinned_memory_c;

hipError_t err;

err = hipHostMalloc(&pinned_memory_a, sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize(), hipHostMallocDefault);
if (err != hipSuccess) {
std::cerr << "hipHostMalloc failed for pinned_memory_a" << std::endl;
return false;
}

err = hipHostMalloc(&pinned_memory_b, sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize(), hipHostMallocDefault);
if (err != hipSuccess) {
std::cerr << "hipHostMalloc failed for pinned_memory_b" << std::endl;
hipHostFree(pinned_memory_a);
return false;
}

err = hipHostMalloc(&pinned_memory_c, sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize(), hipHostMallocDefault);
if (err != hipSuccess) {
std::cerr << "hipHostMalloc failed for pinned_memory_c" << std::endl;
hipHostFree(pinned_memory_a);
hipHostFree(pinned_memory_b);
return false;
}

```

* Block Matrix Multiplication

**Purpose:** To optimize the matrix multiplication by dividing the matrices into smaller blocks, which can be processed more efficiently on the GPU.

```
// Define block size
int block_size = 1024; // Adjust this size for optimization

// Iterate over blocks
for (int i = 0; i < M; i += block_size) {
for (int j = 0; j < N; j += block_size) {
for (int k = 0; k < K; k += block_size) {
int M_sub = std::min(block_size, M - i);
int N_sub = std::min(block_size, N - j);
int K_sub = std::min(block_size, K - k);

auto a_sub = a_m_k.Extract({i, k}, {M_sub, K_sub});
auto b_sub = b_k_n.Extract({k, j}, {K_sub, N_sub});
auto c_sub = c_m_n_device_result.Extract({i, j}, {M_sub, N_sub});

DeviceMem a_sub_device_buf(sizeof(ADataType) * a_sub.mDesc.GetElementSpaceSize());
DeviceMem b_sub_device_buf(sizeof(BDataType) * b_sub.mDesc.GetElementSpaceSize());
DeviceMem c_sub_device_buf(sizeof(CDataType) * c_sub.mDesc.GetElementSpaceSize());

a_sub_device_buf.ToDevice(a_sub.mData.data());
b_sub_device_buf.ToDevice(b_sub.mData.data());
c_sub_device_buf.SetZero();

auto argument = gemm.MakeArgument(
static_cast<ADataType*>(a_sub_device_buf.GetDeviceBuffer()),
static_cast<BDataType*>(b_sub_device_buf.GetDeviceBuffer()),
static_cast<CDataType*>(c_sub_device_buf.GetDeviceBuffer()),
M_sub, N_sub, K_sub, StrideA, StrideB, StrideC,
a_element_op, b_element_op, c_element_op, KBatch);

if (!gemm.IsSupportedArgument(argument)) {
std::cout << gemm.GetTypeString() << " does not support this problem" << std::endl;
hipHostFree(pinned_memory_a);
hipHostFree(pinned_memory_b);
hipHostFree(pinned_memory_c);
return false;
}

invoker.Run(argument, StreamConfig{nullptr, false});
c_sub_device_buf.FromDevice(c_sub.mData.data());
c_m_n_device_result.Update(c_sub, {i, j});
}
}
}

```

### Conclusion

Through the implementation of pinned memory and block matrix multiplication, significant performance improvements were achieved in the matrix multiplication operations on the GPU. The use of pinned memory reduces memory transfer time, while block matrix multiplication leverages the GPU's parallel processing capabilities more effectively. These optimizations demonstrate the potential for significant computational efficiency gains in high-performance computing applications.
227 changes: 217 additions & 10 deletions example/35_splitK_gemm/splitK_gemm_xdl_fp16.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
#include <numeric>
#include <initializer_list>
#include <cstdlib>
#include <cmath> // for std::abs

#include "ck/ck.hpp"
#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
Expand All @@ -29,10 +30,10 @@ using Col = ck::tensor_layout::gemm::ColumnMajor;

using PassThrough = ck::tensor_operation::element_wise::PassThrough;

using ADataType = F16;
using BDataType = F16;
using ADataType = F16;
using BDataType = F16;
using AccDataType = F32;
using CDataType = F16;
using CDataType = F16;

using ALayout = Row;
using BLayout = Col;
Expand All @@ -46,13 +47,219 @@ static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecializa

using DeviceGemmInstance = ck::tensor_operation::device::DeviceGemmXdlSplitKCShuffle
// clang-format off
//######| AData| BData| CData| AccData| ALayout| BLayout| CLayout| A| B| C| GEMM| Block| MPer| NPer| KPer| K1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer|
//######| Type| Type| Type| Type| | | | Elementwise| Elementwise| Elementwise| Spacialization| Size| Block| Block| Block| | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MXdlPerWave_MWaveMPerXdl| ScalarPerVector|
//######| | | | | | | | Operation| Operation| Operation| | | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NXdlPerWave_NWaveNPerXdl| _NWaveNPerXdl|
//######| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | |
< ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout, AElementOp, BElementOp, CElementOp, GemmDefault, 256, 256, 128, 4, 8, 32, 32, 4, 2, S<1, 4, 64, 1>, S<0, 2, 1, 3>, S<0, 2, 1, 3>, 3, 8, 8, true, S<1, 4, 64, 1>, S<0, 1, 3, 2>, S<0, 1, 3, 2>, 3, 8, 8, true, 1, 1, S<1, 32, 1, 8>, 8>;
< ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout,
AElementOp, BElementOp, CElementOp, GemmDefault, 256, 256, 128, 4, 8, 32, 32, 4, 2,
S<1, 4, 64, 1>, S<0, 2, 1, 3>, S<0, 2, 1, 3>, 3, 8, 8, true,
S<1, 4, 64, 1>, S<0, 1, 3, 2>, S<0, 1, 3, 2>, 3, 8, 8, true, 1, 1, S<1, 32, 1, 8>, 8>;
// clang-format on

#include "run_splitK_gemm_example.inc"
struct ProblemSize final {
ck::index_t M = 3840;
ck::index_t N = 4096;
ck::index_t K = 4096;
ck::index_t stride_A = K;
ck::index_t stride_B = K;
ck::index_t stride_C = N;
ck::index_t k_batch = 4;
};

int main(int argc, char* argv[]) { return !run_splitK_gemm_example(argc, argv); }
struct ExecutionConfig final {
bool do_verification = true;
int init_method = 1;
bool time_kernel = false;
};

bool run_splitK_gemm(const ProblemSize& problem_size, const ExecutionConfig& config) {
using namespace ck::literals;

#if defined(BUILD_INT4_EXAMPLE) && defined(CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4)
static_assert(sizeof(ck::int4_t) == sizeof(int8_t));
static_assert(sizeof(ADataType) == sizeof(KernelADataType));
static_assert(sizeof(BDataType) == sizeof(KernelBDataType));
#endif

const auto& [M, N, K, StrideA, StrideB, StrideC, KBatch] = problem_size;

auto create_host_tensor_descriptor = [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
using namespace ck::literals;

if (std::is_same<decltype(layout), ck::tensor_layout::gemm::RowMajor>::value) {
return HostTensorDescriptor({row, col}, {stride, 1_uz});
} else {
return HostTensorDescriptor({row, col}, {1_uz, stride});
}
};

Tensor<ADataType> a_m_k(create_host_tensor_descriptor(M, K, StrideA, ALayout{}));
Tensor<BDataType> b_k_n(create_host_tensor_descriptor(K, N, StrideB, BLayout{}));
Tensor<CDataType> c_m_n_device_result(create_host_tensor_descriptor(M, N, StrideC, CLayout{}));

std::cout << "a_m_k: " << a_m_k.mDesc << std::endl;
std::cout << "b_k_n: " << b_k_n.mDesc << std::endl;
std::cout << "c_m_n: " << c_m_n_device_result.mDesc << std::endl;

switch (config.init_method) {
case 0:
break;
case 1:
a_m_k.GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
b_k_n.GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5});
break;
case 2:
a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});
b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
break;
default:
a_m_k.GenerateTensorValue(GeneratorTensor_Sequential<0>{});
b_k_n.GenerateTensorValue(GeneratorTensor_Sequential<1>{});
}

// Use pinned memory
void* pinned_memory_a;
void* pinned_memory_b;
void* pinned_memory_c;

hipError_t err;

err = hipHostMalloc(&pinned_memory_a, sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize(), hipHostMallocDefault);
if (err != hipSuccess) {
std::cerr << "hipHostMalloc failed for pinned_memory_a" << std::endl;
return false;
}

err = hipHostMalloc(&pinned_memory_b, sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize(), hipHostMallocDefault);
if (err != hipSuccess) {
std::cerr << "hipHostMalloc failed for pinned_memory_b" << std::endl;
err = hipHostFree(pinned_memory_a);
return false;
}

err = hipHostMalloc(&pinned_memory_c, sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize(), hipHostMallocDefault);
if (err != hipSuccess) {
std::cerr << "hipHostMalloc failed for pinned_memory_c" << std::endl;
err = hipHostFree(pinned_memory_a);
err = hipHostFree(pinned_memory_b);
return false;
}

std::memcpy(pinned_memory_a, a_m_k.mData.data(), sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize());
std::memcpy(pinned_memory_b, b_k_n.mData.data(), sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
std::memcpy(pinned_memory_c, c_m_n_device_result.mData.data(), sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize());

DeviceMem a_m_k_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpaceSize());
DeviceMem b_k_n_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpaceSize());
DeviceMem c_m_n_device_buf(sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpaceSize());

#ifdef BUILD_INT4_EXAMPLE
const Tensor<KernelADataType> a_m_k_converted(a_m_k);
const Tensor<KernelBDataType> b_k_n_converted(b_k_n);

a_m_k_device_buf.ToDevice(a_m_k_converted.mData.data());
b_k_n_device_buf.ToDevice(b_k_n_converted.mData.data());
#else
a_m_k_device_buf.ToDevice(pinned_memory_a);
b_k_n_device_buf.ToDevice(pinned_memory_b);
#endif
c_m_n_device_buf.SetZero();

auto a_element_op = AElementOp{};
auto b_element_op = BElementOp{};
auto c_element_op = CElementOp{};

auto gemm = DeviceGemmInstance{};
auto invoker = gemm.MakeInvoker();
auto argument = gemm.MakeArgument(
#ifdef BUILD_INT4_EXAMPLE
static_cast<KernelADataType*>(a_m_k_device_buf.GetDeviceBuffer()),
static_cast<KernelBDataType*>(b_k_n_device_buf.GetDeviceBuffer()),
#else
static_cast<ADataType*>(a_m_k_device_buf.GetDeviceBuffer()),
static_cast<BDataType*>(b_k_n_device_buf.GetDeviceBuffer()),
#endif
static_cast<CDataType*>(c_m_n_device_buf.GetDeviceBuffer()),
M, N, K, StrideA, StrideB, StrideC, a_element_op, b_element_op, c_element_op, KBatch);

if (!gemm.IsSupportedArgument(argument)) {
std::cout << gemm.GetTypeString() << " does not support this problem" << std::endl;
err = hipHostFree(pinned_memory_a);
err = hipHostFree(pinned_memory_b);
err = hipHostFree(pinned_memory_c);
return false;
}

invoker.Run(argument, StreamConfig{nullptr, false});
bool pass = true;

if (config.do_verification) {
c_m_n_device_buf.FromDevice(c_m_n_device_result.mData.data());
using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType, BDataType, CDataType, AccDataType, AElementOp, BElementOp, CElementOp>;

auto ref_gemm = ReferenceGemmInstance{};
auto ref_invoker = ref_gemm.MakeInvoker();

Tensor<CDataType> c_m_n_host_result(create_host_tensor_descriptor(M, N, StrideC, CLayout{}));

auto ref_argument = ref_gemm.MakeArgument(a_m_k, b_k_n, c_m_n_host_result, a_element_op, b_element_op, c_element_op);

ref_invoker.Run(ref_argument);

if (std::is_same<CDataType, ck::half_t>::value) {
pass &= ck::utils::check_err(c_m_n_device_result, c_m_n_host_result, "fp16 incorrect result", 3e-3, 1e-3);
} else {
pass &= ck::utils::check_err(c_m_n_device_result, c_m_n_host_result);
}
}

if (config.time_kernel) {
float ave_time = invoker.Run(argument, StreamConfig{nullptr, config.time_kernel, 1});

std::size_t flop = std::size_t(2) * M * N * K;
std::size_t num_btype = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N;

float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
float gb_per_sec = num_btype / 1.E6 / ave_time;
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " << gemm.GetTypeString() << std::endl;
}

err = hipHostFree(pinned_memory_a);
err = hipHostFree(pinned_memory_b);
err = hipHostFree(pinned_memory_c);

return pass;
}

bool run_splitK_gemm_example(int argc, char* argv[]) {
ProblemSize problem_size;
ExecutionConfig config;

if (argc == 1) {
// use default case
} else if (argc == 5) {
config.do_verification = std::stoi(argv[1]);
config.init_method = std::stoi(argv[2]);
config.time_kernel = std::stoi(argv[3]);
problem_size.k_batch = std::stoi(argv[4]);
} else if (argc == 11) {
config.do_verification = std::stoi(argv[1]);
config.init_method = std::stoi(argv[2]);
config.time_kernel = std::stoi(argv[3]);
problem_size.k_batch = std::stoi(argv[4]);

problem_size.M = std::stoi(argv[5]);
problem_size.N = std::stoi(argv[6]);
problem_size.K = std::stoi(argv[7]);

problem_size.stride_A = std::stoi(argv[8]);
problem_size.stride_B = std::stoi(argv[9]);
problem_size.stride_C = std::stoi(argv[10]);
} else {
std::cerr << "Invalid number of arguments\n";
return false;
}

return run_splitK_gemm(problem_size, config);
}

int main(int argc, char* argv[]) {
return !run_splitK_gemm_example(argc, argv);
}
14 changes: 4 additions & 10 deletions script/cmake-ck-dev.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,19 +5,13 @@ rm -rf CMakeFiles

MY_PROJECT_SOURCE=$1

if [ $# -ge 2 ] ; then
GPU_TARGETS=$2
else
GPU_TARGETS="gfx908;gfx90a;gfx940"
fi

cmake \
-D CMAKE_PREFIX_PATH=/opt/rocm \
-D CMAKE_CXX_COMPILER=/opt/rocm/bin/hipcc \
-D CMAKE_CXX_FLAGS="-Xclang -mllvm -Xclang -enable-post-misched=0 -std=c++17 -O3 -ftemplate-backtrace-limit=0 -fPIE -Wno-gnu-line-marker" \
-D CMAKE_CXX_COMPILER=/opt/rocm-5.7.1/bin/hipcc \
-D CMAKE_CXX_FLAGS="-std=c++17 -O3 -ftemplate-backtrace-limit=0 -fPIE -Wno-gnu-line-marker" \
-D CMAKE_BUILD_TYPE=Release \
-D BUILD_DEV=ON \
-D GPU_TARGETS=$GPU_TARGETS \
-D GPU_TARGETS="gfx908;gfx90a;gfx940" \
-D CMAKE_VERBOSE_MAKEFILE:BOOL=ON \
-D USE_BITINT_EXTENSION_INT4=OFF \
${MY_PROJECT_SOURCE}
${MY_PROJECT_SOURCE}