chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,202 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
Reference in New Issue
Block a user