mirror of https://github.com/xianyi/OpenBLAS.git
441 lines
16 KiB
C
441 lines
16 KiB
C
/***************************************************************************
|
|
* Copyright (c) 2025, The OpenBLAS Project
|
|
* All rights reserved.
|
|
* Redistribution and use in source and binary forms, with or without
|
|
* modification, are permitted provided that the following conditions are
|
|
* met:
|
|
* 1. Redistributions of source code must retain the above copyright
|
|
* notice, this list of conditions and the following disclaimer.
|
|
* 2. Redistributions in binary form must reproduce the above copyright
|
|
* notice, this list of conditions and the following disclaimer in
|
|
* the documentation and/or other materials provided with the
|
|
* distribution.
|
|
* 3. Neither the name of the OpenBLAS project nor the names of
|
|
* its contributors may be used to endorse or promote products
|
|
* derived from this software without specific prior written permission.
|
|
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
|
|
* ARE DISCLAIMED. IN NO EVENT SHALL THE OPENBLAS PROJECT OR CONTRIBUTORS BE
|
|
* LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
|
|
* CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
|
|
* SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
|
* INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
|
|
* CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
|
|
* ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
|
* POSSIBILITY OF SUCH DAMAGE.
|
|
* *****************************************************************************/
|
|
|
|
#include <arm_sve.h>
|
|
#include <arm_neon.h>
|
|
|
|
#include "common.h"
|
|
|
|
#ifdef BGEMM
|
|
|
|
#ifdef ALPHA_ONE
|
|
|
|
#define TO16 vcvth_bf16_f32
|
|
#define TO32 vcvtah_f32_bf16
|
|
|
|
#define UPDATE_C(PG, PTR, DST, SRC) \
|
|
do { \
|
|
svtmp16 = svld1_bf16((pghalf), (PTR)); \
|
|
DST = svreinterpret_f32(svzip1_bf16(zeros, svtmp16)); \
|
|
DST = svadd_z((PG), SRC, DST); \
|
|
svtmp16 = svcvt_bf16_f32_z((PG), DST); \
|
|
svtmp16 = svuzp1_bf16(svtmp16, svtmp16); \
|
|
svst1_bf16((pghalf), (PTR), svtmp16); \
|
|
} while (0);
|
|
#define UPDATE_C2(ptr, tmp, vector) \
|
|
*(ptr) = TO16(vector[0] + TO32(*ptr)); \
|
|
*(ptr + 1) = TO16(vector[1] + TO32(*(ptr + 1)));
|
|
#define UPDATE_C1(ptr, value) *ptr = TO16(TO32(*ptr) + (value))
|
|
|
|
#else
|
|
|
|
#define UPDATE_C(PG, PTR, DST, SRC) \
|
|
do { \
|
|
svtmp16 = svld1_bf16((pghalf), (PTR)); \
|
|
DST = svreinterpret_f32(svzip1_bf16(zeros, svtmp16)); \
|
|
DST = svmad_z((PG), svalpha, SRC, DST); \
|
|
svtmp16 = svcvt_bf16_f32_z((PG), DST); \
|
|
svtmp16 = svuzp1_bf16(svtmp16, svtmp16); \
|
|
svst1_bf16((pghalf), (PTR), svtmp16); \
|
|
} while (0);
|
|
#define UPDATE_C2(ptr, tmp, vector) \
|
|
*(ptr) = TO16(vector[0] * alpha + TO32(*ptr)); \
|
|
*(ptr + 1) = TO16(vector[1] * alpha + TO32(*(ptr + 1)));
|
|
#define UPDATE_C1(ptr, value) *ptr = TO16(TO32(*ptr) + (value) * alpha)
|
|
|
|
#endif
|
|
|
|
#else
|
|
|
|
#ifdef ALPHA_ONE
|
|
|
|
#define UPDATE_C(PG, PTR, DST, SRC) \
|
|
do { \
|
|
DST = svld1_f32((PG), (PTR)); \
|
|
DST = svadd_z((PG), SRC, DST); \
|
|
svst1_f32((PG), (PTR), DST); \
|
|
} while (0);
|
|
#define UPDATE_C2(ptr, tmp, vector) \
|
|
tmp = vld1_f32(ptr); \
|
|
tmp = vadd_f32(vector, tmp); \
|
|
vst1_f32(ptr, tmp);
|
|
#define UPDATE_C1(ptr, value) *ptr = *ptr + (value)
|
|
|
|
#else
|
|
|
|
#define UPDATE_C(PG, PTR, DST, SRC) \
|
|
do { \
|
|
DST = svld1_f32((PG), (PTR)); \
|
|
DST = svmad_z((PG), svalpha, SRC, DST); \
|
|
svst1_f32((PG), (PTR), DST); \
|
|
} while (0);
|
|
#define UPDATE_C2(ptr, tmp, vector) \
|
|
tmp = vld1_f32(ptr); \
|
|
tmp = vmla_n_f32(tmp, vector, alpha); \
|
|
vst1_f32(ptr, tmp);
|
|
#define UPDATE_C1(ptr, value) *ptr = *ptr + (value) * alpha
|
|
|
|
#endif
|
|
|
|
#endif
|
|
|
|
#ifdef BGEMM
|
|
#define OUTPUT_FLOAT bfloat16_t
|
|
#else
|
|
#define OUTPUT_FLOAT float
|
|
#endif
|
|
|
|
#ifdef ALPHA_ONE
|
|
static int bgemm_kernel_neoversev1_alpha_one(BLASLONG m, BLASLONG n, BLASLONG k,
|
|
float alpha, IFLOAT *AA, IFLOAT *BB,
|
|
FLOAT *CC, BLASLONG ldc)
|
|
#else
|
|
static int bgemm_kernel_neoversev1_alpha(BLASLONG m, BLASLONG n, BLASLONG k,
|
|
float alpha, IFLOAT *AA, IFLOAT *BB, FLOAT *CC,
|
|
BLASLONG ldc)
|
|
#endif
|
|
{
|
|
BLASLONG pad_k = (k + 3) & ~3;
|
|
|
|
#ifndef ALPHA_ONE
|
|
svfloat32_t svalpha = svdup_f32(alpha);
|
|
#endif
|
|
|
|
bfloat16_t *ptr_a = (bfloat16_t *)AA;
|
|
bfloat16_t *ptr_b = (bfloat16_t *)BB;
|
|
OUTPUT_FLOAT *ptr_c =(OUTPUT_FLOAT*)CC;
|
|
|
|
bfloat16_t *ptr_a0;
|
|
bfloat16_t *ptr_b0;
|
|
OUTPUT_FLOAT *ptr_c0, *ptr_c1, *ptr_c2, *ptr_c3;
|
|
svfloat32_t tmp0, tmp1, tmp2, tmp3;
|
|
#ifdef BGEMM
|
|
svbfloat16_t zeros = svdup_n_bf16(TO16(0.0));
|
|
svbfloat16_t svtmp16;
|
|
#else
|
|
float32x2_t tmp4, tmp5, tmp6, tmp7;
|
|
#endif
|
|
|
|
const int sve_size_bf16 = svcnth();
|
|
const int num_accumulators = sve_size_bf16 >> 1;
|
|
|
|
svbool_t pgtrue = svptrue_b16();
|
|
#ifdef BGEMM
|
|
// For BF16 load/store we use half the vector size
|
|
svbool_t pghalf = svwhilelt_b16(0, num_accumulators);
|
|
#endif
|
|
|
|
// N values are 4x2 packed matrices
|
|
int n_step = 0;
|
|
const int n2 = n & -2;
|
|
const int n4 = n & -4;
|
|
|
|
// For 256-bit this would be 8
|
|
const int m_acc = (m & -num_accumulators);
|
|
const int m2 = m & -2;
|
|
|
|
for (; n_step < n4; n_step += 4) {
|
|
ptr_a = (bfloat16_t *)AA;
|
|
ptr_c0 = ptr_c;
|
|
ptr_c1 = ptr_c0 + ldc;
|
|
ptr_c2 = ptr_c1 + ldc;
|
|
ptr_c3 = ptr_c2 + ldc;
|
|
ptr_c += 4 * ldc;
|
|
|
|
int m_step = 0;
|
|
for (; m_step < m_acc; m_step += num_accumulators) {
|
|
svfloat32_t acc0 = svdup_f32(0);
|
|
svfloat32_t acc1 = svdup_f32(0);
|
|
svfloat32_t acc2 = svdup_f32(0);
|
|
svfloat32_t acc3 = svdup_f32(0);
|
|
|
|
ptr_a0 = ptr_a;
|
|
ptr_b0 = ptr_b;
|
|
ptr_a += num_accumulators * pad_k;
|
|
|
|
// Load entire 2VL block
|
|
for (BLASLONG p = 0; p < pad_k; p += 4) {
|
|
svbfloat16_t ma0 = svld1_bf16(pgtrue, ptr_a0);
|
|
svbfloat16_t ma1 = svld1_bf16(pgtrue, ptr_a0 + sve_size_bf16);
|
|
svbfloat16_t mb0 = svld1rq_bf16(pgtrue, ptr_b0);
|
|
svbfloat16_t mb1 = svld1rq_bf16(pgtrue, ptr_b0 + 8);
|
|
|
|
acc0 = svbfmmla_f32(acc0, mb0, ma0);
|
|
acc1 = svbfmmla_f32(acc1, mb0, ma1);
|
|
acc2 = svbfmmla_f32(acc2, mb1, ma0);
|
|
acc3 = svbfmmla_f32(acc3, mb1, ma1);
|
|
|
|
ptr_a0 += sve_size_bf16 * 2;
|
|
ptr_b0 += 16;
|
|
}
|
|
|
|
svfloat32_t out0 = svreinterpret_f32_u64(svuzp1_u64(svreinterpret_u64_f32(acc0), svreinterpret_u64_f32(acc1)));
|
|
svfloat32_t out1 = svreinterpret_f32_u64(svuzp2_u64(svreinterpret_u64_f32(acc0), svreinterpret_u64_f32(acc1)));
|
|
|
|
svfloat32_t out2 = svreinterpret_f32_u64(svuzp1_u64(svreinterpret_u64_f32(acc2), svreinterpret_u64_f32(acc3)));
|
|
svfloat32_t out3 = svreinterpret_f32_u64(svuzp2_u64(svreinterpret_u64_f32(acc2), svreinterpret_u64_f32(acc3)));
|
|
|
|
UPDATE_C(pgtrue, ptr_c0, tmp0, out0);
|
|
UPDATE_C(pgtrue, ptr_c1, tmp1, out1);
|
|
UPDATE_C(pgtrue, ptr_c2, tmp2, out2);
|
|
UPDATE_C(pgtrue, ptr_c3, tmp3, out3);
|
|
|
|
ptr_c0 += num_accumulators;
|
|
ptr_c1 += num_accumulators;
|
|
ptr_c2 += num_accumulators;
|
|
ptr_c3 += num_accumulators;
|
|
}
|
|
|
|
for (; m_step < m2; m_step += 2) {
|
|
float32x4_t acc0 = {0,0,0,0};
|
|
float32x4_t acc1 = {0,0,0,0};
|
|
|
|
ptr_a0 = ptr_a;
|
|
ptr_b0 = ptr_b;
|
|
ptr_a += 2 * pad_k;
|
|
|
|
for (BLASLONG p = 0; p < pad_k; p += 4) {
|
|
bfloat16x8_t ma0 = vld1q_bf16(ptr_a0);
|
|
bfloat16x8_t mb0 = vld1q_bf16(ptr_b0);
|
|
bfloat16x8_t mb1 = vld1q_bf16(ptr_b0 + 8);
|
|
|
|
acc0 = vbfmmlaq_f32(acc0, mb0, ma0);
|
|
acc1 = vbfmmlaq_f32(acc1, mb1, ma0);
|
|
|
|
ptr_a0 += 8;
|
|
ptr_b0 += 16;
|
|
}
|
|
|
|
UPDATE_C2(ptr_c0, tmp4, vget_low_f32(acc0));
|
|
UPDATE_C2(ptr_c1, tmp5, vget_high_f32(acc0));
|
|
UPDATE_C2(ptr_c2, tmp6, vget_low_f32(acc1));
|
|
UPDATE_C2(ptr_c3, tmp7, vget_high_f32(acc1));
|
|
|
|
ptr_c0 += 2;
|
|
ptr_c1 += 2;
|
|
ptr_c2 += 2;
|
|
ptr_c3 += 2;
|
|
}
|
|
|
|
// Final row is always a contiguous single row
|
|
if (m & 1) {
|
|
ptr_a0 = ptr_a;
|
|
ptr_b0 = ptr_b;
|
|
float32x4_t acc0 = {0,0,0,0};
|
|
float32x4_t acc1 = {0,0,0,0};
|
|
|
|
for (BLASLONG p = 0; p < pad_k; p += 4) {
|
|
/// Same A value can be used for both B values
|
|
bfloat16x8_t ma0 = vreinterpretq_bf16_u64(vdupq_n_u64(
|
|
*((uint64_t*)ptr_a0)
|
|
));
|
|
bfloat16x8_t mb0 = vld1q_bf16(ptr_b0);
|
|
bfloat16x8_t mb1 = vld1q_bf16(ptr_b0 + 8);
|
|
|
|
acc0 = vbfmmlaq_f32(acc0, mb0, ma0);
|
|
acc1 = vbfmmlaq_f32(acc1, mb1, ma0);
|
|
|
|
ptr_a0 += 4;
|
|
ptr_b0 += 16;
|
|
}
|
|
|
|
UPDATE_C1(ptr_c0, acc0[1]);
|
|
UPDATE_C1(ptr_c1, acc0[3]);
|
|
UPDATE_C1(ptr_c2, acc1[1]);
|
|
UPDATE_C1(ptr_c3, acc1[3]);
|
|
}
|
|
|
|
ptr_b += 4 * pad_k;
|
|
}
|
|
|
|
for (; n_step < n2; n_step += 2) {
|
|
ptr_a = (bfloat16_t *)AA;
|
|
ptr_c0 = ptr_c;
|
|
ptr_c1 = ptr_c0 + ldc;
|
|
ptr_c += 2 * ldc;
|
|
|
|
// Sets of two are contiguously packed so yay
|
|
int m_step = 0;
|
|
for (; m_step < m_acc; m_step += num_accumulators) {
|
|
svfloat32_t acc0 = svdup_f32(0);
|
|
svfloat32_t acc1 = svdup_f32(0);
|
|
ptr_a0 = ptr_a;
|
|
ptr_b0 = ptr_b;
|
|
ptr_a += num_accumulators * pad_k;
|
|
|
|
// Load entire 8x4 block
|
|
for (BLASLONG p = 0; p < pad_k; p += 4) {
|
|
svbfloat16_t ma0 = svld1_bf16(pgtrue, ptr_a0);
|
|
svbfloat16_t ma1 = svld1_bf16(pgtrue, ptr_a0 + sve_size_bf16);
|
|
svbfloat16_t mb0 = svld1rq_bf16(pgtrue, ptr_b0);
|
|
|
|
acc0 = svbfmmla_f32(acc0, mb0, ma0);
|
|
acc1 = svbfmmla_f32(acc1, mb0, ma1);
|
|
|
|
ptr_a0 += sve_size_bf16 * 2;
|
|
ptr_b0 += 8;
|
|
}
|
|
|
|
svfloat32_t out0 = svreinterpret_f32_u64(svuzp1_u64(svreinterpret_u64_f32(acc0), svreinterpret_u64_f32(acc1)));
|
|
svfloat32_t out1 = svreinterpret_f32_u64(svuzp2_u64(svreinterpret_u64_f32(acc0), svreinterpret_u64_f32(acc1)));
|
|
|
|
UPDATE_C(pgtrue, ptr_c0, tmp0, out0);
|
|
UPDATE_C(pgtrue, ptr_c1, tmp1, out1);
|
|
|
|
ptr_c0 += num_accumulators;
|
|
ptr_c1 += num_accumulators;
|
|
}
|
|
|
|
for (; m_step < m2; m_step += 2) {
|
|
float32x4_t acc = {0,0,0,0};
|
|
|
|
ptr_a0 = ptr_a;
|
|
ptr_b0 = ptr_b;
|
|
ptr_a += 2 * pad_k;
|
|
|
|
for (BLASLONG p = 0; p < pad_k; p += 4) {
|
|
bfloat16x8_t ma0 = vld1q_bf16(ptr_a0);
|
|
bfloat16x8_t mb0 = vld1q_bf16(ptr_b0);
|
|
|
|
acc = vbfmmlaq_f32(acc, mb0, ma0);
|
|
|
|
ptr_a0 += 8;
|
|
ptr_b0 += 8;
|
|
}
|
|
|
|
UPDATE_C2(ptr_c0, tmp4, vget_low_f32(acc));
|
|
UPDATE_C2(ptr_c1, tmp5, vget_high_f32(acc));
|
|
|
|
ptr_c0 += 2;
|
|
ptr_c1 += 2;
|
|
}
|
|
|
|
// Final row is always a contiguous single row
|
|
if (m & 1) {
|
|
ptr_a0 = ptr_a;
|
|
ptr_b0 = ptr_b;
|
|
float32x4_t acc = {0,0,0,0};
|
|
|
|
for (BLASLONG p = 0; p < pad_k; p += 4) {
|
|
/// Same A value can be used for both B values
|
|
bfloat16x8_t ma0 = vreinterpretq_bf16_u64(vdupq_n_u64(
|
|
*((uint64_t*)ptr_a0)
|
|
));
|
|
bfloat16x8_t mb0 = vld1q_bf16(ptr_b0);
|
|
|
|
acc = vbfmmlaq_f32(acc, mb0, ma0);
|
|
|
|
ptr_a0 += 4;
|
|
ptr_b0 += 8;
|
|
}
|
|
|
|
UPDATE_C1(ptr_c0, acc[0]);
|
|
UPDATE_C1(ptr_c1, acc[2]);
|
|
}
|
|
|
|
ptr_b += 2 * pad_k;
|
|
}
|
|
|
|
if (n & 1) {
|
|
ptr_a = (bfloat16_t *)AA;
|
|
ptr_c0 = ptr_c;
|
|
|
|
int m_step = 0;
|
|
for (; m_step < m_acc; m_step += num_accumulators) {
|
|
ptr_a0 = ptr_a;
|
|
ptr_b0 = ptr_b;
|
|
ptr_a += num_accumulators * pad_k;
|
|
|
|
svfloat32_t acc0 = svdup_f32(0);
|
|
svfloat32_t acc1 = svdup_f32(0);
|
|
|
|
// Load entire 8x4 block
|
|
for (BLASLONG p = 0; p < pad_k; p += 4) {
|
|
uint64_t* ptr_b0_u64 = (uint64_t*)ptr_b0;
|
|
|
|
svbfloat16_t ma0 = svld1_bf16(pgtrue, ptr_a0);
|
|
svbfloat16_t ma1 = svld1_bf16(pgtrue, ptr_a0 + sve_size_bf16);
|
|
svbfloat16_t mb0 = svreinterpret_bf16_u64(svdup_u64(*ptr_b0_u64));
|
|
|
|
acc0 = svbfmmla_f32(acc0, mb0, ma0);
|
|
acc1 = svbfmmla_f32(acc1, mb0, ma1);
|
|
|
|
ptr_a0 += sve_size_bf16 * 2;
|
|
ptr_b0 += 4;
|
|
}
|
|
|
|
svfloat32_t out0 = svreinterpret_f32_u64(svuzp1_u64(svreinterpret_u64_f32(acc0), svreinterpret_u64_f32(acc1)));
|
|
UPDATE_C(pgtrue, ptr_c0, tmp0, out0);
|
|
|
|
ptr_c0 += num_accumulators;
|
|
}
|
|
|
|
for (; m_step < m2; m_step += 2) {
|
|
float32x4_t acc = {0, 0, 0, 0};
|
|
|
|
ptr_a0 = ptr_a;
|
|
ptr_b0 = ptr_b;
|
|
ptr_a += 2 * pad_k;
|
|
|
|
for (BLASLONG p = 0; p < pad_k; p += 4) {
|
|
bfloat16x8_t ma0 = vld1q_bf16(ptr_a0);
|
|
bfloat16x8_t mb0 = vcombine_bf16(vld1_bf16(ptr_b0), vdup_n_bf16(vcvth_bf16_f32(0.0f)));
|
|
|
|
acc = vbfmmlaq_f32(acc, mb0, ma0);
|
|
|
|
ptr_a0 += 8;
|
|
ptr_b0 += 4;
|
|
}
|
|
|
|
UPDATE_C2(ptr_c0, tmp4, vget_low_f32(acc));
|
|
|
|
ptr_c0 += 2;
|
|
}
|
|
|
|
if (m & 1) {
|
|
ptr_a0 = ptr_a;
|
|
ptr_b0 = ptr_b;
|
|
float32x2_t acc = {0,0};
|
|
|
|
for (BLASLONG p = 0; p < pad_k; p += 4) {
|
|
bfloat16x4_t ma0 = vld1_bf16(ptr_a0);
|
|
bfloat16x4_t mb0 = vld1_bf16(ptr_b0);
|
|
|
|
acc = vbfdot_f32(acc, ma0, mb0);
|
|
|
|
ptr_a0 += 4;
|
|
ptr_b0 += 4;
|
|
}
|
|
|
|
UPDATE_C1(ptr_c0, acc[0] + acc[1]);
|
|
}
|
|
}
|
|
|
|
return 0;
|
|
}
|