177 lines
6.8 KiB
Plaintext
177 lines
6.8 KiB
Plaintext
/*******************************************************************************
|
|
* Copyright (c) 2015-2018 Skymind, Inc.
|
|
*
|
|
* This program and the accompanying materials are made available under the
|
|
* terms of the Apache License, Version 2.0 which is available at
|
|
* https://www.apache.org/licenses/LICENSE-2.0.
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
* License for the specific language governing permissions and limitations
|
|
* under the License.
|
|
*
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
******************************************************************************/
|
|
|
|
//
|
|
// @author raver119@gmail.com
|
|
// @author Yurii Shyrma (iuriish@yahoo.com)
|
|
//
|
|
|
|
#include <exceptions/cuda_exception.h>
|
|
#include <cublas_v2.h>
|
|
#include <specials_cuda.h>
|
|
#include <op_boilerplate.h>
|
|
#include <types/float16.h>
|
|
#include <ops/declarable/helpers/batched_gemm.h>
|
|
#include <PointersManager.h>
|
|
|
|
|
|
namespace nd4j {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
//////////////////////////////////////////////////////////////////////////////
|
|
// bsxMXK x bSxKxN = bSxMxN
|
|
void bgemm(const std::vector<NDArray*>& vA, const std::vector<NDArray*>& vB, std::vector<NDArray*>& vC, const NDArray* alphas, const NDArray* betas, int transA, int transB, int M, int N, int K, const int lda, const int ldb, const int ldc) {
|
|
|
|
const auto bS = vA.size(); // batch size
|
|
|
|
std::vector<NDArray*> pA(bS), pB(bS), pC(bS);
|
|
|
|
std::vector<NDArray*> toDelete;
|
|
|
|
for(int i = 0; i < bS; ++i) {
|
|
|
|
if(vA[i]->ews() != 1) {
|
|
pA[i] = new NDArray(vA[i]->dup('f'));
|
|
toDelete.emplace_back(pA[i]);
|
|
}
|
|
else
|
|
pA[i] = vA[i];
|
|
|
|
if(vB[i]->ews() != 1) {
|
|
pB[i] = new NDArray(vB[i]->dup('f'));
|
|
toDelete.emplace_back(pB[i]);
|
|
}
|
|
else
|
|
pB[i] = vB[i];
|
|
|
|
if(vC[i]->ews() != 1) {
|
|
pC[i] = new NDArray(vC[i]->dup('f'));
|
|
toDelete.emplace_back(pC[i]);
|
|
}
|
|
else
|
|
pC[i] = vC[i];
|
|
|
|
if(pC[i]->ordering() != 'f') {
|
|
auto temp = pA[i];
|
|
pA[i] = new NDArray(pB[i]->permute({1,0}));
|
|
pB[i] = new NDArray(temp ->permute({1,0}));
|
|
pC[i] = new NDArray(pC[i]->permute({1,0}));
|
|
toDelete.push_back(pA[i]);
|
|
toDelete.push_back(pB[i]);
|
|
toDelete.push_back(pC[i]);
|
|
M = pA[i]->sizeAt(0);
|
|
K = pA[i]->sizeAt(1);
|
|
N = pB[i]->sizeAt(1);
|
|
}
|
|
|
|
NDArray::prepareSpecialUse ({pC[i]}, {pA[i], pB[i]});
|
|
NDArray::registerSpecialUse({pC[i]}, {pA[i], pB[i]});
|
|
}
|
|
|
|
NDArray::prepareSpecialUse ({}, {alphas, betas});
|
|
NDArray::registerSpecialUse({}, {alphas, betas});
|
|
|
|
std::vector<void*> pAbuffs(bS), pBbuffs(bS), pCbuffs(bS);
|
|
for(int i = 0; i < bS; ++i) {
|
|
pAbuffs[i] = pA[i]->getSpecialBuffer();
|
|
pBbuffs[i] = pB[i]->getSpecialBuffer();
|
|
pCbuffs[i] = pC[i]->getSpecialBuffer();
|
|
}
|
|
|
|
nd4j::LaunchContext* context = vA[0]->getContext();
|
|
PointersManager manager(context, "helpers::bgemm cuda");
|
|
|
|
const void** aBuffers = reinterpret_cast<const void**>(manager.replicatePointer(pAbuffs.data(), bS * sizeof(void*)));
|
|
const void** bBuffers = reinterpret_cast<const void**>(manager.replicatePointer(pBbuffs.data(), bS * sizeof(void*)));
|
|
void** cBuffers = reinterpret_cast<void**>(manager.replicatePointer(pCbuffs.data(), bS * sizeof(void*)));
|
|
|
|
// const auto aOrder = pA->ordering();
|
|
// const auto bOrder = pB->ordering();
|
|
|
|
// const bool transA = aOrder != 'f';
|
|
// const bool transB = bOrder != 'f';
|
|
|
|
const cublasOperation_t transAblas = transA == 112 ? CUBLAS_OP_T : CUBLAS_OP_N;
|
|
const cublasOperation_t transBblas = transB == 112 ? CUBLAS_OP_T : CUBLAS_OP_N;
|
|
|
|
// const int lda = aOrder == 'f' ? M : K;
|
|
// const int ldb = bOrder == 'f' ? K : N;
|
|
// const int ldc = M; // cOrder == 'f' ? M : N;
|
|
|
|
const auto aType = pA[0]->dataType();
|
|
const auto bType = pB[0]->dataType();
|
|
const auto cType = pC[0]->dataType();
|
|
|
|
auto handle = reinterpret_cast<cublasHandle_t*>(context->getCublasHandle());
|
|
auto stream = context->getCudaStream();
|
|
|
|
auto status = cublasSetStream_v2(*handle, *stream);
|
|
|
|
if (status != CUBLAS_STATUS_SUCCESS)
|
|
throw cuda_exception::build("MmulHelper::mmulMxM cuda failed !", status);
|
|
|
|
const bool AB(aType == bType), AC(aType == cType), ABC(AB && AC);
|
|
|
|
// choose appropriate cuda gemm api depending on data types
|
|
if(ABC && aType == DataType::DOUBLE) {
|
|
double alpha = alphas->e<double>(0);
|
|
double beta = betas->e<double>(0);
|
|
status = cublasDgemmBatched(*handle, transAblas, transBblas, M, N, K, &alpha, (const double**)aBuffers, lda, (const double**)bBuffers, ldb, &beta, (double**)cBuffers, ldc, bS);
|
|
}
|
|
else if(ABC && aType == DataType::FLOAT32) {
|
|
float alpha = alphas->e<float>(0);
|
|
float beta = betas->e<float>(0);
|
|
status = cublasSgemmBatched(*handle, transAblas, transBblas, M, N, K, &alpha, (const float**)aBuffers, lda, (const float**)bBuffers, ldb, &beta, (float**)cBuffers, ldc, bS);
|
|
}
|
|
else if(ABC && aType == DataType::HALF) {
|
|
__half alpha = alphas->e<float>(0);
|
|
__half beta = betas->e<float>(0);
|
|
status = cublasHgemmBatched(*handle, transAblas, transBblas, M, N, K, &alpha, (const __half**)aBuffers, lda, (const __half**)bBuffers, ldb, &beta, (__half**)cBuffers, ldc, bS);
|
|
}
|
|
else if(AB && aType == DataType::INT8 && cType == DataType::FLOAT32) {
|
|
float alpha = alphas->e<float>(0);
|
|
float beta = betas->e<float>(0);
|
|
status = cublasGemmBatchedEx(*handle, transAblas, transBblas, M, N, K, &alpha, aBuffers, CUDA_R_8I, lda, bBuffers, CUDA_R_8I, ldb, &beta, cBuffers, CUDA_R_32F, ldc, bS, CUDA_R_32F, CUBLAS_GEMM_DEFAULT);
|
|
}
|
|
else if(AB && aType == DataType::HALF && cType == DataType::FLOAT32) {
|
|
float alpha = alphas->e<float>(0);
|
|
float beta = betas->e<float>(0);
|
|
status = cublasGemmBatchedEx(*handle, transAblas, transBblas, M, N, K, &alpha, aBuffers, CUDA_R_16F, lda, bBuffers, CUDA_R_16F, ldb, &beta, cBuffers, CUDA_R_32F, ldc, bS, CUDA_R_32F, CUBLAS_GEMM_DEFAULT);
|
|
}
|
|
else
|
|
throw std::runtime_error("batched gemm cuda: this mode is not implemented yet !");
|
|
|
|
if (status != CUBLAS_STATUS_SUCCESS)
|
|
throw cuda_exception::build("MmulHelper::mmulMxM cuda failed !", status);
|
|
|
|
auto cudaResult = cudaStreamSynchronize(*stream);
|
|
if (cudaResult != 0)
|
|
throw cuda_exception::build("MmulHelper::mmulMxM cuda failed !", cudaResult);
|
|
|
|
for(int i = 0; i < bS; ++i)
|
|
if(vC[i]->ews() != 1)
|
|
vC[i]->assign(pC[i]);
|
|
|
|
for(int i = toDelete.size() - 1; i >= 0; --i)
|
|
delete toDelete[i];
|
|
}
|
|
|
|
}
|
|
}
|
|
}
|
|
|