1// Copyright 2019 Google LLC 2// 3// This source code is licensed under the BSD-style license found in the 4// LICENSE file in the root directory of this source tree. 5 6$assert MR % 8 == 0 7$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" 8#include <assert.h> 9 10#include <arm_neon.h> 11 12#include <xnnpack/spmm.h> 13 14 15void xnn_f16_spmm_minmax_ukernel_${MR}x${NR}__neonfp16arith${"_x%d" % UNROLL if UNROLL > 1 else ""}( 16 size_t mc, 17 size_t nc, 18 const void*restrict input, 19 const void*restrict weights, 20 const int32_t*restrict widx_dmap, 21 const uint32_t*restrict nidx_nnzmap, 22 void*restrict output, 23 size_t output_stride, 24 const union xnn_f16_minmax_params params[restrict XNN_MIN_ELEMENTS(1)]) 25{ 26 assert(mc != 0); 27 assert(mc % sizeof(__fp16) == 0); 28 assert(nc != 0); 29 30 const __fp16*restrict i = (const __fp16*) input; 31 __fp16*restrict o = (__fp16*) output; 32 33 const float16x8_t vmax = vreinterpretq_f16_u16(vld1q_dup_u16(¶ms->neon.max)); 34 const float16x8_t vmin = vreinterpretq_f16_u16(vld1q_dup_u16(¶ms->neon.min)); 35 36 size_t output_decrement = output_stride * nc - ${MR} * sizeof(__fp16); 37 while XNN_LIKELY(mc >= ${MR} * sizeof(__fp16)) { 38 const __fp16*restrict w = (const __fp16*) weights; 39 const int32_t* dmap = widx_dmap; 40 const uint32_t* nnzmap = nidx_nnzmap; 41 size_t n = nc; 42 do { 43 uint32_t nnz = *nnzmap++; 44 $if UNROLL > 1: 45 float16x8_t vacc01234567x0 = vld1q_dup_f16(w); w += 1; 46 $for K in range(1, UNROLL): 47 float16x8_t vacc01234567x${K} = vmovq_n_f16(0.0f); 48 $for M in range(8, MR, 8): 49 float16x8_t vacc${ABC[M:M+8]}x0 = vacc01234567x0; 50 $for K in range(1, UNROLL): 51 float16x8_t vacc${ABC[M:M+8]}x${K} = vmovq_n_f16(0.0f); 52 for (; nnz >= ${UNROLL}; nnz -= ${UNROLL}) { 53 $for K in range(UNROLL): 54 const intptr_t diff${K} = dmap[${K}]; 55 dmap += ${UNROLL}; 56 $for K in range(UNROLL): 57 const float16x8_t va01234567x${K} = vld1q_f16(i); 58 $for M in range(8, MR, 8): 59 const float16x8_t va${ABC[M:M+8]}x${K} = vld1q_f16(i + ${M}); 60 i = (const __fp16*restrict) ((uintptr_t) i + (uintptr_t) diff${K}); 61 const float16x8_t vb${K} = vld1q_dup_f16(w); w += 1; 62 $for M in range(0, MR, 8): 63 vacc${ABC[M:M+8]}x${K} = vfmaq_f16(vacc${ABC[M:M+8]}x${K}, va${ABC[M:M+8]}x${K}, vb${K}); 64 } 65 $for M in range(0, MR, 8): 66 float16x8_t vacc${ABC[M:M+8]} = vacc${ABC[M:M+8]}x0; 67 $for K in range(1, UNROLL): 68 $for M in range(0, MR, 8): 69 vacc${ABC[M:M+8]} = vaddq_f16(vacc${ABC[M:M+8]}, vacc${ABC[M:M+8]}x${K}); 70 $else: 71 float16x8_t vacc01234567 = vld1q_dup_f16(w); w += 1; 72 $for M in range(8, MR, 8): 73 float16x8_t vacc${ABC[M:M+8]} = vacc01234567; 74 if XNN_LIKELY(nnz != 0) { 75 do { 76 const intptr_t diff = *dmap++; 77 const float16x8_t va01234567 = vld1q_f16(i); 78 $for M in range(8, MR, 8): 79 const float16x8_t va${ABC[M:M+8]} = vld1q_f16(i + ${M}); 80 i = (const __fp16*restrict) ((uintptr_t) i + (uintptr_t) diff); 81 const float16x8_t vb = vld1q_dup_f16(w); w += 1; 82 $for M in range(0, MR, 8): 83 vacc${ABC[M:M+8]} = vfmaq_f16(vacc${ABC[M:M+8]}, va${ABC[M:M+8]}, vb); 84 } while (--nnz != 0); 85 } 86 $for M in range(0, MR, 8): 87 float16x8_t vout${ABC[M:M+8]} = vminq_f16(vacc${ABC[M:M+8]}, vmax); 88 $for M in range(0, MR, 8): 89 vout${ABC[M:M+8]} = vmaxq_f16(vout${ABC[M:M+8]}, vmin); 90 vst1q_f16(o, vout01234567); 91 $for M in range(8, MR, 8): 92 vst1q_f16(o + ${M}, vout${ABC[M:M+8]}); 93 o = (__fp16*restrict) ((uintptr_t) o + output_stride); 94 } while (--n != 0); 95 o = (__fp16*restrict) ((uintptr_t) o - output_decrement); 96 i += ${MR}; 97 mc -= ${MR} * sizeof(__fp16); 98 } 99 if XNN_UNLIKELY(mc != 0) { 100 $for LOG2M in reversed(range((MR - 1).bit_length())): 101 $SUBMR = 1 << LOG2M 102 $if SUBMR * 2 >= MR: 103 output_decrement += ${MR - SUBMR} * sizeof(__fp16); 104 $else: 105 output_decrement += ${SUBMR} * sizeof(__fp16); 106 if (mc & (${SUBMR} * sizeof(__fp16))) { 107 const __fp16*restrict w = (const __fp16*) weights; 108 const int32_t* dmap = widx_dmap; 109 const uint32_t* nnzmap = nidx_nnzmap; 110 size_t n = nc; 111 do { 112 uint32_t nnz = *nnzmap++; 113 $if SUBMR <= 4: 114 float16x4_t vacc${ABC[0:SUBMR]} = vld1_dup_f16(w); w += 1; 115 $else: 116 float16x8_t vacc01234567 = vld1q_dup_f16(w); w += 1; 117 $for M in range(8, SUBMR, 8): 118 float16x8_t vacc${ABC[M:M+8]} = vacc01234567; 119 if XNN_LIKELY(nnz != 0) { 120 do { 121 const intptr_t diff = *dmap++; 122 $if SUBMR == 1: 123 const float16x4_t va0 = vld1_dup_f16(i); 124 $elif SUBMR == 2: 125 const float16x4_t va01 = vreinterpret_f16_f32(vld1_dup_f32((const void*) i)); 126 $elif SUBMR == 4: 127 const float16x4_t va0123 = vld1_f16(i); 128 $else: 129 const float16x8_t va01234567 = vld1q_f16(i); 130 $for M in range(8, SUBMR, 8): 131 const float16x8_t va${ABC[M:M+8]} = vld1q_f16(i + ${M}); 132 i = (const __fp16*restrict) ((uintptr_t) i + (uintptr_t) diff); 133 $if SUBMR <= 4: 134 const float16x4_t vb = vld1_dup_f16(w); w += 1; 135 $else: 136 const float16x8_t vb = vld1q_dup_f16(w); w += 1; 137 $if SUBMR <= 4: 138 vacc${ABC[0:SUBMR]} = vfma_f16(vacc${ABC[0:SUBMR]}, va${ABC[0:SUBMR]}, vb); 139 $else: 140 $for M in range(0, SUBMR, 8): 141 vacc${ABC[M:M+8]} = vfmaq_f16(vacc${ABC[M:M+8]}, va${ABC[M:M+8]}, vb); 142 } while (--nnz != 0); 143 } 144 $if SUBMR <= 4: 145 float16x4_t vout${ABC[0:SUBMR]} = vmin_f16(vacc${ABC[0:SUBMR]}, vget_low_f16(vmax)); 146 vout${ABC[0:SUBMR]} = vmax_f16(vout${ABC[0:SUBMR]}, vget_low_f16(vmin)); 147 $if SUBMR == 1: 148 vst1_lane_f16(o, vout${ABC[0]}, 0); 149 $elif SUBMR == 2: 150 vst1_lane_f32((void*) o, vreinterpret_f32_f16(vout${ABC[0:SUBMR]}), 0); 151 $else: 152 vst1_f16(o, vout${ABC[0:SUBMR]}); 153 $else: 154 $for M in range(0, SUBMR, 8): 155 float16x8_t vout${ABC[M:M+8]} = vminq_f16(vacc${ABC[M:M+8]}, vmax); 156 $for M in range(0, SUBMR, 8): 157 vout${ABC[M:M+8]} = vmaxq_f16(vout${ABC[M:M+8]}, vmin); 158 vst1q_f16(o, vout01234567); 159 $for M in range(8, SUBMR, 8): 160 vst1q_f16(o + ${M}, vout${ABC[M:M+8]}); 161 o = (__fp16*restrict) ((uintptr_t) o + output_stride); 162 } while (--n != 0); 163 o = (__fp16*restrict) ((uintptr_t) o - output_decrement); 164 i += ${SUBMR}; 165 } 166 } 167} 168