• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 CHANNEL_TILE % 8 == 0
7$assert KERNEL_TILE >= 2
8$assert ACCUMULATORS >= 1
9$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
10#include <assert.h>
11
12#include <immintrin.h>
13
14#include <xnnpack/dwconv.h>
15
16
17static const int32_t mask_table[14] = {-1, -1, -1, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0};
18
19$ISA = {0: "avx", 3: "fma3"}[FMA]
20void xnn_f32_dwconv_minmax_ukernel_up${CHANNEL_TILE}x${KERNEL_TILE}__${ISA}${"" if ACCUMULATORS == 1 else "_acc%d" % ACCUMULATORS}(
21    size_t channels,
22    size_t output_width,
23    const float** input,
24    const float* weights,
25    float* output,
26    size_t input_stride,
27    size_t output_increment,
28    size_t input_offset,
29    const float* zero,
30    const union xnn_f32_minmax_params params[restrict XNN_MIN_ELEMENTS(1)])
31{
32  assert(channels != 0);
33  assert(output_width != 0);
34
35  const __m256 vmax = _mm256_broadcast_ps((const __m128*) params->sse.max);
36  const __m256 vmin = _mm256_broadcast_ps((const __m128*) params->sse.min);
37  do {
38    $for K in range(KERNEL_TILE):
39      const float* i${K} = input[${K}];
40      assert(i${K} != NULL);
41      if XNN_UNPREDICTABLE(i${K} != zero) {
42        i${K} = (const float*) ((uintptr_t) i${K} + input_offset);
43      }
44    input = (const float**) ((uintptr_t) input + input_stride);
45
46    size_t c = channels;
47    const float* w = weights;
48    for (; c >= ${CHANNEL_TILE}; c -= ${CHANNEL_TILE}) {
49      __m256 vacc${ABC[0:8]}p0 = _mm256_load_ps(w);
50      $for C in range(8, CHANNEL_TILE, 8):
51        __m256 vacc${ABC[C:C+8]}p0 = _mm256_load_ps(w + ${C});
52
53      $for K in range(KERNEL_TILE):
54
55        const __m256 vi${K}x${ABC[0:8]} = _mm256_loadu_ps(i${K});
56        $for C in range(8, CHANNEL_TILE, 8):
57          const __m256 vi${K}x${ABC[C:C+8]} = _mm256_loadu_ps(i${K} + ${C});
58        i${K} += ${CHANNEL_TILE};
59
60        $for C in range(0, CHANNEL_TILE, 8):
61          const __m256 vk${K}x${ABC[C:C+8]} = _mm256_load_ps(w + ${(K + 1) * CHANNEL_TILE + C});
62        $for C in range(0, CHANNEL_TILE, 8):
63          $if 1 <= K < ACCUMULATORS:
64            __m256 vacc${ABC[C:C+8]}p${K} = _mm256_mul_ps(vi${K}x${ABC[C:C+8]}, vk${K}x${ABC[C:C+8]});
65          $elif FMA == 3:
66            vacc${ABC[C:C+8]}p${K % ACCUMULATORS} = _mm256_fmadd_ps(vi${K}x${ABC[C:C+8]}, vk${K}x${ABC[C:C+8]}, vacc${ABC[C:C+8]}p${K % ACCUMULATORS});
67          $else:
68            vacc${ABC[C:C+8]}p${K % ACCUMULATORS} = _mm256_add_ps(vacc${ABC[C:C+8]}p${K % ACCUMULATORS}, _mm256_mul_ps(vi${K}x${ABC[C:C+8]}, vk${K}x${ABC[C:C+8]}));
69
70      w += ${(KERNEL_TILE + 1) * CHANNEL_TILE};
71
72      $if ACCUMULATORS > 1:
73        // Add up all accumulators to vacc${ABC[0:CHANNEL_TILE]}p0
74        $ACC_SLICE = 1
75        $while ACC_SLICE < ACCUMULATORS:
76          $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
77            $if A + ACC_SLICE < ACCUMULATORS:
78              $for C in range(0, CHANNEL_TILE, 8):
79                vacc${ABC[C:C+8]}p${A} = _mm256_add_ps(vacc${ABC[C:C+8]}p${A}, vacc${ABC[C:C+8]}p${A + ACC_SLICE});
80          $ACC_SLICE *= 2
81
82      $for C in range(0, CHANNEL_TILE, 8):
83        __m256 vacc${ABC[C:C+8]} = _mm256_max_ps(vacc${ABC[C:C+8]}p0, vmin);
84      $for C in range(0, CHANNEL_TILE, 8):
85        vacc${ABC[C:C+8]} = _mm256_min_ps(vacc${ABC[C:C+8]}, vmax);
86
87      _mm256_storeu_ps(output, vacc${ABC[0:8]});
88      $for C in range(8, CHANNEL_TILE, 8):
89        _mm256_storeu_ps(output + ${C}, vacc${ABC[C:C+8]});
90      output += ${CHANNEL_TILE};
91    }
92    $if CHANNEL_TILE > 8:
93      for (; c >= 8; c -= 8) {
94        __m256 vacc01234567p0 = _mm256_load_ps(w);
95        $for K in range(KERNEL_TILE):
96
97          const __m256 vi${K}x01234567 = _mm256_loadu_ps(i${K});
98          i${K} += 8;
99
100          const __m256 vk${K}x01234567 = _mm256_load_ps(w + ${(K + 1) * CHANNEL_TILE});
101          $if 1 <= K < ACCUMULATORS:
102            __m256 vacc01234567p${K} = _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567);
103          $elif FMA == 3:
104            vacc01234567p${K % ACCUMULATORS} = _mm256_fmadd_ps(vi${K}x01234567, vk${K}x01234567, vacc01234567p${K % ACCUMULATORS});
105          $else:
106            vacc01234567p${K % ACCUMULATORS} = _mm256_add_ps(vacc01234567p${K % ACCUMULATORS}, _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567));
107
108        w += 8;
109
110        $if ACCUMULATORS > 1:
111          // Add up all accumulators to vacc${ABC[0:8]}p0
112          $ACC_SLICE = 1
113          $while ACC_SLICE < ACCUMULATORS:
114            $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
115              $if A + ACC_SLICE < ACCUMULATORS:
116                vacc01234567p${A} = _mm256_add_ps(vacc01234567p${A}, vacc01234567p${A + ACC_SLICE});
117            $ACC_SLICE *= 2
118
119        __m256 vacc01234567 = _mm256_max_ps(vacc01234567p0, vmin);
120        vacc01234567 = _mm256_min_ps(vacc01234567, vmax);
121
122        _mm256_storeu_ps(output, vacc01234567);
123        output += 8;
124      }
125    if XNN_UNLIKELY(c != 0) {
126      assert(c >= 1);
127      assert(c <= 7);
128      __m256i vmask = _mm256_loadu_si256((const __m256i*) &mask_table[7 - c]);
129
130      __m256 vacc01234567p0 = _mm256_load_ps(w);
131      $for K in range(KERNEL_TILE):
132
133        const __m256 vi${K}x01234567 = _mm256_maskload_ps(i${K}, vmask);
134        const __m256 vk${K}x01234567 = _mm256_load_ps(w + ${(K + 1) * CHANNEL_TILE});
135        $if 1 <= K < ACCUMULATORS:
136          __m256 vacc01234567p${K} = _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567);
137        $elif FMA == 3:
138          vacc01234567p${K % ACCUMULATORS} = _mm256_fmadd_ps(vi${K}x01234567, vk${K}x01234567, vacc01234567p${K % ACCUMULATORS});
139        $else:
140          vacc01234567p${K % ACCUMULATORS} = _mm256_add_ps(vacc01234567p${K % ACCUMULATORS}, _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567));
141
142      $if ACCUMULATORS > 1:
143        // Add up all accumulators to vacc${ABC[0:8]}p0
144        $ACC_SLICE = 1
145        $while ACC_SLICE < ACCUMULATORS:
146          $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
147            $if A + ACC_SLICE < ACCUMULATORS:
148              vacc01234567p${A} = _mm256_add_ps(vacc01234567p${A}, vacc01234567p${A + ACC_SLICE});
149          $ACC_SLICE *= 2
150
151      __m256 vacc01234567 = _mm256_max_ps(vacc01234567p0, vmin);
152      vacc01234567 = _mm256_min_ps(vacc01234567, vmax);
153
154      // _mm256_maskstore_ps(output, vmask, vacc01234567); output += c; could be used here, but triggers msan failures (probably an msan bug).
155      __m128 vacc0123 = _mm256_castps256_ps128(vacc01234567);
156      if (c & 4) {
157        _mm_storeu_ps(output, vacc0123);
158        vacc0123 = _mm256_extractf128_ps(vacc01234567, 1);
159        output += 4;
160      }
161      if (c & 2) {
162        _mm_storel_pi((__m64*) output, vacc0123);
163        vacc0123 = _mm_movehl_ps(vacc0123, vacc0123);
164        output += 2;
165      }
166      if (c & 1) {
167        _mm_store_ss(output, vacc0123);
168        output += 1;
169      }
170    }
171
172    output = (float*) ((uintptr_t) output + output_increment);
173  } while (--output_width != 0);
174}
175