wehub-resource-sync
202 lines
8.0 KiB
C++
202 lines
8.0 KiB
C++
/* ******************************************************************************
|
|
*
|
|
*
|
|
* 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
|
|
//
|
|
#include <execution/Threads.h>
|
|
#include <helpers/BlasHelper.h>
|
|
#include <ops/declarable/helpers/batched_gemm.h>
|
|
#include <system/op_boilerplate.h>
|
|
#include <types/float16.h>
|
|
#include <indexing/NDIndexUtils.h>
|
|
#include <ops/declarable/CustomOperations.h>
|
|
|
|
#if NOT_EXCLUDED(OP_batched_gemm)
|
|
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();
|
|
allIndex = allLocal;
|
|
}
|
|
|
|
int batchSize = a->sizeAt(0);
|
|
std::vector<NDArray *>inputs;
|
|
std::vector<NDArray *> bInputs;
|
|
std::vector<NDArray *> outputs;
|
|
|
|
ops::create_view createView;
|
|
|
|
//divide by 2: queries and keys
|
|
for(int i = 0; i < batchSize; i++) {
|
|
auto point = NDIndexUtils::createPoint(i);
|
|
auto aSlice = createView.evaluate({a,point,allIndex,allIndex},{},{});
|
|
auto bSlice = createView.evaluate({b,point,allIndex,allIndex},{},{});
|
|
auto outSlice = createView.evaluate({c,point,allIndex,allIndex},{},{});
|
|
inputs.push_back(aSlice.at(0));
|
|
bInputs.push_back(bSlice.at(0));
|
|
outputs.push_back(outSlice.at(0));
|
|
delete point;
|
|
}
|
|
|
|
delete allIndex;
|
|
|
|
|
|
|
|
bgemm(inputs, bInputs,outputs,alphas,betas,transA,transB,M,N,K,lda,ldb,ldc);
|
|
|
|
}
|
|
|
|
template <typename T>
|
|
static 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) {
|
|
int batchSize = vA.size();
|
|
|
|
|
|
|
|
// Use batched BLAS only when: 1) batched GEMM is available AND 2) BLAS is enabled
|
|
// Previously used || which incorrectly entered BLAS path when BLAS was disabled
|
|
if (BlasHelper::getInstance().hasBatchedGEMM<T>() && Environment::getInstance().isEnableBlas()) {
|
|
auto arr = vA.at(0);
|
|
CBLAS_TRANSPOSE *tA, *tB;
|
|
int *tM, *tN, *tK, *tldA, *tldB, *tldC, *tsize;
|
|
// mkl requires mnk etc as arrays, cuda doesn't
|
|
ALLOCATE(tA, arr->getContext()->getWorkspace(), batchSize, CBLAS_TRANSPOSE);
|
|
ALLOCATE(tB, arr->getContext()->getWorkspace(), batchSize, CBLAS_TRANSPOSE);
|
|
ALLOCATE(tM, arr->getContext()->getWorkspace(), batchSize, int);
|
|
ALLOCATE(tN, arr->getContext()->getWorkspace(), batchSize, int);
|
|
ALLOCATE(tK, arr->getContext()->getWorkspace(), batchSize, int);
|
|
ALLOCATE(tldA, arr->getContext()->getWorkspace(), batchSize, int);
|
|
ALLOCATE(tldB, arr->getContext()->getWorkspace(), batchSize, int);
|
|
ALLOCATE(tldC, arr->getContext()->getWorkspace(), batchSize, int);
|
|
ALLOCATE(tsize, arr->getContext()->getWorkspace(), batchSize, int);
|
|
|
|
shape::fill(tA, (CBLAS_TRANSPOSE)transA, batchSize);
|
|
shape::fill(tB, (CBLAS_TRANSPOSE)transB, batchSize);
|
|
|
|
shape::fill(tM, M, batchSize);
|
|
shape::fill(tN, N, batchSize);
|
|
shape::fill(tK, K, batchSize);
|
|
shape::fill(tldA, lda, batchSize);
|
|
shape::fill(tldB, ldb, batchSize);
|
|
shape::fill(tldC, ldc, batchSize);
|
|
shape::fill(tsize, 1, batchSize);
|
|
|
|
std::vector<T *> buffersA;
|
|
std::vector<T *> buffersB;
|
|
std::vector<T *> buffersC;
|
|
|
|
|
|
|
|
for (int e = 0; e < batchSize; e++) {
|
|
buffersA.push_back(reinterpret_cast<T *>(vA[e]->buffer()));
|
|
buffersB.push_back(reinterpret_cast<T *>(vB[e]->buffer()));
|
|
buffersC.push_back(reinterpret_cast<T *>(vC[e]->buffer()));
|
|
}
|
|
|
|
// Acquire BLAS lock to prevent OpenBLAS TLS corruption and race conditions
|
|
auto blasLock = BlasHelper::getInstance().lockBlas();
|
|
|
|
// Inside BLAS block, only check type - BLAS enablement was already verified in outer condition
|
|
if (std::is_same<T, double>::value) {
|
|
BlasHelper::getInstance().dgemmBatched()(CblasColMajor, tA, tB, tM, tN, tK, (double *)alphas->buffer(),
|
|
(double **)buffersA.data(), tldA, (double **)buffersB.data(), tldB,
|
|
(double *)betas->buffer(), (double **)buffersC.data(), tldC, vA.size(),
|
|
tsize);
|
|
} else if (std::is_same<T, float>::value) {
|
|
BlasHelper::getInstance().sgemmBatched()(
|
|
CblasColMajor, tA, tB, tM, tN, tK, (float *)alphas->buffer(), (float **)buffersA.data(), tldA,
|
|
(float **)buffersB.data(), tldB, (float *)betas->buffer(), (float **)buffersC.data(), tldC, vA.size(), tsize);
|
|
}
|
|
|
|
|
|
|
|
// release temporary arrays
|
|
RELEASE(tA, arr->getContext()->getWorkspace());
|
|
RELEASE(tB, arr->getContext()->getWorkspace());
|
|
RELEASE(tM, arr->getContext()->getWorkspace());
|
|
RELEASE(tN, arr->getContext()->getWorkspace());
|
|
RELEASE(tK, arr->getContext()->getWorkspace());
|
|
RELEASE(tldA, arr->getContext()->getWorkspace());
|
|
RELEASE(tldB, arr->getContext()->getWorkspace());
|
|
RELEASE(tldC, arr->getContext()->getWorkspace());
|
|
RELEASE(tsize, arr->getContext()->getWorkspace());
|
|
} else {
|
|
|
|
CBLAS_TRANSPOSE tA = (CBLAS_TRANSPOSE)transA;
|
|
CBLAS_TRANSPOSE tB = (CBLAS_TRANSPOSE)transB;
|
|
int vaSize = vA.size();
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto p = start; p < stop; p++) {
|
|
auto A = reinterpret_cast<T *>(vA.at(p)->buffer());
|
|
auto B = reinterpret_cast<T *>(vB.at(p)->buffer());
|
|
auto C = reinterpret_cast<T *>(vC.at(p)->buffer());
|
|
// Handle scalar, single-element, or empty arrays (use defaults for empty)
|
|
auto alpha = (alphas->isScalar() || alphas->lengthOf() <= 1)
|
|
? (alphas->lengthOf() > 0 ? alphas->e<T>(0) : static_cast<T>(1))
|
|
: alphas->e<T>(p);
|
|
auto beta = (betas->isScalar() || betas->lengthOf() <= 1)
|
|
? (betas->lengthOf() > 0 ? betas->e<T>(0) : static_cast<T>(0))
|
|
: betas->e<T>(p);
|
|
for (int m = 0; m < M; m++) {
|
|
for (int n = 0; n < N; n++) {
|
|
T c_mnp = static_cast<T>(0);
|
|
PRAGMA_OMP_SIMD
|
|
for (int k = 0; k < K; k++) {
|
|
c_mnp += A[tA == CblasNoTrans ? (m + k * lda) : (m * lda + k)] *
|
|
B[tB == CblasNoTrans ? (k + n * ldb) : (k * ldb + n)];
|
|
}
|
|
C[m + n * ldc] = alpha * c_mnp + beta * C[m + n * ldc];
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_tad(func, 0, vaSize);
|
|
|
|
}
|
|
}
|
|
|
|
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) {
|
|
auto xType = vA.at(0)->dataType();
|
|
BUILD_SINGLE_SELECTOR(xType, bgemm_, (vA, vB, vC, alphas, betas, transA, transB, M, N, K, lda, ldb, ldc),
|
|
SD_FLOAT_TYPES);
|
|
}
|
|
|
|
BUILD_SINGLE_TEMPLATE( 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),
|
|
SD_FLOAT_TYPES);
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|
|
#endif |