This changes the kernels to pack full SVE vectors and reduces the overall complexity of the inner GEMM loop.pull/5419/head
@@ -34,15 +34,15 @@ SGEMVTKERNEL = gemv_t_sve_v1x3.c | |||
DGEMVTKERNEL = gemv_t_sve_v1x3.c | |||
ifeq ($(BUILD_BFLOAT16), 1) | |||
BGEMM_BETA = bgemm_beta_neon.c | |||
BGEMMKERNEL = bgemm_kernel_$(BGEMM_UNROLL_M)x$(BGEMM_UNROLL_N)_neoversev1.c | |||
BGEMMKERNEL = bgemm_kernel_2vlx4_neoversev1.c | |||
ifneq ($(BGEMM_UNROLL_M), $(BGEMM_UNROLL_N)) | |||
BGEMMINCOPY = sbgemm_ncopy_$(SBGEMM_UNROLL_M)_neoversev1.c | |||
BGEMMITCOPY = sbgemm_tcopy_$(SBGEMM_UNROLL_M)_neoversev1.c | |||
BGEMMINCOPY = bgemm_ncopy_2vl_neoversev1.c | |||
BGEMMITCOPY = bgemm_tcopy_2vl_neoversev1.c | |||
BGEMMINCOPYOBJ = bgemm_incopy$(TSUFFIX).$(SUFFIX) | |||
BGEMMITCOPYOBJ = bgemm_itcopy$(TSUFFIX).$(SUFFIX) | |||
endif | |||
BGEMMONCOPY = sbgemm_ncopy_$(BGEMM_UNROLL_N)_neoversev1.c | |||
BGEMMOTCOPY = sbgemm_tcopy_$(BGEMM_UNROLL_N)_neoversev1.c | |||
BGEMMONCOPY = bgemm_ncopy_4_neoversev1.c | |||
BGEMMOTCOPY = bgemm_tcopy_4_neoversev1.c | |||
BGEMMONCOPYOBJ = bgemm_oncopy$(TSUFFIX).$(SUFFIX) | |||
BGEMMOTCOPYOBJ = bgemm_otcopy$(TSUFFIX).$(SUFFIX) | |||
@@ -50,15 +50,15 @@ BGEMVTKERNEL = sbgemv_t_bfdot.c | |||
BGEMVNKERNEL = bgemv_n_sve_v3x4.c | |||
SBGEMM_BETA = sbgemm_beta_neoversev1.c | |||
SBGEMMKERNEL = sbgemm_kernel_$(SBGEMM_UNROLL_M)x$(SBGEMM_UNROLL_N)_neoversev1.c | |||
SBGEMMKERNEL = bgemm_kernel_2vlx4_neoversev1.c | |||
ifneq ($(SBGEMM_UNROLL_M), $(SBGEMM_UNROLL_N)) | |||
SBGEMMINCOPY = sbgemm_ncopy_$(SBGEMM_UNROLL_M)_neoversev1.c | |||
SBGEMMITCOPY = sbgemm_tcopy_$(SBGEMM_UNROLL_M)_neoversev1.c | |||
SBGEMMINCOPY = bgemm_ncopy_2vl_neoversev1.c | |||
SBGEMMITCOPY = bgemm_tcopy_2vl_neoversev1.c | |||
SBGEMMINCOPYOBJ = sbgemm_incopy$(TSUFFIX).$(SUFFIX) | |||
SBGEMMITCOPYOBJ = sbgemm_itcopy$(TSUFFIX).$(SUFFIX) | |||
endif | |||
SBGEMMONCOPY = sbgemm_ncopy_$(SBGEMM_UNROLL_N)_neoversev1.c | |||
SBGEMMOTCOPY = sbgemm_tcopy_$(SBGEMM_UNROLL_N)_neoversev1.c | |||
SBGEMMONCOPY = bgemm_ncopy_4_neoversev1.c | |||
SBGEMMOTCOPY = bgemm_tcopy_4_neoversev1.c | |||
SBGEMMONCOPYOBJ = sbgemm_oncopy$(TSUFFIX).$(SUFFIX) | |||
SBGEMMOTCOPYOBJ = sbgemm_otcopy$(TSUFFIX).$(SUFFIX) | |||
@@ -0,0 +1,57 @@ | |||
/*************************************************************************** | |||
* 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" | |||
#define ALPHA_ONE | |||
#include "bgemm_kernel_2vlx4_neoversev1_impl.c" | |||
#undef ALPHA_ONE | |||
#undef UPDATE_C | |||
#undef UPDATE_C2 | |||
#undef UPDATE_C1 | |||
#include "bgemm_kernel_2vlx4_neoversev1_impl.c" | |||
int CNAME(BLASLONG m, BLASLONG n, BLASLONG k, FLOAT alpha, IFLOAT *A, IFLOAT *B, | |||
FLOAT *C, BLASLONG ldc) { | |||
#ifdef BGEMM | |||
bfloat16_t alpha_bf16; | |||
memcpy(&alpha_bf16, &alpha, sizeof(bfloat16_t)); | |||
float alpha_f32 = vcvtah_f32_bf16(alpha_bf16); | |||
#else | |||
float alpha_f32 = alpha; | |||
#endif | |||
if (alpha_f32 == 1.0f) | |||
return bgemm_kernel_neoversev1_alpha_one(m, n, k, alpha_f32, A, B, C, ldc); | |||
else | |||
return bgemm_kernel_neoversev1_alpha(m, n, k, alpha_f32, A, B, C, ldc); | |||
return 0; | |||
} |
@@ -0,0 +1,437 @@ | |||
/*************************************************************************** | |||
* 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 { \ | |||
DST = svreinterpret_f32_u32(svld1uh_u32((pghalf), (uint16_t*)PTR)); \ | |||
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 { \ | |||
DST = svreinterpret_f32_u32(svld1uh_u32((pghalf), (uint16_t*)PTR)); \ | |||
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 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; | |||
} |
@@ -0,0 +1,135 @@ | |||
/*************************************************************************** | |||
* 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" | |||
int CNAME(BLASLONG n, BLASLONG m, IFLOAT *input, BLASLONG lda, IFLOAT *output) { | |||
const int sve_size_bf16 = svcnth(); | |||
const int num_accumulators = sve_size_bf16 >> 1; | |||
const int m_sve_accumulators = m & -num_accumulators; | |||
const int n4 = n & -4; | |||
const int n_rest = n - n4; | |||
const int m2 = m & -2; | |||
const int m_rest = m - m2; | |||
size_t m_step = 0; | |||
for (; m_step < m_sve_accumulators; m_step += num_accumulators) { | |||
const uint16_t* inner_input = input; | |||
// Potential for vld1q here with transpose | |||
for (int n_step = 0; n_step < n4; n_step += 4) { | |||
for (int line = 0; line < num_accumulators; line += 4) { | |||
uint16x4_t a_vec0 = vld1_u16(inner_input + line * lda); | |||
uint16x4_t a_vec1 = vld1_u16(inner_input + (line + 1) * lda); | |||
uint16x4_t a_vec2 = vld1_u16(inner_input + (line + 2) * lda); | |||
uint16x4_t a_vec3 = vld1_u16(inner_input + (line + 3) * lda); | |||
vst1_u16(output, a_vec0); | |||
vst1_u16(output + 4, a_vec1); | |||
vst1_u16(output + 8, a_vec2); | |||
vst1_u16(output + 12, a_vec3); | |||
output += 16; | |||
} | |||
inner_input += 4; | |||
} | |||
// Bit of padding up to 4 for any remaining K | |||
// by the time we get here we hope the memory bandwidth is saturated | |||
if (n_rest) { | |||
for (BLASLONG line = 0; line < num_accumulators; line++) { | |||
output[0] = inner_input[0]; | |||
output[1] = n_rest == 1 ? 0 : inner_input[1]; | |||
output[2] = n_rest <= 2 ? 0 : inner_input[2]; | |||
output[3] = n_rest <= 3 ? 0 : inner_input[3]; | |||
inner_input += lda; | |||
output += 4; | |||
} | |||
} | |||
input += lda * num_accumulators; | |||
} | |||
// Any remaining blocks are done 2 at a time for ASIMD processing | |||
for (; m_step < m2; m_step += 2) { | |||
const uint16_t* inner_input = input; | |||
for (size_t n_step = 0; n_step < n4; n_step += 4) { | |||
uint16x4_t a_vec0 = vld1_u16(inner_input); | |||
uint16x4_t a_vec1 = vld1_u16(inner_input + lda); | |||
vst1_u16(output, a_vec0); | |||
vst1_u16(output + 4, a_vec1); | |||
inner_input += 4; | |||
output += 8; | |||
} | |||
if (n_rest) { | |||
for (BLASLONG line = 0; line < 2; line++) { | |||
output[0] = inner_input[0]; | |||
output[1] = n_rest == 1 ? 0 : inner_input[1]; | |||
output[2] = n_rest <= 2 ? 0 : inner_input[2]; | |||
output[3] = n_rest <= 3 ? 0 : inner_input[3]; | |||
inner_input += lda; | |||
output += 4; | |||
} | |||
} | |||
input += lda * 2; | |||
} | |||
// Final row is just there | |||
if (m_rest & 1) { | |||
for (size_t n_step = 0; n_step < n4; n_step += 4) { | |||
uint16x4_t a_vec0 = vld1_u16(input); | |||
vst1_u16(output, a_vec0); | |||
input += 4; | |||
output += 4; | |||
} | |||
if (n_rest) { | |||
output[0] = input[0]; | |||
output[1] = n_rest == 1 ? 0 : input[1]; | |||
output[2] = n_rest <= 2 ? 0 : input[2]; | |||
output[3] = n_rest <= 3 ? 0 : input[3]; | |||
} | |||
} | |||
return 0; | |||
} |
@@ -0,0 +1,124 @@ | |||
/*************************************************************************** | |||
* 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" | |||
int CNAME(BLASLONG n, BLASLONG m, IFLOAT *input, BLASLONG lda, IFLOAT *output) { | |||
const int num_accumulators = 4; | |||
const int m_accumulators = m & -4; | |||
const int n4 = n & -4; | |||
const int n_rest = n - n4; | |||
const int m_rest = m - m_accumulators; | |||
for (size_t m_step = 0; m_step < m_accumulators; m_step += num_accumulators) { | |||
const uint16_t* inner_input = input; | |||
// Potential for vld1q here with transpose | |||
for (size_t n_step = 0; n_step < n4; n_step += 4) { | |||
uint16x4_t a_vec0 = vld1_u16(inner_input + 0 * lda); | |||
uint16x4_t a_vec1 = vld1_u16(inner_input + 1 * lda); | |||
uint16x4_t a_vec2 = vld1_u16(inner_input + 2 * lda); | |||
uint16x4_t a_vec3 = vld1_u16(inner_input + 3 * lda); | |||
vst1_u16(output, a_vec0); | |||
vst1_u16(output + 4, a_vec1); | |||
vst1_u16(output + 8, a_vec2); | |||
vst1_u16(output + 12, a_vec3); | |||
output += 16; | |||
inner_input += 4; | |||
} | |||
if (n_rest) { | |||
for (BLASLONG line = 0; line < num_accumulators; line++) { | |||
output[0] = inner_input[0]; | |||
output[1] = n_rest == 1 ? 0 : inner_input[1]; | |||
output[2] = n_rest <= 2 ? 0 : inner_input[2]; | |||
output[3] = n_rest <= 3 ? 0 : inner_input[3]; | |||
inner_input += lda; | |||
output += 4; | |||
} | |||
} | |||
input += lda * num_accumulators; | |||
} | |||
if (m_rest & 2) { | |||
const uint16_t* inner_input = input; | |||
for (size_t n_step = 0; n_step < n4; n_step += 4) { | |||
uint16x4_t a_vec0 = vld1_u16(inner_input); | |||
uint16x4_t a_vec1 = vld1_u16(inner_input + lda); | |||
vst1_u16(output, a_vec0); | |||
vst1_u16(output + 4, a_vec1); | |||
inner_input += 4; | |||
output += 8; | |||
} | |||
if (n_rest) { | |||
for (BLASLONG line = 0; line < 2; line++) { | |||
output[0] = inner_input[0]; | |||
output[1] = n_rest == 1 ? 0 : inner_input[1]; | |||
output[2] = n_rest <= 2 ? 0 : inner_input[2]; | |||
output[3] = n_rest <= 3 ? 0 : inner_input[3]; | |||
inner_input += lda; | |||
output += 4; | |||
} | |||
} | |||
input += lda * 2; | |||
} | |||
if (m_rest & 1) { | |||
for (size_t n_step = 0; n_step < n4; n_step += 4) { | |||
uint16x4_t a_vec0 = vld1_u16(input); | |||
vst1_u16(output, a_vec0); | |||
input += 4; | |||
output += 4; | |||
} | |||
if (n_rest) { | |||
output[0] = input[0]; | |||
output[1] = n_rest == 1 ? 0 : input[1]; | |||
output[2] = n_rest <= 2 ? 0 : input[2]; | |||
output[3] = n_rest <= 3 ? 0 : input[3]; | |||
} | |||
} | |||
return 0; | |||
} |
@@ -0,0 +1,154 @@ | |||
/*************************************************************************** | |||
* 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" | |||
int CNAME(BLASLONG m, BLASLONG n, IFLOAT *input, BLASLONG lda, IFLOAT *output) { | |||
const int sve_size_bf16 = svcnth(); | |||
const int num_accumulators_sve = sve_size_bf16 >> 1; | |||
const int num_accumulators = num_accumulators_sve; | |||
const int incr_accumulators = 4; | |||
const int n_sve_accumulators = (n & -num_accumulators); | |||
const int n2 = n & -2; | |||
const int n_rest = n - n2; | |||
const int m4 = m & -4; | |||
const int m_rest = m - m4; | |||
size_t n_step = 0; | |||
for (; n_step < n_sve_accumulators; n_step += num_accumulators) { | |||
const uint16_t* inner_input = input; | |||
// Full 4x4 item transposes down the M dimension | |||
for (size_t m_step = 0; m_step < m4; m_step += 4) { | |||
const uint16_t* tile = inner_input; | |||
for (size_t line = 0; line < num_accumulators; line += incr_accumulators) { | |||
// Load 4x4 block | |||
uint16x4_t a_vec0 = vld1_u16(tile); | |||
uint16x4_t a_vec1 = vld1_u16(tile + lda); | |||
uint16x4_t a_vec2 = vld1_u16(tile + 2 * lda); | |||
uint16x4_t a_vec3 = vld1_u16(tile + 3 * lda); | |||
// Transpose 4x4 blocks | |||
uint16x4_t out_vec0 = vzip1_u16(a_vec0, a_vec1); | |||
uint16x4_t out_vec1 = vzip2_u16(a_vec0, a_vec1); | |||
uint16x4_t out_vec2 = vzip1_u16(a_vec2, a_vec3); | |||
uint16x4_t out_vec3 = vzip2_u16(a_vec2, a_vec3); | |||
// Transpose 8x4 blocks | |||
a_vec0 = vreinterpret_u16_u32(vzip1_u32(vreinterpret_u32_u16(out_vec0), vreinterpret_u32_u16(out_vec2))); | |||
a_vec1 = vreinterpret_u16_u32(vzip2_u32(vreinterpret_u32_u16(out_vec0), vreinterpret_u32_u16(out_vec2))); | |||
a_vec2 = vreinterpret_u16_u32(vzip1_u32(vreinterpret_u32_u16(out_vec1), vreinterpret_u32_u16(out_vec3))); | |||
a_vec3 = vreinterpret_u16_u32(vzip2_u32(vreinterpret_u32_u16(out_vec1), vreinterpret_u32_u16(out_vec3))); | |||
vst1_u16(output, a_vec0); | |||
vst1_u16(output + 4, a_vec1); | |||
vst1_u16(output + 8, a_vec2); | |||
vst1_u16(output + 12, a_vec3); | |||
tile += incr_accumulators; | |||
output += 16; | |||
} | |||
inner_input += incr_accumulators * lda; | |||
} | |||
if (m_rest) { | |||
for (BLASLONG line = 0; line < num_accumulators; line++) { | |||
output[0] = inner_input[0]; | |||
output[1] = m_rest == 1 ? 0 : *(inner_input + lda); | |||
output[2] = m_rest <= 2 ? 0 : *(inner_input + 2 * lda); | |||
output[3] = m_rest <= 3 ? 0 : *(inner_input + 3 * lda); | |||
inner_input++; | |||
output += 4; | |||
} | |||
} | |||
input += num_accumulators; | |||
} | |||
for (; n_step < n2; n_step += 2) { | |||
const uint16_t* inner_input = input; | |||
for (size_t m_step = 0; m_step < m4; m_step += 4) { | |||
for (BLASLONG line = 0; line < 2; line++) { | |||
output[0] = *(inner_input + line); | |||
output[1] = *(inner_input + line + lda); | |||
output[2] = *(inner_input + line + 2 * lda); | |||
output[3] = *(inner_input + line + 3 * lda); | |||
output += 4; | |||
} | |||
inner_input += 4 * lda; | |||
} | |||
if (m_rest) { | |||
for (BLASLONG line = 0; line < 2; line++) { | |||
output[0] = *(inner_input + line); | |||
output[1] = m_rest == 1 ? 0 : *(inner_input + line + lda); | |||
output[2] = m_rest <= 2 ? 0 : *(inner_input + line + 2 * lda); | |||
output[3] = m_rest <= 3 ? 0 : *(inner_input + line + 3 * lda); | |||
output += 4; | |||
} | |||
} | |||
input += 2; | |||
} | |||
if (n_rest & 1) { | |||
const uint16_t* inner_input = input; | |||
for (size_t m_step = 0; m_step < m4; m_step += 4) { | |||
output[0] = *inner_input; | |||
output[1] = *(inner_input + lda); | |||
output[2] = *(inner_input + 2 * lda); | |||
output[3] = *(inner_input + 3 * lda); | |||
inner_input += 4 * lda; | |||
output += 4; | |||
} | |||
if (m_rest) { | |||
output[0] = inner_input[0]; | |||
output[1] = m_rest == 1 ? 0 : *(inner_input + lda); | |||
output[2] = m_rest <= 2 ? 0 : *(inner_input + 2 * lda); | |||
output[3] = m_rest <= 3 ? 0 : *(inner_input + 3 * lda); | |||
output += 4; | |||
} | |||
} | |||
return 0; | |||
} |
@@ -0,0 +1,143 @@ | |||
/*************************************************************************** | |||
* 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" | |||
int CNAME(BLASLONG m, BLASLONG n, IFLOAT *input, BLASLONG lda, IFLOAT *output) { | |||
const int num_accumulators = 4; | |||
const int n_accumulators = (n & -num_accumulators); | |||
const int n_rest = n - n_accumulators; | |||
const int m4 = m & -4; | |||
const int m_rest = m - m4; | |||
for (size_t n_step = 0; n_step < n_accumulators; n_step += num_accumulators) { | |||
const uint16_t* inner_input = input; | |||
// Full 4x4 item transposes down the M dimension | |||
for (size_t m_step = 0; m_step < m4; m_step += 4) { | |||
// Load 4x4 block | |||
uint16x4_t a_vec0 = vld1_u16(inner_input); | |||
uint16x4_t a_vec1 = vld1_u16(inner_input + lda); | |||
uint16x4_t a_vec2 = vld1_u16(inner_input + 2 * lda); | |||
uint16x4_t a_vec3 = vld1_u16(inner_input + 3 * lda); | |||
// Transpose 4x4 blocks | |||
uint16x4_t out_vec0 = vzip1_u16(a_vec0, a_vec1); | |||
uint16x4_t out_vec1 = vzip2_u16(a_vec0, a_vec1); | |||
uint16x4_t out_vec2 = vzip1_u16(a_vec2, a_vec3); | |||
uint16x4_t out_vec3 = vzip2_u16(a_vec2, a_vec3); | |||
// Transpose 8x4 blocks | |||
a_vec0 = vreinterpret_u16_u32(vzip1_u32(vreinterpret_u32_u16(out_vec0), vreinterpret_u32_u16(out_vec2))); | |||
a_vec1 = vreinterpret_u16_u32(vzip2_u32(vreinterpret_u32_u16(out_vec0), vreinterpret_u32_u16(out_vec2))); | |||
a_vec2 = vreinterpret_u16_u32(vzip1_u32(vreinterpret_u32_u16(out_vec1), vreinterpret_u32_u16(out_vec3))); | |||
a_vec3 = vreinterpret_u16_u32(vzip2_u32(vreinterpret_u32_u16(out_vec1), vreinterpret_u32_u16(out_vec3))); | |||
vst1_u16(output, a_vec0); | |||
vst1_u16(output + 4, a_vec1); | |||
vst1_u16(output + 8, a_vec2); | |||
vst1_u16(output + 12, a_vec3); | |||
inner_input += 4 * lda; | |||
output += 16; | |||
} | |||
if (m_rest) { | |||
for (BLASLONG line = 0; line < num_accumulators; line++) { | |||
output[0] = inner_input[0]; | |||
output[1] = m_rest == 1 ? 0 : *(inner_input + lda); | |||
output[2] = m_rest <= 2 ? 0 : *(inner_input + 2 * lda); | |||
output[3] = m_rest <= 3 ? 0 : *(inner_input + 3 * lda); | |||
inner_input++; | |||
output += 4; | |||
} | |||
} | |||
input += num_accumulators; | |||
} | |||
// Extract two remaining rows as 128-bit vector paired | |||
if (n_rest & 2) { | |||
const uint16_t* inner_input = input; | |||
for (size_t m_step = 0; m_step < m4; m_step += 4) { | |||
for (BLASLONG line = 0; line < 2; line++) { | |||
output[0] = *(inner_input + line); | |||
output[1] = *(inner_input + line + lda); | |||
output[2] = *(inner_input + line + 2 * lda); | |||
output[3] = *(inner_input + line + 3 * lda); | |||
output += 4; | |||
} | |||
inner_input += 4 * lda; | |||
} | |||
if (m_rest) { | |||
for (BLASLONG line = 0; line < 2; line++) { | |||
output[0] = *(inner_input + line); | |||
output[1] = m_rest == 1 ? 0 : *(inner_input + line + lda); | |||
output[2] = m_rest <= 2 ? 0 : *(inner_input + line + 2 * lda); | |||
output[3] = m_rest <= 3 ? 0 : *(inner_input + line + 3 * lda); | |||
output += 4; | |||
} | |||
} | |||
input += 2; | |||
} | |||
// Flatten final row | |||
if (n_rest & 1) { | |||
const uint16_t* inner_input = input; | |||
for (size_t m_step = 0; m_step < m4; m_step += 4) { | |||
output[0] = *inner_input; | |||
output[1] = *(inner_input + lda); | |||
output[2] = *(inner_input + 2 * lda); | |||
output[3] = *(inner_input + 3 * lda); | |||
inner_input += 4 * lda; | |||
output += 4; | |||
} | |||
if (m_rest) { | |||
output[0] = inner_input[0]; | |||
output[1] = m_rest == 1 ? 0 : *(inner_input + lda); | |||
output[2] = m_rest <= 2 ? 0 : *(inner_input + 2 * lda); | |||
output[3] = m_rest <= 3 ? 0 : *(inner_input + 3 * lda); | |||
output += 4; | |||
} | |||
} | |||
return 0; | |||
} |
@@ -3598,15 +3598,15 @@ is a big desktop or server with abundant cache rather than a phone or embedded d | |||
#undef BGEMM_ALIGN_K | |||
#undef BGEMM_DEFAULT_UNROLL_M | |||
#undef BGEMM_DEFAULT_UNROLL_N | |||
#define BGEMM_ALIGN_K 8 | |||
#define BGEMM_ALIGN_K 4 | |||
#define BGEMM_DEFAULT_UNROLL_M 8 | |||
#define BGEMM_DEFAULT_UNROLL_N 4 | |||
#define BGEMM_DEFAULT_UNROLL_M 4 | |||
#undef SBGEMM_ALIGN_K | |||
#undef SBGEMM_DEFAULT_UNROLL_M | |||
#undef SBGEMM_DEFAULT_UNROLL_N | |||
#define SBGEMM_ALIGN_K 8 | |||
#define SBGEMM_DEFAULT_UNROLL_M 4 | |||
#define SBGEMM_ALIGN_K 4 | |||
#define SBGEMM_DEFAULT_UNROLL_M 8 | |||
#define SBGEMM_DEFAULT_UNROLL_N 4 | |||
#define SGEMM_DEFAULT_UNROLL_M 16 | |||