1// Copyright 2020 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 NR % 4 == 0 7$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" 8#include <assert.h> 9 10#include <wasm_simd128.h> 11 12#include <xnnpack/igemm.h> 13 14 15$assert ACTIVATION in ["LINEAR", "RELU", "MINMAX"] 16$ACTIVATION_SUFFIX = {"LINEAR": ""}.get(ACTIVATION, "_" + ACTIVATION.lower()) 17$ARCH_SUFFIX = "" if ACTIVATION in ["LINEAR", "RELU"] else "_x86" if X86 else "_arm" 18$PARAMS = {"LINEAR": "xnn_f32_default_params", "RELU": "xnn_f32_relu_params", "MINMAX": "xnn_f32_minmax_params"}[ACTIVATION] 19void xnn_f32_igemm${ACTIVATION_SUFFIX}_ukernel_${MR}x${NR}__wasmsimd${ARCH_SUFFIX}_loadsplat( 20 size_t mr, 21 size_t nc, 22 size_t kc, 23 size_t ks, 24 const float**restrict a, 25 const float*restrict w, 26 float*restrict c, 27 size_t cm_stride, 28 size_t cn_stride, 29 size_t a_offset, 30 const float* zero, 31 const union ${PARAMS} params[restrict XNN_MIN_ELEMENTS(1)]) 32{ 33 assert(mr != 0); 34 assert(mr <= ${MR}); 35 assert(nc != 0); 36 assert(kc != 0); 37 assert(kc % sizeof(float) == 0); 38 assert(ks != 0); 39 assert(ks % (${MR} * sizeof(void*)) == 0); 40 assert(a_offset % sizeof(float) == 0); 41 assert(a != NULL); 42 assert(w != NULL); 43 assert(c != NULL); 44 45 float* c0 = c; 46 $for M in range(1, MR): 47 float* c${M} = (float*) ((uintptr_t) c${M-1} + cm_stride); 48 $if M % 2 == 0: 49 if XNN_UNPREDICTABLE(mr <= ${M}) { 50 c${M} = c${M-1}; 51 } 52 $elif M + 1 == MR: 53 if XNN_UNPREDICTABLE(mr != ${M+1}) { 54 c${M} = c${M-1}; 55 } 56 $else: 57 if XNN_UNPREDICTABLE(mr < ${M+1}) { 58 c${M} = c${M-1}; 59 } 60 61 $if ACTIVATION == "MINMAX" and not X86: 62 const v128_t vmin = wasm_v32x4_load_splat(¶ms->scalar.min); 63 const v128_t vmax = wasm_v32x4_load_splat(¶ms->scalar.max); 64 do { 65 v128_t vacc0x${ABC[0:4]} = wasm_v128_load(w); 66 $for N in range(4, NR, 4): 67 v128_t vacc0x${ABC[N:N+4]} = wasm_v128_load(w + ${N}); 68 $for M in range(1, MR): 69 $for N in range(0, NR, 4): 70 v128_t vacc${M}x${ABC[N:N+4]} = vacc0x${ABC[N:N+4]}; 71 w += ${NR}; 72 73 size_t p = ks; 74 do { 75 $for M in range(MR): 76 const float* restrict a${M} = a[${M}]; 77 assert(a${M} != NULL); 78 if XNN_UNPREDICTABLE(a${M} != zero) { 79 a${M} = (const float*) ((uintptr_t) a${M} + a_offset); 80 } 81 a += ${MR}; 82 83 size_t k = kc; 84 do { 85 const v128_t vb${ABC[0:4]} = wasm_v128_load(w); 86 $for N in range(4, NR, 4): 87 const v128_t vb${ABC[N:N+4]} = wasm_v128_load(w + ${N}); 88 w += ${NR}; 89 90 $for M in range(MR): 91 const v128_t va${M} = wasm_v32x4_load_splat(a${M}); 92 a${M} += 1; 93 94 $for M in range(MR): 95 $for N in range(0, NR, 4): 96 vacc${M}x${ABC[N:N+4]} = wasm_f32x4_add(vacc${M}x${ABC[N:N+4]}, wasm_f32x4_mul(va${M}, vb${ABC[N:N+4]})); 97 k -= sizeof(float); 98 } while (k != 0); 99 p -= ${MR} * sizeof(void*); 100 } while (p != 0); 101 102 $if ACTIVATION == "MINMAX": 103 $if X86: 104 const v128_t vmin = wasm_v32x4_load_splat(¶ms->scalar.min); 105 $for N in range(0, NR, 4): 106 $for M in range(MR): 107 vacc${M}x${ABC[N:N+4]} = wasm_v128_bitselect(vmin, vacc${M}x${ABC[N:N+4]}, wasm_f32x4_lt(vacc${M}x${ABC[N:N+4]}, vmin)); 108 $else: 109 $for N in range(0, NR, 4): 110 $for M in range(MR): 111 vacc${M}x${ABC[N:N+4]} = wasm_f32x4_max(vacc${M}x${ABC[N:N+4]}, vmin); 112 113 $if X86: 114 const v128_t vmax = wasm_v32x4_load_splat(¶ms->scalar.max); 115 $for N in range(0, NR, 4): 116 $for M in range(MR): 117 vacc${M}x${ABC[N:N+4]} = wasm_v128_bitselect(vacc${M}x${ABC[N:N+4]}, vmax, wasm_f32x4_le(vacc${M}x${ABC[N:N+4]}, vmax)); 118 $else: 119 $for N in range(0, NR, 4): 120 $for M in range(MR): 121 vacc${M}x${ABC[N:N+4]} = wasm_f32x4_min(vacc${M}x${ABC[N:N+4]}, vmax); 122 $elif ACTIVATION == "RELU": 123 const v128_t vzero = wasm_f32x4_splat(0.0f); 124 $for N in range(0, NR, 4): 125 $for M in range(MR): 126 vacc${M}x${ABC[N:N+4]} = wasm_i32x4_max(vacc${M}x${ABC[N:N+4]}, vzero); 127 128 if XNN_LIKELY(nc >= ${NR}) { 129 $for M in reversed(range(MR)): 130 wasm_v128_store(c${M}, vacc${M}x${ABC[0:4]}); 131 $for N in range(4, NR, 4): 132 wasm_v128_store(c${M} + ${N}, vacc${M}x${ABC[N:N+4]}); 133 c${M} = (float*) ((uintptr_t) c${M} + cn_stride); 134 135 a = (const float**restrict) ((uintptr_t) a - ks); 136 nc -= ${NR}; 137 } else { 138 $for LOG2N in reversed(range(NR.bit_length())): 139 $if NR != 1 << LOG2N: 140 if (nc & ${1 << LOG2N}) { 141 $if LOG2N >= 2: 142 $for M in reversed(range(MR)): 143 wasm_v128_store(c${M}, vacc${M}x${ABC[0:4]}); 144 $for N in range(4, 1 << LOG2N, 4): 145 wasm_v128_store(c${M} + ${N}, vacc${M}x${ABC[N:N+4]}); 146 147 $for M in reversed(range(MR)): 148 $for N in range(0, 1 << (LOG2N - 1), 4): 149 vacc${M}x${ABC[N:N+4]} = vacc${M}x${ABC[N + (1 << LOG2N):N + (1 << LOG2N)+4]}; 150 151 $for M in reversed(range(MR)): 152 c${M} += ${1 << LOG2N}; 153 $elif LOG2N == 1: 154 $for M in reversed(range(MR)): 155 *((double*) c${M}) = wasm_f64x2_extract_lane(vacc${M}x${ABC[0:4]}, 0); 156 157 $for M in reversed(range(MR)): 158 vacc${M}x${ABC[0:4]} = wasm_v32x4_shuffle(vacc${M}x${ABC[0:4]}, vacc${M}x${ABC[0:4]}, 2, 3, 2, 3); 159 160 $for M in reversed(range(MR)): 161 c${M} += 2; 162 $elif LOG2N == 0: 163 $for M in reversed(range(MR)): 164 *c${M} = wasm_f32x4_extract_lane(vacc${M}x${ABC[0:4]}, 0); 165 } 166 167 nc = 0; 168 } 169 } while (nc != 0); 170} 171