wehub-resource-sync
209 lines
8.1 KiB
Plaintext
209 lines
8.1 KiB
Plaintext
/* ******************************************************************************
|
|
*
|
|
*
|
|
* 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.
|
|
*
|
|
* See the NOTICE file distributed with this work for additional
|
|
* information regarding copyright ownership.
|
|
* 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 <cublas_v2.h>
|
|
#include <exceptions/cuda_exception.h>
|
|
#include <helpers/PointersManager.h>
|
|
#include <ops/declarable/helpers/batched_gemm.h>
|
|
#include <ops/specials_cuda.h>
|
|
#include <system/op_boilerplate.h>
|
|
#include <types/float16.h>
|
|
|
|
#include <indexing/NDIndexUtils.h>
|
|
#include <ops/declarable/CustomOperations.h>
|
|
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
|
|
void bgemm(NDArray *a, NDArray *b, NDArray *c, NDArray *alphas, NDArray *betas,
|
|
int transA, int transB, int M, int N, int K, int lda, int ldb, int ldc, NDArray *all) {
|
|
NDArray *allIndex = nullptr;
|
|
if(all != nullptr)
|
|
allIndex = all;
|
|
else {
|
|
NDArray allLocal = NDIndexUtils::createAll();
|
|
all = &allLocal;
|
|
}
|
|
|
|
|
|
int batchSize = a->sizeAt(0) / 2;
|
|
std::vector<NDArray *>inputs;
|
|
std::vector<NDArray *> keyInputs;
|
|
std::vector<NDArray *> outputs;
|
|
|
|
create_view createView;
|
|
|
|
//add alpha and beta before the batch gemm, this just needs to be broadcasted
|
|
inputs.push_back(alphas);
|
|
inputs.push_back(betas);
|
|
|
|
//divide by 2: queries and keys
|
|
for(int i = 0; i < batchSize; i++) {
|
|
auto point = NDIndexUtils::createPoint(i);
|
|
auto aSlice = createView.evaluate({a,&point,all,all},{},{});
|
|
auto bSlice = createView.evaluate({b,&point,all,all},{},{});
|
|
auto outSlice = createView.evaluate({c,&point,all,all},{},{});
|
|
inputs.push_back(aSlice.at(0));
|
|
keyInputs.push_back(bSlice.at(0));
|
|
outputs.push_back(outSlice.at(0));
|
|
}
|
|
|
|
bgemm(inputs,keyInputs,outputs,alphas,betas,transA,transB,M,N,K,lda,ldb,ldc);
|
|
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////////
|
|
// bsxMXK x bSxKxN = bSxMxN
|
|
void bgemm( std::vector<NDArray *> &vA, std::vector<NDArray *> &vB, std::vector<NDArray *> &vC,
|
|
NDArray *alphas, NDArray *betas, int transA, int transB, int M, int N, int K, int lda,
|
|
int ldb, 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) {
|
|
pA[i] = vA[i]->dup('f'); // dup() already returns NDArray*
|
|
toDelete.emplace_back(pA[i]);
|
|
|
|
|
|
pB[i] = vB[i]->dup('f'); // dup() already returns NDArray*
|
|
toDelete.emplace_back(pB[i]);
|
|
|
|
pC[i] = vC[i]->dup('f'); // dup() already returns NDArray*
|
|
toDelete.emplace_back(pC[i]);
|
|
|
|
|
|
if (pC[i]->ordering() != 'f') {
|
|
auto temp = pA[i];
|
|
std::vector<sd::LongType> permute = {1,0};
|
|
pA[i] = pB[i]->permute(permute, false, false); // permute() already returns NDArray*
|
|
pB[i] = temp->permute(permute, false, false);
|
|
pC[i] = pC[i]->permute(permute, false, false);
|
|
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]->specialBuffer();
|
|
pBbuffs[i] = pB[i]->specialBuffer();
|
|
pCbuffs[i] = pC[i]->specialBuffer();
|
|
}
|
|
|
|
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 cublasOperation_t transAblas = transA == 112 ? CUBLAS_OP_T : CUBLAS_OP_N;
|
|
const cublasOperation_t transBblas = transB == 112 ? CUBLAS_OP_T : CUBLAS_OP_N;
|
|
|
|
if(M < 0) THROW_EXCEPTION("M < 0");
|
|
if(N < 0) THROW_EXCEPTION("N < 0");
|
|
if(K < 0) THROW_EXCEPTION("K < 0");
|
|
|
|
|
|
const auto aType = pA[0]->dataType();
|
|
const auto bType = pB[0]->dataType();
|
|
const auto cType = pC[0]->dataType();
|
|
|
|
std::lock_guard<std::mutex> lock(*LaunchContext::deviceMutex());
|
|
|
|
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 set stream failed ! Please double check the passed in handle.", status);
|
|
|
|
const bool AB(aType == bType), AC(aType == cType), ABC(AB && AC);
|
|
|
|
// choose appropriate cuda gemm api depending on data types
|
|
// Handle empty arrays by using default values (alpha=1.0, beta=0.0)
|
|
if (ABC && aType == DOUBLE) {
|
|
double alpha = alphas->lengthOf() > 0 ? alphas->e<double>(0) : 1.0;
|
|
double beta = betas->lengthOf() > 0 ? betas->e<double>(0) : 0.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 == FLOAT32) {
|
|
float alpha = alphas->lengthOf() > 0 ? alphas->e<float>(0) : 1.0f;
|
|
float beta = betas->lengthOf() > 0 ? betas->e<float>(0) : 0.0f;
|
|
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 == HALF) {
|
|
__half alpha = alphas->lengthOf() > 0 ? alphas->e<float>(0) : 1.0f;
|
|
__half beta = betas->lengthOf() > 0 ? betas->e<float>(0) : 0.0f;
|
|
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 == INT8 && cType == FLOAT32) {
|
|
float alpha = alphas->lengthOf() > 0 ? alphas->e<float>(0) : 1.0f;
|
|
float beta = betas->lengthOf() > 0 ? betas->e<float>(0) : 0.0f;
|
|
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 == HALF && cType == FLOAT32) {
|
|
float alpha = alphas->lengthOf() > 0 ? alphas->e<float>(0) : 1.0f;
|
|
float beta = betas->lengthOf() > 0 ? betas->e<float>(0) : 0.0f;
|
|
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_EXCEPTION("batched gemm cuda: this mode is not implemented yet !");
|
|
|
|
if (status != CUBLAS_STATUS_SUCCESS) {
|
|
throw cuda_exception::build("MmulHelper::mmulMxM cuda execution failed !", status);
|
|
}
|
|
|
|
auto cudaResult = cudaStreamSynchronize(*stream);
|
|
if (cudaResult != 0) {
|
|
throw cuda_exception::build("MmulHelper::mmulMxM cuda stream synchronize failed !", cudaResult);
|
|
}
|
|
|
|
for (int i = 0; i < bS; ++i)
|
|
vC[i]->assign(pC[i]);
|
|
|
|
for (int i = toDelete.size() - 1; i >= 0; --i) delete toDelete[i];
|
|
}
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|