wehub-resource-sync
201 lines
6.3 KiB
C++
201 lines
6.3 KiB
C++
// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://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.
|
|
|
|
#pragma once
|
|
|
|
#if !defined(PADDLE_WITH_ARM) && !defined(PADDLE_WITH_SW) && \
|
|
!defined(PADDLE_WITH_MIPS) && !defined(PADDLE_WITH_LOONGARCH)
|
|
#include <immintrin.h>
|
|
#endif
|
|
#include <cfloat>
|
|
#include <cmath>
|
|
#include <cstring>
|
|
|
|
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
|
|
namespace phi {
|
|
namespace funcs {
|
|
|
|
template <typename DeviceContext, typename T>
|
|
void call_gemm(const funcs::BlasT<DeviceContext, T>& blas,
|
|
const CBLAS_TRANSPOSE TransA,
|
|
const CBLAS_TRANSPOSE TransB,
|
|
const int M,
|
|
const int N,
|
|
const int K,
|
|
const T alpha,
|
|
const T* A,
|
|
const T* B,
|
|
const T beta,
|
|
T* C) {
|
|
int lda = (TransA == CblasNoTrans) ? K : M;
|
|
int ldb = (TransB == CblasNoTrans) ? N : K;
|
|
blas.GEMM(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, C, N);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void call_gemm(const Context& dev_ctx,
|
|
const CBLAS_TRANSPOSE TransA,
|
|
const CBLAS_TRANSPOSE TransB,
|
|
const int M,
|
|
const int N,
|
|
const int K,
|
|
const T alpha,
|
|
const T* A,
|
|
const T* B,
|
|
const T beta,
|
|
T* C) {
|
|
int lda = (TransA == CblasNoTrans) ? K : M;
|
|
int ldb = (TransB == CblasNoTrans) ? N : K;
|
|
// auto& dev_ctx = dev_ctx.template device_context<CPUContext>();
|
|
auto blas = funcs::GetBlas<CPUContext, T>(dev_ctx);
|
|
blas.GEMM(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, C, N);
|
|
}
|
|
|
|
template <typename DeviceContext, typename T>
|
|
void call_gemm_with_lda(const funcs::BlasT<DeviceContext, T>& blas,
|
|
const CBLAS_TRANSPOSE TransA,
|
|
const CBLAS_TRANSPOSE TransB,
|
|
const int M,
|
|
const int N,
|
|
const int K,
|
|
const T alpha,
|
|
const T* A,
|
|
const T* B,
|
|
const T beta,
|
|
T* C,
|
|
int lda) {
|
|
int ldb = (TransB == CblasNoTrans) ? N : K;
|
|
|
|
blas.GEMM(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, C, N);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void call_gemm_batched(const Context& dev_ctx,
|
|
const CBLAS_TRANSPOSE TransA,
|
|
const CBLAS_TRANSPOSE TransB,
|
|
const int M,
|
|
const int N,
|
|
const int K,
|
|
const T alpha,
|
|
const T** A,
|
|
const T** B,
|
|
const T beta,
|
|
T** C,
|
|
const int batch) {
|
|
for (int i = 0; i < batch; ++i) {
|
|
call_gemm(dev_ctx, TransA, TransB, M, N, K, alpha, A[i], B[i], beta, C[i]);
|
|
}
|
|
}
|
|
|
|
#if !defined(PADDLE_WITH_ARM) && !defined(PADDLE_WITH_SW) && \
|
|
!defined(PADDLE_WITH_MIPS) && !defined(PADDLE_WITH_LOONGARCH)
|
|
|
|
#define __m256x __m256
|
|
|
|
static const unsigned int AVX_STEP_SIZE = 8;
|
|
static const unsigned int AVX_CUT_LEN_MASK = 7U;
|
|
|
|
#define _mm256_mul_px _mm256_mul_ps
|
|
#define _mm256_add_px _mm256_add_ps
|
|
#define _mm256_load_px _mm256_loadu_ps
|
|
#define _mm256_store_px _mm256_storeu_ps
|
|
#define _mm256_broadcast_sx _mm256_broadcast_ss
|
|
|
|
#define __m128x __m128
|
|
|
|
static const unsigned int SSE_STEP_SIZE = 2;
|
|
static const unsigned int SSE_CUT_LEN_MASK = 1U;
|
|
|
|
#define _mm_add_px _mm_add_ps
|
|
#define _mm_mul_px _mm_mul_ps
|
|
#define _mm_load_px _mm_loadu_ps
|
|
#define _mm_store_px _mm_storeu_ps
|
|
#define _mm_load1_px _mm_load1_ps
|
|
|
|
#endif
|
|
|
|
template <typename T>
|
|
inline void axpy(const T* x, T* y, size_t len, const T alpha) {
|
|
unsigned int jjj, lll;
|
|
jjj = lll = 0;
|
|
|
|
#ifdef PADDLE_WITH_AVX
|
|
lll = len & ~AVX_CUT_LEN_MASK;
|
|
__m256x mm_alpha = _mm256_broadcast_sx(&alpha);
|
|
for (jjj = 0; jjj < lll; jjj += AVX_STEP_SIZE) {
|
|
_mm256_store_px(
|
|
y + jjj,
|
|
_mm256_add_px(_mm256_load_px(y + jjj),
|
|
_mm256_mul_px(mm_alpha, _mm256_load_px(x + jjj))));
|
|
}
|
|
#elif defined(PADDLE_WITH_ARM) || defined(PADDLE_WITH_SW) || \
|
|
defined(PADDLE_WITH_MIPS) || defined(PADDLE_WITH_LOONGARCH)
|
|
PADDLE_THROW(common::errors::Unimplemented("axpy is not supported"));
|
|
#else
|
|
lll = len & ~SSE_CUT_LEN_MASK;
|
|
__m128x mm_alpha = _mm_load1_px(&alpha);
|
|
for (jjj = 0; jjj < lll; jjj += SSE_STEP_SIZE) {
|
|
_mm_store_px(y + jjj,
|
|
_mm_add_px(_mm_load_px(y + jjj),
|
|
_mm_mul_px(mm_alpha, _mm_load_px(x + jjj))));
|
|
}
|
|
|
|
#endif
|
|
|
|
for (; jjj < len; jjj++) {
|
|
y[jjj] += alpha * x[jjj];
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
inline void axpy_noadd(const T* x, T* y, size_t len, const T alpha) {
|
|
unsigned int jjj, lll;
|
|
jjj = lll = 0;
|
|
|
|
#ifdef PADDLE_WITH_AVX
|
|
lll = len & ~AVX_CUT_LEN_MASK;
|
|
__m256x mm_alpha = _mm256_broadcast_sx(&alpha);
|
|
for (jjj = 0; jjj < lll; jjj += AVX_STEP_SIZE) {
|
|
_mm256_store_px(y + jjj, _mm256_mul_px(mm_alpha, _mm256_load_px(x + jjj)));
|
|
}
|
|
#elif defined(PADDLE_WITH_ARM) || defined(PADDLE_WITH_SW) || \
|
|
defined(PADDLE_WITH_MIPS) || defined(PADDLE_WITH_LOONGARCH)
|
|
PADDLE_THROW(common::errors::Unimplemented("axpy_noadd is not supported"));
|
|
#else
|
|
lll = len & ~SSE_CUT_LEN_MASK;
|
|
__m128x mm_alpha = _mm_load1_px(&alpha);
|
|
for (jjj = 0; jjj < lll; jjj += SSE_STEP_SIZE) {
|
|
_mm_store_px(y + jjj, _mm_mul_px(mm_alpha, _mm_load_px(x + jjj)));
|
|
}
|
|
|
|
#endif
|
|
|
|
for (; jjj < len; jjj++) {
|
|
y[jjj] = alpha * x[jjj];
|
|
}
|
|
}
|
|
|
|
inline void axpy_noadd(const int8_t* x,
|
|
int8_t* y,
|
|
size_t len,
|
|
const float alpha) {
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"int8_t input of axpy_noadd is not supported"));
|
|
}
|
|
|
|
} // namespace funcs
|
|
} // namespace phi
|