src/f32-vmulcaddc/psimd.c.in

// Copyright 2019 Google LLC
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

$assert CR % 4 == 0
$ABC = "0123456789ABCDEFGHIJKLMN"
#include <assert.h>

#include <psimd.h>

#include <xnnpack/math.h>
#include <xnnpack/vmulcaddc.h>


void xnn_f32_vmulcaddc_ukernel_c${CR}__psimd_x${MR}(
    size_t m,
    size_t channels,
    const float*restrict x,
    size_t x_stride,
    const float*restrict weights,
    float*restrict y,
    size_t y_stride,
    const union xnn_f32_output_params params[restrict static 1])
{
  assert(m != 0);
  assert(channels != 0);
  assert(channels % sizeof(float) == 0);

  const size_t x_increment = x_stride * ${MR} - (channels & -(${CR} * sizeof(float)));
  const size_t y_increment = y_stride * ${MR} - channels;

  const float* x0 = x;
  float* y0 = y;
  $for M in range(1, MR):
    const float* x${M} = (const float*) ((uintptr_t) x${M-1} + x_stride);
    float* y${M} = (float*) ((uintptr_t) y${M-1} + y_stride);
    $if M % 2 == 0:
      if XNN_UNPREDICTABLE(m <= ${M}) {
        x${M} = x${M-1};
        y${M} = y${M-1};
      }
    $else:
      if XNN_UNPREDICTABLE(m < ${M+1}) {
        x${M} = x${M-1};
        y${M} = y${M-1};
      }

  const psimd_f32 vmin = psimd_load_splat_f32(&params->scalar.min);
  const psimd_f32 vmax = psimd_load_splat_f32(&params->scalar.max);
  do {
    const float* w = weights;
    size_t c = channels;
    for (; c >= ${CR} * sizeof(float); c -= ${CR} * sizeof(float)) {
      const psimd_f32 vscale${ABC[0:4]} = psimd_load_f32(w);
      $for C in range(4, CR, 4):
        const psimd_f32 vscale${ABC[C:C+4]} = psimd_load_f32(w + ${C});

      $for M in range(MR):
        const psimd_f32 vx${M}x${ABC[0:4]} = psimd_load_f32(x${M});
        $for C in range(4, CR, 4):
          const psimd_f32 vx${M}x${ABC[C:C+4]} = psimd_load_f32(x${M} + ${C});
        x${M} += ${CR};

      $for C in range(0, CR, 4):
        const psimd_f32 vbias${ABC[C:C+4]} = psimd_load_f32(w + ${C + CR});

      $for M in range(MR):
        $for C in range(0, CR, 4):
          psimd_f32 vacc${M}x${ABC[C:C+4]} = psimd_qfma_f32(vbias${ABC[C:C+4]}, vx${M}x${ABC[C:C+4]}, vscale${ABC[C:C+4]});

      $for M in range(MR):
        $for C in range(0, CR, 4):
          vacc${M}x${ABC[C:C+4]} = psimd_max_f32(vacc${M}x${ABC[C:C+4]}, vmin);

      $for M in range(MR):
        $for C in range(0, CR, 4):
          vacc${M}x${ABC[C:C+4]} = psimd_min_f32(vacc${M}x${ABC[C:C+4]}, vmax);

      $for M in range(MR):
        psimd_store_f32(y${M}, vacc${M}x${ABC[0:4]});
        $for C in range(4, CR, 4):
          psimd_store_f32(y${M} + ${C}, vacc${M}x${ABC[C:C+4]});
        y${M} += ${CR};

      w += ${CR * 2};
    }
    if XNN_UNLIKELY(c != 0) {
      const psimd_f32 vscale${ABC[0:4]} = psimd_load_f32(w);
      $for C in range(4, CR, 4):
        const psimd_f32 vscale${ABC[C:C+4]} = psimd_load_f32(w + ${C});

      $for M in range(MR):
        const psimd_f32 vx${M}x${ABC[0:4]} = psimd_load_f32(x${M});
        $for C in range(4, CR, 4):
          const psimd_f32 vx${M}x${ABC[C:C+4]} = psimd_load_f32(x${M} + ${C});

      $for C in range(0, CR, 4):
        const psimd_f32 vbias${ABC[C:C+4]} = psimd_load_f32(w + ${C + CR});

      $for M in range(MR):
        $for C in range(0, CR, 4):
          psimd_f32 vacc${M}x${ABC[C:C+4]} = psimd_qfma_f32(vbias${ABC[C:C+4]}, vx${M}x${ABC[C:C+4]}, vscale${ABC[C:C+4]});

      $for M in range(MR):
        $for C in range(0, CR, 4):
          vacc${M}x${ABC[C:C+4]} = psimd_max_f32(vacc${M}x${ABC[C:C+4]}, vmin);

      $for M in range(MR):
        $for C in range(0, CR, 4):
          vacc${M}x${ABC[C:C+4]} = psimd_min_f32(vacc${M}x${ABC[C:C+4]}, vmax);

      w += ${CR * 2};

      $for LOG2C in reversed(range(CR.bit_length())):
        $if CR != 1 << LOG2C:
          if (c & (${1 << LOG2C} * sizeof(float))) {
            $if LOG2C >= 2:
              $for M in range(MR):
                psimd_store_f32(y${M}, vacc${M}x${ABC[C:C+4]});
                $for C in range(4, 1 << LOG2C, 4):
                  psimd_store_f32(y${M} + ${C}, vacc${M}x${ABC[C:C+4]});

              $for M in range(MR):
                y${M} += ${1 << LOG2C};

              $for M in range(MR):
                $for C in range(0, 1 << (LOG2C - 1), 4):
                  vacc${M}x${ABC[C:C+4]} = vacc${M}x${ABC[C+(1<<LOG2C):C+(1<<LOG2C)+4]}
            $elif LOG2C == 1:
              $for M in range(MR):
                psimd_store2_f32(y${M}, vacc${M}x${ABC[0:4]});

              $for M in range(MR):
                y${M} += 2;

              $for M in range(MR):
                vacc${M}x${ABC[0:4]} = psimd_concat_hi_f32(vacc${M}x${ABC[0:4]}, vacc${M}x${ABC[0:4]});
            $elif LOG2C == 0:
              $for M in range(MR):
                psimd_store1_f32(y${M}, vacc${M}x${ABC[0:4]});

              $for M in range(MR):
                y${M} += 1;
          }
    }
    $for M in range(MR):
      x${M} = (const float*) ((uintptr_t) x${M} + x_increment);
      y${M} = (float*) ((uintptr_t) y${M} + y_increment);
      $if M % 2 == 1:
        if XNN_UNPREDICTABLE(m < ${MR + M+1}) {
          x${M} = x${M-1};
          y${M} = y${M-1};
        }
      $elif M != 0:
        if XNN_UNPREDICTABLE(m <= ${MR + M}) {
          x${M} = x${M-1};
          y${M} = y${M-1};
        }
    m = doz(m, ${MR});
  } while (m != 0);
}