Implemented direct version of strided-batched GEMM kernel

pull/239/head
Cedric Nugteren 2018-01-07 14:58:45 +01:00
parent 9fb2c61b25
commit 13f0f6fc6e
2 changed files with 121 additions and 37 deletions

View File

@ -17,15 +17,16 @@
R"(
// =================================================================================================
#if defined(ROUTINE_GEMMBATCHED)
// Direct version of the batched GEMM kernel with [A, B] = [non-transposed, non-transposed]
__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
void XgemmDirectBatchedNN(const int kSizeM, const int kSizeN, const int kSizeK,
const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
__global real* cgm, const __constant int* c_offsets, const int c_ld,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
__global real* cgm, const __constant int* c_offsets, const int c_ld,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const int batch = get_group_id(2);
const real_arg arg_alpha = arg_alphas[batch];
const real_arg arg_beta = arg_betas[batch];
@ -42,11 +43,11 @@ void XgemmDirectBatchedNN(const int kSizeM, const int kSizeN, const int kSizeK,
// Direct version of the batched GEMM kernel with [A, B] = [non-transposed, transposed]
__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
void XgemmDirectBatchedNT(const int kSizeM, const int kSizeN, const int kSizeK,
const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
__global real* cgm, const __constant int* c_offsets, const int c_ld,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
__global real* cgm, const __constant int* c_offsets, const int c_ld,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const int batch = get_group_id(2);
const real_arg arg_alpha = arg_alphas[batch];
const real_arg arg_beta = arg_betas[batch];
@ -63,11 +64,11 @@ void XgemmDirectBatchedNT(const int kSizeM, const int kSizeN, const int kSizeK,
// Direct version of the batched GEMM kernel with [A, B] = [transposed, non-transposed]
__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
void XgemmDirectBatchedTN(const int kSizeM, const int kSizeN, const int kSizeK,
const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
__global real* cgm, const __constant int* c_offsets, const int c_ld,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
__global real* cgm, const __constant int* c_offsets, const int c_ld,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const int batch = get_group_id(2);
const real_arg arg_alpha = arg_alphas[batch];
const real_arg arg_beta = arg_betas[batch];
@ -84,11 +85,11 @@ void XgemmDirectBatchedTN(const int kSizeM, const int kSizeN, const int kSizeK,
// Direct version of the batched GEMM kernel with [A, B] = [transposed, transposed]
__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
void XgemmDirectBatchedTT(const int kSizeM, const int kSizeN, const int kSizeK,
const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
__global real* cgm, const __constant int* c_offsets, const int c_ld,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const __constant real_arg* arg_alphas, const __constant real_arg* arg_betas,
const __global realMD* restrict agm, const __constant int* a_offsets, const int a_ld,
const __global realND* restrict bgm, const __constant int* b_offsets, const int b_ld,
__global real* cgm, const __constant int* c_offsets, const int c_ld,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const int batch = get_group_id(2);
const real_arg arg_alpha = arg_alphas[batch];
const real_arg arg_beta = arg_betas[batch];
@ -102,6 +103,87 @@ void XgemmDirectBatchedTT(const int kSizeM, const int kSizeN, const int kSizeK,
alm, blm, 1, 1, c_transpose, a_conjugate, b_conjugate);
}
#endif
// =================================================================================================
#if defined(ROUTINE_GEMMSTRIDEDBATCHED)
// Direct version of the strided-batched GEMM kernel with [A, B] = [non-transposed, non-transposed]
__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
void XgemmDirectStridedBatchedNN(const int kSizeM, const int kSizeN, const int kSizeK,
const real_arg arg_alpha, const real_arg arg_beta,
const __global realMD* restrict agm, const int a_offset, const int a_ld, const int a_stride,
const __global realND* restrict bgm, const int b_offset, const int b_ld, const int b_stride,
__global real* cgm, const int c_offset, const int c_ld, const int c_stride,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const int batch = get_group_id(2);
const int a_offset_batch = a_offset + a_stride * batch;
const int b_offset_batch = b_offset + b_stride * batch;
const int c_offset_batch = c_offset + c_stride * batch;
__local real alm[WGD * (WGD + PADA)];
__local real blm[WGD * (WGD + PADB)];
XgemmDirect(kSizeM, kSizeN, kSizeK, arg_alpha, arg_beta,
agm, a_offset_batch, a_ld, bgm, b_offset_batch, b_ld, cgm, c_offset_batch, c_ld,
alm, blm, 0, 0, c_transpose, a_conjugate, b_conjugate);
}
// Direct version of the strided-batched GEMM kernel with [A, B] = [non-transposed, transposed]
__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
void XgemmDirectStridedBatchedNT(const int kSizeM, const int kSizeN, const int kSizeK,
const real_arg arg_alpha, const real_arg arg_beta,
const __global realMD* restrict agm, const int a_offset, const int a_ld, const int a_stride,
const __global realND* restrict bgm, const int b_offset, const int b_ld, const int b_stride,
__global real* cgm, const int c_offset, const int c_ld, const int c_stride,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const int batch = get_group_id(2);
const int a_offset_batch = a_offset + a_stride * batch;
const int b_offset_batch = b_offset + b_stride * batch;
const int c_offset_batch = c_offset + c_stride * batch;
__local real alm[WGD * (WGD + PADA)];
__local real blm[WGD * (WGD + PADB)];
XgemmDirect(kSizeM, kSizeN, kSizeK, arg_alpha, arg_beta,
agm, a_offset_batch, a_ld, bgm, b_offset_batch, b_ld, cgm, c_offset_batch, c_ld,
alm, blm, 0, 1, c_transpose, a_conjugate, b_conjugate);
}
// Direct version of the strided-batched GEMM kernel with [A, B] = [transposed, non-transposed]
__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
void XgemmDirectStridedBatchedTN(const int kSizeM, const int kSizeN, const int kSizeK,
const real_arg arg_alpha, const real_arg arg_beta,
const __global realMD* restrict agm, const int a_offset, const int a_ld, const int a_stride,
const __global realND* restrict bgm, const int b_offset, const int b_ld, const int b_stride,
__global real* cgm, const int c_offset, const int c_ld, const int c_stride,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const int batch = get_group_id(2);
const int a_offset_batch = a_offset + a_stride * batch;
const int b_offset_batch = b_offset + b_stride * batch;
const int c_offset_batch = c_offset + c_stride * batch;
__local real alm[WGD * (WGD + PADA)];
__local real blm[WGD * (WGD + PADB)];
XgemmDirect(kSizeM, kSizeN, kSizeK, arg_alpha, arg_beta,
agm, a_offset_batch, a_ld, bgm, b_offset_batch, b_ld, cgm, c_offset_batch, c_ld,
alm, blm, 1, 0, c_transpose, a_conjugate, b_conjugate);
}
// Direct version of the strided-batched GEMM kernel with [A, B] = [transposed, transposed]
__kernel __attribute__((reqd_work_group_size(MDIMCD, NDIMCD, 1)))
void XgemmDirectStridedBatchedTT(const int kSizeM, const int kSizeN, const int kSizeK,
const real_arg arg_alpha, const real_arg arg_beta,
const __global realMD* restrict agm, const int a_offset, const int a_ld, const int a_stride,
const __global realND* restrict bgm, const int b_offset, const int b_ld, const int b_stride,
__global real* cgm, const int c_offset, const int c_ld, const int c_stride,
const int c_transpose, const int a_conjugate, const int b_conjugate) {
const int batch = get_group_id(2);
const int a_offset_batch = a_offset + a_stride * batch;
const int b_offset_batch = b_offset + b_stride * batch;
const int c_offset_batch = c_offset + c_stride * batch;
__local real alm[WGD * (WGD + PADA)];
__local real blm[WGD * (WGD + PADB)];
XgemmDirect(kSizeM, kSizeN, kSizeK, arg_alpha, arg_beta,
agm, a_offset_batch, a_ld, bgm, b_offset_batch, b_ld, cgm, c_offset_batch, c_ld,
alm, blm, 1, 1, c_transpose, a_conjugate, b_conjugate);
}
#endif
// =================================================================================================
// End of the C++11 raw string literal

View File

@ -112,6 +112,7 @@ void XgemmStridedBatched<T>::BatchedGemmIndirect(const size_t m, const size_t n,
const size_t b_one, const size_t b_two,
const size_t c_one, const size_t c_two,
const size_t batch_count) {
/* TODO
// Calculates the ceiled versions of m, n, and k
const auto m_ceiled = Ceil(Ceil(m, db_["MWG"]), db_["VWM"]);
const auto n_ceiled = Ceil(Ceil(n, db_["NWG"]), db_["VWN"]);
@ -123,7 +124,6 @@ void XgemmStridedBatched<T>::BatchedGemmIndirect(const size_t m, const size_t n,
Xgemm<T>::CalculateInternalDimensions(m, n, k, db_["MWG"], db_["NWG"], db_["KWG"],
a_one_i, a_two_i, b_one_i, b_two_i, c_one_i, c_two_i);
/* TODO
// Sets the "internal" offsets, i.e. the perfect offsets
auto a_offsets_i = 0;//std::vector<int>(batch_count);
auto b_offsets_i = 0;//std::vector<int>(batch_count);
@ -244,30 +244,33 @@ void XgemmStridedBatched<T>::BatchedGemmDirect(const size_t m, const size_t n, c
const bool a_do_transpose, const bool b_do_transpose, const bool c_do_transpose,
const bool a_conjugate, const bool b_conjugate,
const size_t batch_count) {
/* TODO
// Retrieves the proper XgemmDirect kernel from the compiled binary
const auto name = (a_do_transpose) ? (b_do_transpose ? "XgemmDirectBatchedTT" : "XgemmDirectBatchedTN") :
(b_do_transpose ? "XgemmDirectBatchedNT" : "XgemmDirectBatchedNN");
const auto name = (a_do_transpose) ? (b_do_transpose ? "XgemmDirectStridedBatchedTT" : "XgemmDirectStridedBatchedTN") :
(b_do_transpose ? "XgemmDirectStridedBatchedNT" : "XgemmDirectStridedBatchedNN");
auto kernel = Kernel(program_, name);
// Sets the kernel arguments
kernel.SetArgument(0, static_cast<int>(m));
kernel.SetArgument(1, static_cast<int>(n));
kernel.SetArgument(2, static_cast<int>(k));
kernel.SetArgument(3, alpha);
kernel.SetArgument(4, beta);
kernel.SetArgument(3, GetRealArg(alpha));
kernel.SetArgument(4, GetRealArg(beta));
kernel.SetArgument(5, a_buffer());
kernel.SetArgument(6, a_offset);
kernel.SetArgument(6, static_cast<int>(a_offset));
kernel.SetArgument(7, static_cast<int>(a_ld));
kernel.SetArgument(8, b_buffer());
kernel.SetArgument(9, b_offset);
kernel.SetArgument(10, static_cast<int>(b_ld));
kernel.SetArgument(11, c_buffer());
kernel.SetArgument(12, c_offset);
kernel.SetArgument(13, static_cast<int>(c_ld));
kernel.SetArgument(14, static_cast<int>(c_do_transpose));
kernel.SetArgument(15, static_cast<int>(a_conjugate));
kernel.SetArgument(16, static_cast<int>(b_conjugate));
kernel.SetArgument(8, static_cast<int>(a_stride));
kernel.SetArgument(9, b_buffer());
kernel.SetArgument(10, static_cast<int>(b_offset));
kernel.SetArgument(11, static_cast<int>(b_ld));
kernel.SetArgument(12, static_cast<int>(b_stride));
kernel.SetArgument(13, c_buffer());
kernel.SetArgument(14, static_cast<int>(c_offset));
kernel.SetArgument(15, static_cast<int>(c_ld));
kernel.SetArgument(16, static_cast<int>(c_stride));
kernel.SetArgument(17, static_cast<int>(c_do_transpose));
kernel.SetArgument(18, static_cast<int>(a_conjugate));
kernel.SetArgument(19, static_cast<int>(b_conjugate));
// Computes the global and local thread sizes
const auto m_ceiled = Ceil(m, db_["WGD"]);
@ -281,7 +284,6 @@ void XgemmStridedBatched<T>::BatchedGemmDirect(const size_t m, const size_t n, c
// Launches the kernel
RunKernel(kernel, queue_, device_, global, local, event_);
*/
}
// =================================================================================================