• 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 % 4 == 0
7$assert KERNEL_TILE >= 2
8$assert ACCUMULATORS >= 1
9$ABC = "0123456789ABCDEF"
10#include <assert.h>
11
12#include <psimd.h>
13
14#include <xnnpack/dwconv.h>
15
16
17void xnn_f32_dwconv_ukernel_up${CHANNEL_TILE}x${KERNEL_TILE}__psimd${"" if ACCUMULATORS == 1 else "_acc%d" % ACCUMULATORS}(
18    size_t channels,
19    size_t output_width,
20    const float** input,
21    const float* weights,
22    float* output,
23    size_t input_stride,
24    size_t output_increment,
25    const union xnn_f32_output_params params[restrict static 1])
26{
27  assert(channels != 0);
28  assert(output_width != 0);
29
30  const psimd_f32 vmax = psimd_load_splat_f32(&params->scalar.max);
31  const psimd_f32 vmin = psimd_load_splat_f32(&params->scalar.min);
32  do {
33    $for K in range(KERNEL_TILE):
34      const float* i${K} = input[${K}];
35      assert(i${K} != NULL);
36    input = (const float**) ((uintptr_t) input + input_stride);
37
38    size_t c = channels;
39    const float* w = weights;
40    for (; c >= ${CHANNEL_TILE}; c -= ${CHANNEL_TILE}) {
41      psimd_f32 vacc${ABC[0:4]}p0 = psimd_load_f32(w);
42      $for C in range(4, CHANNEL_TILE, 4):
43        psimd_f32 vacc${ABC[C:C+4]}p0 = psimd_load_f32(w + ${C});
44
45      $for K in range(KERNEL_TILE):
46
47        const psimd_f32 vi${K}x${ABC[0:4]} = psimd_load_f32(i${K});
48        $for C in range(4, CHANNEL_TILE, 4):
49          const psimd_f32 vi${K}x${ABC[C:C+4]} = psimd_load_f32(i${K} + ${C});
50        i${K} += ${CHANNEL_TILE};
51
52        $for C in range(0, CHANNEL_TILE, 4):
53          const psimd_f32 vk${K}x${ABC[C:C+4]} = psimd_load_f32(w + ${(K + 1) * CHANNEL_TILE + C});
54        $for C in range(0, CHANNEL_TILE, 4):
55          $if 1 <= K < ACCUMULATORS:
56            psimd_f32 vacc${ABC[C:C+4]}p${K} = psimd_mul_f32(vi${K}x${ABC[C:C+4]}, vk${K}x${ABC[C:C+4]});
57          $else:
58            vacc${ABC[C:C+4]}p${K % ACCUMULATORS} = psimd_qfma_f32(vacc${ABC[C:C+4]}p${K % ACCUMULATORS}, vi${K}x${ABC[C:C+4]}, vk${K}x${ABC[C:C+4]});
59
60      w += ${(KERNEL_TILE + 1) * CHANNEL_TILE};
61
62      $if ACCUMULATORS > 1:
63        // Add up all accumulators to vacc${ABC[0:CHANNEL_TILE]}p0
64        $ACC_SLICE = 1
65        $while ACC_SLICE < ACCUMULATORS:
66          $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
67            $if A + ACC_SLICE < ACCUMULATORS:
68              $for C in range(0, CHANNEL_TILE, 4):
69                vacc${ABC[C:C+4]}p${A} = psimd_add_f32(vacc${ABC[C:C+4]}p${A}, vacc${ABC[C:C+4]}p${A + ACC_SLICE});
70          $ACC_SLICE *= 2
71
72      $for C in range(0, CHANNEL_TILE, 4):
73        psimd_f32 vacc${ABC[C:C+4]} = psimd_max_f32(vacc${ABC[C:C+4]}p0, vmin);
74      $for C in range(0, CHANNEL_TILE, 4):
75        vacc${ABC[C:C+4]} = psimd_min_f32(vacc${ABC[C:C+4]}, vmax);
76
77      psimd_store_f32(output, vacc${ABC[0:4]});
78      $for C in range(4, CHANNEL_TILE, 4):
79        psimd_store_f32(output + ${C}, vacc${ABC[C:C+4]});
80      output += ${CHANNEL_TILE};
81    }
82    $if CHANNEL_TILE > 4:
83      for (; c >= 4; c -= 4) {
84        psimd_f32 vacc0123p0 = psimd_load_f32(w);
85        $for K in range(KERNEL_TILE):
86
87          const psimd_f32 vi${K}x0123 = psimd_load_f32(i${K});
88          i${K} += 4;
89
90          const psimd_f32 vk${K}x0123 = psimd_load_f32(w + ${(K + 1) * CHANNEL_TILE});
91          $if 1 <= K < ACCUMULATORS:
92            psimd_f32 vacc0123p${K} = psimd_mul_f32(vi${K}x0123, vk${K}x0123);
93          $else:
94            vacc0123p${K % ACCUMULATORS} = psimd_qfma_f32(vacc0123p${K % ACCUMULATORS}, vi${K}x0123, vk${K}x0123);
95
96        w += 4;
97
98        $if ACCUMULATORS > 1:
99          // Add up all accumulators to vacc${ABC[0:CHANNEL_TILE]}p0
100          $ACC_SLICE = 1
101          $while ACC_SLICE < ACCUMULATORS:
102            $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
103              $if A + ACC_SLICE < ACCUMULATORS:
104                vacc0123p${A} = psimd_add_f32(vacc0123p${A}, vacc0123p${A + ACC_SLICE});
105            $ACC_SLICE *= 2
106
107        psimd_f32 vacc0123 = psimd_max_f32(vacc0123p0, vmin);
108        vacc0123 = psimd_min_f32(vacc0123, vmax);
109
110        psimd_store_f32(output, vacc0123);
111        output += 4;
112      }
113    if XNN_UNLIKELY(c != 0) {
114      psimd_f32 vacc0123p0 = psimd_load_f32(w);
115      $for K in range(KERNEL_TILE):
116
117        const psimd_f32 vi${K}x0123 = psimd_load_f32(i${K});
118        const psimd_f32 vk${K}x0123 = psimd_load_f32(w + ${(K+1) * CHANNEL_TILE});
119        $if 1 <= K < ACCUMULATORS:
120          psimd_f32 vacc0123p${K} = psimd_mul_f32(vi${K}x0123, vk${K}x0123);
121        $else:
122          vacc0123p${K % ACCUMULATORS} = psimd_qfma_f32(vacc0123p${K % ACCUMULATORS}, vi${K}x0123, vk${K}x0123);
123
124      $if ACCUMULATORS > 1:
125        // Add up all accumulators to vacc${ABC[0:CHANNEL_TILE]}p0
126        $ACC_SLICE = 1
127        $while ACC_SLICE < ACCUMULATORS:
128          $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
129            $if A + ACC_SLICE < ACCUMULATORS:
130              vacc0123p${A} = psimd_add_f32(vacc0123p${A}, vacc0123p${A + ACC_SLICE});
131          $ACC_SLICE *= 2
132
133      psimd_f32 vacc0123 = psimd_max_f32(vacc0123p0, vmin);
134      vacc0123 = psimd_min_f32(vacc0123, vmax);
135
136      if (c & 2) {
137        psimd_store2_f32(output, vacc0123);
138        vacc0123 = psimd_concat_hi_f32(vacc0123, vacc0123);
139        output += 2;
140      }
141      if (c & 1) {
142        psimd_store1_f32(output, vacc0123);
143        output += 1;
144      }
145    }
146
147    output = (float*) ((uintptr_t) output + output_increment);
148  } while (--output_width != 0);
149}
150