wehub-resource-sync
519 lines
17 KiB
C++
519 lines
17 KiB
C++
// Copyright (c) 2023 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.
|
|
|
|
#include <immintrin.h>
|
|
#include <omp.h>
|
|
#include <cassert>
|
|
#include <cmath>
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <iostream>
|
|
#include <new>
|
|
#include <string>
|
|
#include "glog/logging.h"
|
|
|
|
#ifdef PADDLE_WITH_DNNL
|
|
#include "dnnl.hpp" //NOLINT
|
|
#endif
|
|
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/tensor_utils.h"
|
|
|
|
namespace phi {
|
|
namespace fusion {
|
|
|
|
template <typename T, typename Tt>
|
|
void arraycpy(T* dst, const Tt* src, int n) {
|
|
#ifdef PADDLE_WITH_MKLML
|
|
#pragma omp simd
|
|
#endif
|
|
for (int i = 0; i < n; i++) {
|
|
dst[i] = static_cast<T>(src[i]);
|
|
}
|
|
}
|
|
|
|
// batches x tokens x 3 x head x heads -> 3 x batches x head x tokens x heads
|
|
// (2 0 3 1 4)
|
|
template <typename T, typename Tt>
|
|
void transpose_before_bmm1(const T* qkvBuffer,
|
|
Tt* qkvTransBuffer,
|
|
int batchSize,
|
|
int tokenSize,
|
|
int headNum,
|
|
int headSize) {
|
|
int hiddenSize = headNum * headSize;
|
|
int blocksize = tokenSize * hiddenSize; // dst buffer stride in each batch
|
|
|
|
const T* qBuffer = qkvBuffer;
|
|
const T* kBuffer = qkvBuffer + hiddenSize;
|
|
const T* vBuffer = qkvBuffer + hiddenSize * 2;
|
|
|
|
Tt* q_buffer = qkvTransBuffer;
|
|
Tt* k_buffer = qkvTransBuffer + batchSize * blocksize;
|
|
Tt* v_buffer = qkvTransBuffer + batchSize * blocksize * 2;
|
|
|
|
int bmHead = headNum;
|
|
int cols_per_bmHead = hiddenSize / headNum; // 768/12 = 64
|
|
|
|
#ifdef PADDLE_WITH_MKLML
|
|
#pragma omp parallel for collapse(3)
|
|
#endif
|
|
for (int i = 0; i < batchSize; i++) {
|
|
for (int k = 0; k < bmHead; k++) {
|
|
for (int j = 0; j < tokenSize; j++) {
|
|
const T* q_src_each_batch =
|
|
reinterpret_cast<const T*>(qBuffer) + blocksize * 3 * i;
|
|
const T* k_src_each_batch =
|
|
reinterpret_cast<const T*>(kBuffer) + blocksize * 3 * i;
|
|
const T* v_src_each_batch =
|
|
reinterpret_cast<const T*>(vBuffer) + blocksize * 3 * i;
|
|
|
|
int dst_offset_each_bmHead = k * tokenSize * cols_per_bmHead;
|
|
int src_offset_each_line = k * cols_per_bmHead;
|
|
|
|
int dst_offset_each_line = j * cols_per_bmHead;
|
|
int src_offset_each_bmHead = j * hiddenSize * 3;
|
|
|
|
Tt* q_dst_each_line = q_buffer + i * blocksize +
|
|
dst_offset_each_bmHead + dst_offset_each_line;
|
|
const T* q_src_each_line =
|
|
q_src_each_batch + src_offset_each_bmHead + src_offset_each_line;
|
|
|
|
Tt* k_dst_each_line = k_buffer + i * blocksize +
|
|
dst_offset_each_bmHead + dst_offset_each_line;
|
|
const T* k_src_each_line =
|
|
k_src_each_batch + src_offset_each_bmHead + src_offset_each_line;
|
|
|
|
Tt* v_dst_each_line = v_buffer + i * blocksize +
|
|
dst_offset_each_bmHead + dst_offset_each_line;
|
|
const T* v_src_each_line =
|
|
v_src_each_batch + src_offset_each_bmHead + src_offset_each_line;
|
|
arraycpy<Tt, T>(q_dst_each_line, q_src_each_line, cols_per_bmHead);
|
|
arraycpy<Tt, T>(k_dst_each_line, k_src_each_line, cols_per_bmHead);
|
|
arraycpy<Tt, T>(v_dst_each_line, v_src_each_line, cols_per_bmHead);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// batches x head x tokens x heads -> batches x tokens x head x heads (0 2 1 3)
|
|
template <typename T, typename Tt>
|
|
void transpose_after_bmm2(T* Buffer,
|
|
Tt* TransBuffer,
|
|
int batchSize,
|
|
int tokenSize,
|
|
int headNum,
|
|
int headSize) {
|
|
int hiddenSize = headNum * headSize;
|
|
int blocksize = tokenSize * hiddenSize; // dst buffer stride in each batch
|
|
|
|
int bmHead = headNum;
|
|
int cols_per_bmHead = hiddenSize / headNum; // 768/12 = 64
|
|
|
|
#ifdef PADDLE_WITH_MKLML
|
|
#pragma omp parallel for collapse(2)
|
|
#endif
|
|
for (int i = 0; i < batchSize; i++) {
|
|
for (int k = 0; k < tokenSize; k++) {
|
|
int src_offset_each_head = k * cols_per_bmHead;
|
|
int dst_offset_each_line = k * hiddenSize;
|
|
|
|
for (int j = 0; j < bmHead; j++) {
|
|
int src_offset_each_line = j * tokenSize * cols_per_bmHead;
|
|
int dst_offset_each_head = j * cols_per_bmHead;
|
|
|
|
Tt* q_dst_each_line = TransBuffer + dst_offset_each_head +
|
|
dst_offset_each_line + i * blocksize;
|
|
const T* q_src_each_line = Buffer + src_offset_each_line +
|
|
src_offset_each_head + i * blocksize;
|
|
|
|
arraycpy<Tt, T>(q_dst_each_line, q_src_each_line, cols_per_bmHead);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// C = A * B
|
|
// bTranspose: B need to be transposed or not
|
|
void sgemm(const float* A,
|
|
const float* B,
|
|
float* C,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
bool transa,
|
|
bool transb) {
|
|
#ifdef PADDLE_WITH_DNNL
|
|
int lda = (transa ? m : k);
|
|
int ldb = (transb ? k : n);
|
|
int ldc = n;
|
|
float alpha = 1;
|
|
float beta = 0;
|
|
std::array<char, 2> ta = {"N"};
|
|
std::array<char, 2> tb = {"N"};
|
|
if (transa) ta[0] = 'T';
|
|
if (transb) tb[0] = 'T';
|
|
|
|
dnnl_sgemm(ta[0], tb[0], m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
|
|
#else
|
|
LOG(ERROR) << "scaled_dp_atten not supported without WITH_MKL!";
|
|
#endif
|
|
}
|
|
|
|
// exp based-on jit code
|
|
static inline __m512 vexp(const __m512& _x) {
|
|
__m512 p16f_1 = _mm512_set1_ps(1.0f);
|
|
__m512 p16f_half = _mm512_set1_ps(0.5f);
|
|
__m512 p16f_127 = _mm512_set1_ps(127.f);
|
|
__m512 p16f_exp_hi = _mm512_set1_ps(88.3762626647950f);
|
|
__m512 p16f_exp_lo = _mm512_set1_ps(-88.3762626647949f);
|
|
|
|
__m512 p16f_cephes_LOG2EF = _mm512_set1_ps(1.44269504088896341f);
|
|
|
|
__m512 p16f_cephes_exp_p0 = _mm512_set1_ps(1.9875691500E-4f);
|
|
__m512 p16f_cephes_exp_p1 = _mm512_set1_ps(1.3981999507E-3f);
|
|
__m512 p16f_cephes_exp_p2 = _mm512_set1_ps(8.3334519073E-3f);
|
|
__m512 p16f_cephes_exp_p3 = _mm512_set1_ps(4.1665795894E-2f);
|
|
__m512 p16f_cephes_exp_p4 = _mm512_set1_ps(1.6666665459E-1f);
|
|
__m512 p16f_cephes_exp_p5 = _mm512_set1_ps(5.0000001201E-1f);
|
|
|
|
// Clamp x.
|
|
__m512 x = _mm512_max_ps(_mm512_min_ps(_x, p16f_exp_hi), p16f_exp_lo);
|
|
|
|
// Express exp(x) as exp(m*ln(2) + r), start by extracting
|
|
// m = floor(x/ln(2) + 0.5).
|
|
__m512 m = _mm512_floor_ps(_mm512_fmadd_ps(x, p16f_cephes_LOG2EF, p16f_half));
|
|
|
|
// Get r = x - m*ln(2).
|
|
__m512 p16f_nln2 = _mm512_set1_ps(-0.6931471805599453f);
|
|
__m512 r = _mm512_fmadd_ps(m, p16f_nln2, x);
|
|
|
|
__m512 r2 = _mm512_mul_ps(r, r);
|
|
|
|
__m512 y = p16f_cephes_exp_p0;
|
|
y = _mm512_fmadd_ps(y, r, p16f_cephes_exp_p1);
|
|
y = _mm512_fmadd_ps(y, r, p16f_cephes_exp_p2);
|
|
y = _mm512_fmadd_ps(y, r, p16f_cephes_exp_p3);
|
|
y = _mm512_fmadd_ps(y, r, p16f_cephes_exp_p4);
|
|
y = _mm512_fmadd_ps(y, r, p16f_cephes_exp_p5);
|
|
y = _mm512_fmadd_ps(y, r2, r);
|
|
y = _mm512_add_ps(y, p16f_1);
|
|
|
|
// Build emm0 = 2^m.
|
|
__m512i emm0 = _mm512_cvttps_epi32(_mm512_add_ps(m, p16f_127));
|
|
emm0 = _mm512_slli_epi32(emm0, 23);
|
|
|
|
// Return 2^m * exp(r).
|
|
return _mm512_max_ps(_mm512_mul_ps(y, _mm512_castsi512_ps(emm0)), _x);
|
|
}
|
|
|
|
// need to do for res.
|
|
void softmax_sum_max(float* AB,
|
|
float* sum,
|
|
float* max,
|
|
float* pre_sum,
|
|
float* pre_max,
|
|
float refac,
|
|
int m,
|
|
int k) {
|
|
float max_val = std::numeric_limits<float>::lowest();
|
|
__m512 vrefac = _mm512_set1_ps(refac);
|
|
for (int i = 0; i < m; ++i) {
|
|
float* buf = AB + i * k;
|
|
// max val for avoiding inf and nan
|
|
__m512 vmax = _mm512_set1_ps(max_val);
|
|
for (int off = 0; off < k; off += 16) {
|
|
int remain = k - off;
|
|
__mmask16 mask = (remain >= 16 ? 0xffff : (1 << remain) - 1);
|
|
__m512 vx = _mm512_maskz_loadu_ps(mask, buf + off);
|
|
|
|
vmax = _mm512_mask_max_ps(vmax, mask, vmax, vx);
|
|
}
|
|
float _max = _mm512_reduce_max_ps(vmax);
|
|
|
|
_max *= refac;
|
|
_max = _max > max[i] ? _max : max[i];
|
|
__m512 merr = _mm512_set1_ps(max[i] - _max);
|
|
merr = vexp(merr);
|
|
max[i] = _max;
|
|
|
|
// exp and get sum
|
|
__m512 vsum = _mm512_set1_ps(0);
|
|
vmax = _mm512_set1_ps(_max);
|
|
for (int off = 0; off < k; off += 16) {
|
|
int remain = k - off;
|
|
__mmask16 mask = (remain >= 16 ? 0xffff : (1 << remain) - 1);
|
|
|
|
__m512 vx = _mm512_maskz_loadu_ps(mask, buf + off);
|
|
vx = _mm512_mask_mul_ps(vx, mask, vx, vrefac);
|
|
vx = _mm512_mask_sub_ps(vx, mask, vx, vmax);
|
|
vx = vexp(vx);
|
|
|
|
_mm512_mask_storeu_ps(buf + off, mask, vx);
|
|
|
|
vsum = _mm512_mask_add_ps(vsum, mask, vsum, vx);
|
|
}
|
|
float _sum = _mm512_reduce_add_ps(vsum);
|
|
float fac = _mm512_cvtss_f32(merr);
|
|
sum[i] = sum[i] * fac + _sum;
|
|
_sum = sum[i];
|
|
|
|
// Compute exp/sum(exp) and store
|
|
__m512 vrsum = _mm512_set1_ps(1.0f / _sum);
|
|
for (int off = 0; off < k; off += 16) {
|
|
int remain = k - off;
|
|
__mmask16 mask = (remain >= 16 ? 0xffff : (1 << remain) - 1);
|
|
|
|
__m512 vx = _mm512_maskz_loadu_ps(mask, buf + off);
|
|
vx = _mm512_mask_mul_ps(vx, mask, vx, vrsum);
|
|
_mm512_mask_storeu_ps(buf + off, mask, vx);
|
|
}
|
|
}
|
|
}
|
|
|
|
void update_out_blk(float* output,
|
|
const float* exp_ABC,
|
|
float* pre_sum,
|
|
float* sum,
|
|
float* pre_max,
|
|
float* max,
|
|
int m,
|
|
int n) {
|
|
for (int i = 0; i < m; ++i) {
|
|
const float* buf = exp_ABC + i * n;
|
|
float* outbuf = output + i * n;
|
|
__m512 merr = _mm512_set1_ps(pre_max[i] - max[i]);
|
|
merr = vexp(merr);
|
|
__m512 vfac = _mm512_set1_ps(pre_sum[i] / sum[i]);
|
|
for (int off = 0; off < n; off += 16) {
|
|
int remain = n - off;
|
|
__mmask16 mask = (remain >= 16 ? 0xffff : (1 << remain) - 1);
|
|
__m512 vout = _mm512_maskz_loadu_ps(mask, outbuf + off);
|
|
__m512 vabc = _mm512_maskz_loadu_ps(mask, buf + off);
|
|
vout = _mm512_mask_mul_ps(vout, mask, vout, merr);
|
|
vout = _mm512_mask_mul_ps(vout, mask, vout, vfac);
|
|
__m512 vupt = _mm512_set1_ps(0.0f);
|
|
vupt = _mm512_mask_add_ps(vupt, mask, vout, vabc);
|
|
_mm512_mask_storeu_ps(outbuf + off, mask, vupt);
|
|
}
|
|
pre_sum[i] = sum[i];
|
|
pre_max[i] = max[i];
|
|
}
|
|
}
|
|
|
|
// hard code: axis = 1
|
|
// sum += sum(exp(A[i]))
|
|
// output = output * pre_sum / sum + (exp(A) / sum) x B
|
|
// pre_sum = sum
|
|
void incremental_tile_attention(const float* A,
|
|
const float* B,
|
|
const float* C,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
float* pre_sum,
|
|
float* sum,
|
|
float* pre_max,
|
|
float* max,
|
|
float refac,
|
|
float* AB,
|
|
float* exp_ABC,
|
|
float* output) {
|
|
sgemm(A, B, AB, m, k, n, false, true);
|
|
softmax_sum_max(AB, sum, max, pre_sum, pre_max, refac, m, k);
|
|
sgemm(AB, C, exp_ABC, m, n, k, false, false);
|
|
update_out_blk(output, exp_ABC, pre_sum, sum, pre_max, max, m, n);
|
|
}
|
|
|
|
// scaled dot-product attention: bmm1 + softmax + bmm2
|
|
void scaled_dp_attention(const float* query,
|
|
const float* key,
|
|
const float* value,
|
|
float scale,
|
|
int batch_size,
|
|
int itsize,
|
|
int otsize,
|
|
int num_head,
|
|
int head_size,
|
|
float* output) {
|
|
// output = trans(softmax(query * trans(key)) * value)
|
|
int iblk = std::min(512, itsize / 1);
|
|
int oblk = std::min(512, otsize / 1);
|
|
float refac = scale;
|
|
|
|
#ifdef PADDLE_WITH_MKLML
|
|
int nth = omp_get_max_threads();
|
|
#else
|
|
int nth = 1;
|
|
#endif
|
|
|
|
float** pre_sum;
|
|
float** sum;
|
|
float** pre_max;
|
|
float** max;
|
|
float** qk_arr;
|
|
float** exp_qkv_arr;
|
|
pre_sum = new float*[nth];
|
|
sum = new float*[nth];
|
|
pre_max = new float*[nth];
|
|
max = new float*[nth];
|
|
qk_arr = new float*[nth];
|
|
exp_qkv_arr = new float*[nth];
|
|
for (int i = 0; i < nth; ++i) {
|
|
pre_sum[i] = new float[iblk];
|
|
sum[i] = new float[iblk];
|
|
pre_max[i] = new float[iblk];
|
|
max[i] = new float[iblk];
|
|
qk_arr[i] = new float[iblk * oblk];
|
|
exp_qkv_arr[i] = new float[iblk * head_size];
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_MKLML
|
|
#pragma omp parallel for collapse(3)
|
|
#endif
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
for (int j = 0; j < num_head; ++j) {
|
|
for (int m = 0; m < itsize; m += iblk) {
|
|
#ifdef PADDLE_WITH_MKLML
|
|
int tid = omp_get_thread_num();
|
|
#else
|
|
int tid = 0;
|
|
#endif
|
|
int ooffset =
|
|
i * num_head * otsize * head_size + j * otsize * head_size;
|
|
const float* k = key + ooffset;
|
|
const float* v = value + ooffset;
|
|
|
|
int q_rblk = std::min(iblk, itsize - m);
|
|
int ioffset =
|
|
i * num_head * otsize * head_size + j * otsize * head_size;
|
|
const float* q = query + ioffset + m * head_size;
|
|
float* out = output + ioffset + m * head_size;
|
|
|
|
// reset out
|
|
for (int ii = 0; ii < q_rblk; ++ii) {
|
|
#ifdef PADDLE_WITH_MKLML
|
|
#pragma omp simd
|
|
#endif
|
|
for (int jj = 0; jj < head_size; ++jj) {
|
|
out[ii * head_size + jj] = 0; // reset output
|
|
}
|
|
}
|
|
// reset sum
|
|
#ifdef PADDLE_WITH_MKLML
|
|
#pragma omp simd
|
|
#endif
|
|
for (int ii = 0; ii < q_rblk; ++ii) {
|
|
pre_sum[tid][ii] = 0;
|
|
sum[tid][ii] = 0;
|
|
pre_max[tid][ii] = std::numeric_limits<float>::lowest();
|
|
max[tid][ii] = std::numeric_limits<float>::lowest();
|
|
}
|
|
//
|
|
for (int b = 0; b < otsize; b += oblk) {
|
|
int kv_rblk = std::min(oblk, otsize - b);
|
|
const float* blk_k = k + b * head_size;
|
|
const float* blk_v = v + b * head_size;
|
|
|
|
incremental_tile_attention(q,
|
|
blk_k,
|
|
blk_v,
|
|
q_rblk,
|
|
head_size,
|
|
kv_rblk,
|
|
pre_sum[tid],
|
|
sum[tid],
|
|
pre_max[tid],
|
|
max[tid],
|
|
refac,
|
|
qk_arr[tid],
|
|
exp_qkv_arr[tid],
|
|
out);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
for (int i = 0; i < nth; ++i) {
|
|
delete[] pre_sum[i];
|
|
delete[] sum[i];
|
|
delete[] pre_max[i];
|
|
delete[] max[i];
|
|
delete[] qk_arr[i];
|
|
delete[] exp_qkv_arr[i];
|
|
}
|
|
delete[] pre_sum;
|
|
delete[] sum;
|
|
delete[] pre_max;
|
|
delete[] max;
|
|
delete[] qk_arr;
|
|
delete[] exp_qkv_arr;
|
|
|
|
return;
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void SelfDPAttenKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const float alpha,
|
|
const int head_number,
|
|
DenseTensor* out) {
|
|
auto* input_d = x.data<T>();
|
|
auto* output_d = dev_ctx.template Alloc<T>(out);
|
|
float scale = static_cast<float>(alpha);
|
|
auto input_dims = x.dims();
|
|
// in shouble be (batch * seq * 3 * head_num * head_size)
|
|
// out shouble be (batch * seq * head_num * head_size)
|
|
int batch_size = static_cast<int>(input_dims[0]);
|
|
int seq_len = static_cast<int>(input_dims[1]);
|
|
int head_size = static_cast<int>(input_dims[4]);
|
|
|
|
DenseTensor temp1, temp2;
|
|
temp1.Resize(input_dims);
|
|
float* trans_input = dev_ctx.template Alloc<float>(&temp1);
|
|
temp2.Resize(input_dims);
|
|
float* trans_output = dev_ctx.template Alloc<float>(&temp2);
|
|
|
|
transpose_before_bmm1<T, float>(
|
|
input_d, trans_input, batch_size, seq_len, head_number, head_size);
|
|
float* query = trans_input;
|
|
float* key = trans_input + batch_size * head_number * seq_len * head_size;
|
|
float* value =
|
|
trans_input + batch_size * head_number * seq_len * head_size * 2;
|
|
|
|
scaled_dp_attention(query,
|
|
key,
|
|
value,
|
|
scale,
|
|
batch_size,
|
|
seq_len,
|
|
seq_len,
|
|
head_number,
|
|
head_size,
|
|
trans_output);
|
|
transpose_after_bmm2<float, T>(
|
|
trans_output, output_d, batch_size, seq_len, head_number, head_size);
|
|
}
|
|
|
|
} // namespace fusion
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(self_dp_attention,
|
|
CPU,
|
|
ALL_LAYOUT,
|
|
phi::fusion::SelfDPAttenKernel,
|
|
float,
|
|
double) {}
|