Implemented direct version of strided-batched GEMM kernel
parent
9fb2c61b25
commit
13f0f6fc6e
|
@ -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
|
||||
|
|
|
@ -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_);
|
||||
*/
|
||||
}
|
||||
|
||||
// =================================================================================================
|
||||
|
|
Loading…
Reference in New Issue