src/f32-dwconv/multipass-avx512.c.in

// Copyright 2022 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.

$CHANNEL_SUBTILE = 16
$assert CHANNEL_TILE % CHANNEL_SUBTILE == 0
$CHANNEL_ROUND = 1
$assert MIDDLE_PASS_TILE <= LAST_PASS_TILE
$assert FIRST_PASS_TILE >= 1
$assert MIDDLE_PASS_TILE >= 1
$assert LAST_PASS_TILE >= 1
$assert ACCUMULATORS >= 1
$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
#include <assert.h>
#include <stddef.h>
#include <stdint.h>

#include <immintrin.h>

#include <xnnpack/dwconv.h>


void xnn_f32_dwconv_minmax_ukernel_${FIRST_PASS_TILE}f${MIDDLE_PASS_TILE}m${LAST_PASS_TILE}l${CHANNEL_TILE}c${CHANNEL_SUBTILE}s${CHANNEL_ROUND}r__avx512f${"" if ACCUMULATORS == 1 else "_acc%d" % ACCUMULATORS}(
    size_t channels,
    size_t output_width,
    const float** input,
    const float* weights,
    float* output,
    intptr_t input_stride,
    size_t output_increment,
    size_t input_offset,
    const float* zero,
    size_t kernel_size,
    float* buffer,
    const union xnn_f32_minmax_params params[restrict XNN_MIN_ELEMENTS(1)])
{
  assert(channels != 0);
  assert(output_width != 0);
  assert(kernel_size > ${FIRST_PASS_TILE});

  const __m512 vmin = _mm512_set1_ps(params->scalar.min);
  const __m512 vmax = _mm512_set1_ps(params->scalar.max);
  do {
    const float* w = weights;

    // First pass to process ${FIRST_PASS_TILE} inputs.
    {
      float* b = buffer;
      $for K in range(FIRST_PASS_TILE):
        const float* i${K} = input[${K}];
        assert(i${K} != NULL);
        if XNN_UNPREDICTABLE(i${K} != zero) {
          i${K} = (const float*) ((uintptr_t) i${K} + input_offset);
        }
      input += ${FIRST_PASS_TILE};

      // Process c channels and write to buffer.
      size_t c = channels;
      $if CHANNEL_TILE > 16:
        for (; c >= ${CHANNEL_TILE}; c -= ${CHANNEL_TILE}) {
          $for C in range(0, CHANNEL_TILE, 16):
            $if C == 0:
              __m512 vacc${ABC[C//16]}p0 = _mm512_load_ps(w);
            $else:
              __m512 vacc${ABC[C//16]}p0 = _mm512_load_ps(w + ${C});

          $for K in range(FIRST_PASS_TILE):

            $for C in range(0, CHANNEL_TILE, 16):
              $if C == 0:
                const __m512 vi${K}x${ABC[C//16]} = _mm512_loadu_ps(i${K});
              $else:
                const __m512 vi${K}x${ABC[C//16]} = _mm512_loadu_ps(i${K} + ${C});
            i${K} += ${CHANNEL_TILE};

            $for C in range(0, CHANNEL_TILE, 16):
              const __m512 vk${K}x${ABC[C//16]} = _mm512_load_ps(w + ${(K + 1) * CHANNEL_TILE + C});
            $for C in range(0, CHANNEL_TILE, 16):
              $if 1 <= K < ACCUMULATORS:
                __m512 vacc${ABC[C//16]}p${K} = _mm512_mul_ps(vi${K}x${ABC[C//16]}, vk${K}x${ABC[C//16]});
              $else:
                vacc${ABC[C//16]}p${K % ACCUMULATORS} = _mm512_fmadd_ps(vi${K}x${ABC[C//16]}, vk${K}x${ABC[C//16]}, vacc${ABC[C//16]}p${K % ACCUMULATORS});

          w += ${(FIRST_PASS_TILE + 1) * CHANNEL_TILE};

          $if ACCUMULATORS > 1:
            // Add up all accumulators to vacc0p0
            $ACC_SLICE = 1
            $while ACC_SLICE < ACCUMULATORS:
              $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
                $if A + ACC_SLICE < ACCUMULATORS:
                  $for C in range(0, CHANNEL_TILE, 16):
                    vacc${ABC[C//16]}p${A} = _mm512_add_ps(vacc${ABC[C//16]}p${A}, vacc${ABC[C//16]}p${A + ACC_SLICE});
              $ACC_SLICE *= 2

          $for C in range(0, CHANNEL_TILE, 16):
            $if C == 0:
              _mm512_store_ps(b, vacc${ABC[C//16]}p0);
            $else:
              _mm512_store_ps(b + ${C}, vacc${ABC[C//16]}p0);
          b += ${CHANNEL_TILE};
        }

      for (; c >= ${CHANNEL_SUBTILE}; c -= ${CHANNEL_SUBTILE}) {
        __m512 vaccp0 = _mm512_load_ps(w);

        $for K in range(FIRST_PASS_TILE):

          const __m512 vi${K}x0 = _mm512_loadu_ps(i${K});
          i${K} += ${CHANNEL_SUBTILE};

          const __m512 vk${K}x0 = _mm512_load_ps(w + ${(K + 1) * CHANNEL_SUBTILE});
          $if 1 <= K < ACCUMULATORS:
            __m512 vaccp${K} = _mm512_mul_ps(vi${K}x0, vk${K}x0);
          $else:
            vaccp${K % ACCUMULATORS} = _mm512_fmadd_ps(vi${K}x0, vk${K}x0, vaccp${K % ACCUMULATORS});

        w += ${(FIRST_PASS_TILE + 1) * CHANNEL_SUBTILE};

        $if ACCUMULATORS > 1:
          // Add up all accumulators to vaccp0
          $ACC_SLICE = 1
          $while ACC_SLICE < ACCUMULATORS:
            $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
              $if A + ACC_SLICE < ACCUMULATORS:
                vaccp${A} = _mm512_add_ps(vaccp${A}, vaccp${A + ACC_SLICE});
            $ACC_SLICE *= 2

        _mm512_store_ps(b, vaccp0);
        b += ${CHANNEL_SUBTILE};
      }

      if (c != 0) {
        assert(c >= 1);
        assert(c <= ${CHANNEL_SUBTILE-1});
        const __mmask16 vmask = _cvtu32_mask16((uint16_t) ((uint32_t) (UINT32_C(1) << c) - UINT32_C(1)));
        __m512 vaccp0 = _mm512_load_ps(w);

        $for K in range(FIRST_PASS_TILE):

          const __m512 vi${K}x0 = _mm512_maskz_loadu_ps(vmask, i${K});

          const __m512 vk${K}x0 = _mm512_load_ps(w + ${(K + 1) * CHANNEL_SUBTILE});
          $if 1 <= K < ACCUMULATORS:
            __m512 vaccp${K} = _mm512_mul_ps(vi${K}x0, vk${K}x0);
          $else:
            vaccp${K % ACCUMULATORS} = _mm512_fmadd_ps(vi${K}x0, vk${K}x0, vaccp${K % ACCUMULATORS});

        w += ${(FIRST_PASS_TILE + 1) * CHANNEL_SUBTILE};

        $if ACCUMULATORS > 1:
          // Add up all accumulators to vaccp0
          $ACC_SLICE = 1
          $while ACC_SLICE < ACCUMULATORS:
            $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
              $if A + ACC_SLICE < ACCUMULATORS:
                vaccp${A} = _mm512_add_ps(vaccp${A}, vaccp${A + ACC_SLICE});
            $ACC_SLICE *= 2

        _mm512_store_ps(b, vaccp0);
      }
    }

    // Middle pass to process ${MIDDLE_PASS_TILE} inputs in each iteration.
    for (size_t ks = kernel_size - ${FIRST_PASS_TILE}; ks > ${LAST_PASS_TILE}; ks -= ${MIDDLE_PASS_TILE}) {
      float* b = buffer;
      $for K in range(MIDDLE_PASS_TILE):
        const float* i${K} = input[${K}];
        assert(i${K} != NULL);
        if XNN_UNPREDICTABLE(i${K} != zero) {
          i${K} = (const float*) ((uintptr_t) i${K} + input_offset);
        }
      input += ${MIDDLE_PASS_TILE};

      size_t c = channels;
      $if CHANNEL_TILE > 16:
        for (; c >= ${CHANNEL_TILE}; c -= ${CHANNEL_TILE}) {
          $for C in range(0, CHANNEL_TILE, 16):
            $if C == 0:
              __m512 vacc${ABC[C//16]}p0 = _mm512_load_ps(b);
            $else:
              __m512 vacc${ABC[C//16]}p0 = _mm512_load_ps(b + ${C});

          $for K in range(MIDDLE_PASS_TILE):

            $for C in range(0, CHANNEL_TILE, 16):
              $if C == 0:
                const __m512 vi${K}x${ABC[C//16]} = _mm512_loadu_ps(i${K});
              $else:
                const __m512 vi${K}x${ABC[C//16]} = _mm512_loadu_ps(i${K} + ${C});
            i${K} += ${CHANNEL_TILE};

            $for C in range(0, CHANNEL_TILE, 16):
              $if K == 0 and C == 0:
                const __m512 vk${K}x${ABC[C//16]} = _mm512_load_ps(w);
              $else:
                const __m512 vk${K}x${ABC[C//16]} = _mm512_load_ps(w + ${K * CHANNEL_TILE + C});
            $for C in range(0, CHANNEL_TILE, 16):
              $if 1 <= K < ACCUMULATORS:
                __m512 vacc${ABC[C//16]}p${K} = _mm512_mul_ps(vi${K}x${ABC[C//16]}, vk${K}x${ABC[C//16]});
              $else:
                vacc${ABC[C//16]}p${K % ACCUMULATORS} = _mm512_fmadd_ps(vi${K}x${ABC[C//16]}, vk${K}x${ABC[C//16]}, vacc${ABC[C//16]}p${K % ACCUMULATORS});

          w += ${MIDDLE_PASS_TILE * CHANNEL_TILE};

          $if ACCUMULATORS > 1:
            // Add up all accumulators to vacc0p0
            $ACC_SLICE = 1
            $while ACC_SLICE < ACCUMULATORS:
              $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
                $if A + ACC_SLICE < ACCUMULATORS:
                  $for C in range(0, CHANNEL_TILE, 16):
                    vacc${ABC[C//16]}p${A} = _mm512_add_ps(vacc${ABC[C//16]}p${A}, vacc${ABC[C//16]}p${A + ACC_SLICE});
              $ACC_SLICE *= 2

          $for C in range(0, CHANNEL_TILE, 16):
            $if C == 0:
              _mm512_store_ps(b, vacc${ABC[C//16]}p0);
            $else:
              _mm512_store_ps(b + ${C}, vacc${ABC[C//16]}p0);
          b += ${CHANNEL_TILE};
        }

      for (; c >= 16; c -= 16) {
        __m512 vaccp0 = _mm512_load_ps(b);

        $for K in range(MIDDLE_PASS_TILE):

          const __m512 vi${K}x0 = _mm512_loadu_ps(i${K});
          i${K} += ${CHANNEL_SUBTILE};

          $if K == 0:
            const __m512 vk${K}x0 = _mm512_load_ps(w);
          $else:
            const __m512 vk${K}x0 = _mm512_load_ps(w + ${K * CHANNEL_SUBTILE});
          $if 1 <= K < ACCUMULATORS:
            __m512 vaccp${K} = _mm512_mul_ps(vi${K}x0, vk${K}x0);
          $else:
            vaccp${K % ACCUMULATORS} = _mm512_fmadd_ps(vi${K}x0, vk${K}x0, vaccp${K % ACCUMULATORS});

        w += ${MIDDLE_PASS_TILE * CHANNEL_SUBTILE};

        $if ACCUMULATORS > 1:
          // Add up all accumulators to vaccp0
          $ACC_SLICE = 1
          $while ACC_SLICE < ACCUMULATORS:
            $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
              $if A + ACC_SLICE < ACCUMULATORS:
                vaccp${A} = _mm512_add_ps(vaccp${A}, vaccp${A + ACC_SLICE});
            $ACC_SLICE *= 2

        _mm512_store_ps(b, vaccp0);
        b += ${CHANNEL_SUBTILE};
      }

      if (c != 0) {
        assert(c >= 1);
        assert(c <= ${CHANNEL_SUBTILE-1});
        const __mmask16 vmask = _cvtu32_mask16((uint16_t) ((uint32_t) (UINT32_C(1) << c) - UINT32_C(1)));
        __m512 vaccp0 = _mm512_load_ps(b);

        $for K in range(MIDDLE_PASS_TILE):

          const __m512 vi${K}x0 = _mm512_maskz_loadu_ps(vmask, i${K});

          $if K == 0:
            const __m512 vk${K}x0 = _mm512_load_ps(w);
          $else:
            const __m512 vk${K}x0 = _mm512_load_ps(w + ${(K) * CHANNEL_SUBTILE});
          $if 1 <= K < ACCUMULATORS:
            __m512 vaccp${K} = _mm512_mul_ps(vi${K}x0, vk${K}x0);
          $else:
            vaccp${K % ACCUMULATORS} = _mm512_fmadd_ps(vi${K}x0, vk${K}x0, vaccp${K % ACCUMULATORS});

        w += ${(MIDDLE_PASS_TILE) * CHANNEL_SUBTILE};

        $if ACCUMULATORS > 1:
          // Add up all accumulators to vaccp0
          $ACC_SLICE = 1
          $while ACC_SLICE < ACCUMULATORS:
            $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
              $if A + ACC_SLICE < ACCUMULATORS:
                vaccp${A} = _mm512_add_ps(vaccp${A}, vaccp${A + ACC_SLICE});
            $ACC_SLICE *= 2

        _mm512_store_ps(b, vaccp0);
      }
    }

    // Last pass to process up to ${LAST_PASS_TILE} inputs.
    {
      float* b = buffer;
      $for K in range(0, LAST_PASS_TILE):
        const float* i${K} = input[${K}];
        assert(i${K} != NULL);
        if XNN_UNPREDICTABLE(i${K} != zero) {
          i${K} = (const float*) ((uintptr_t) i${K} + input_offset);
        }

      size_t c = channels;
      $if CHANNEL_TILE > 16:
        for (; c >= ${CHANNEL_TILE}; c -= ${CHANNEL_TILE}) {
          $for C in range(0, CHANNEL_TILE, 16):
            $if C == 0:
              __m512 vacc${ABC[C//16]}p0 = _mm512_load_ps(b);
            $else:
              __m512 vacc${ABC[C//16]}p0 = _mm512_load_ps(b + ${C});
          b += ${CHANNEL_TILE};

          $for K in range(LAST_PASS_TILE):

            $for C in range(0, CHANNEL_TILE, 16):
              $if C == 0:
                const __m512 vi${K}x${ABC[C//16]} = _mm512_loadu_ps(i${K});
              $else:
                const __m512 vi${K}x${ABC[C//16]} = _mm512_loadu_ps(i${K} + ${C});
            i${K} += ${CHANNEL_TILE};

            $for C in range(0, CHANNEL_TILE, 16):
              $if K == 0 and C == 0:
                __m512 vk${K}x${ABC[C//16]} = _mm512_load_ps(w);
              $else:
                __m512 vk${K}x${ABC[C//16]} = _mm512_load_ps(w + ${K * CHANNEL_TILE + C});

            $for C in range(0, CHANNEL_TILE, 16):
              $if 1 <= K < ACCUMULATORS:
                __m512 vacc${ABC[C//16]}p${K} = _mm512_mul_ps(vi${K}x${ABC[C//16]}, vk${K}x${ABC[C//16]});
              $else:
                vacc${ABC[C//16]}p${K % ACCUMULATORS} = _mm512_fmadd_ps(vi${K}x${ABC[C//16]}, vk${K}x${ABC[C//16]}, vacc${ABC[C//16]}p${K % ACCUMULATORS});

          w += ${LAST_PASS_TILE * CHANNEL_TILE};

          $if ACCUMULATORS > 1:
            // Add up all accumulators to vacc0p0
            $ACC_SLICE = 1
            $while ACC_SLICE < ACCUMULATORS:
              $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
                $if A + ACC_SLICE < ACCUMULATORS:
                  $for C in range(0, CHANNEL_TILE, 16):
                    vacc${ABC[C//16]}p${A} = _mm512_add_ps(vacc${ABC[C//16]}p${A}, vacc${ABC[C//16]}p${A + ACC_SLICE});
              $ACC_SLICE *= 2

          $for C in range(0, CHANNEL_TILE, 16):
            __m512 vacc${ABC[C//16]} = _mm512_max_ps(vmin, vacc${ABC[C//16]}p0);

          $for C in range(0, CHANNEL_TILE, 16):
            vacc${ABC[C//16]} = _mm512_min_ps(vmax, vacc${ABC[C//16]});

          $for C in range(0, CHANNEL_TILE, 16):
            $if C == 0:
              _mm512_storeu_ps(output, vacc${ABC[C//16]});
            $else:
              _mm512_storeu_ps(output + ${C}, vacc${ABC[C//16]});
          output += ${CHANNEL_TILE};
        }


      for (; c >= 16; c -= 16) {
        __m512 vaccp0 = _mm512_load_ps(b);
        b += 16;

        $for K in range(LAST_PASS_TILE):

          const __m512 vi${K}x0 = _mm512_loadu_ps(i${K});
          i${K} += 16;

          $if K == 0:
            __m512 vk${K}x0 = _mm512_load_ps(w);
          $else:
            __m512 vk${K}x0 = _mm512_load_ps(w + ${K * 16});

          $if 1 <= K < ACCUMULATORS:
            __m512 vaccp${K} = _mm512_mul_ps(vi${K}x0, vk${K}x0);
          $else:
            vaccp${K % ACCUMULATORS} = _mm512_fmadd_ps(vi${K}x0, vk${K}x0, vaccp${K % ACCUMULATORS});

        $if CHANNEL_TILE > 16:
          w += ${LAST_PASS_TILE * 16};
        $else:
          w += ${LAST_PASS_TILE * CHANNEL_TILE};


        $if ACCUMULATORS > 1:
          // Add up all accumulators to vaccp0
          $ACC_SLICE = 1
          $while ACC_SLICE < ACCUMULATORS:
            $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
              $if A + ACC_SLICE < ACCUMULATORS:
                vaccp${A} = _mm512_add_ps(vaccp${A}, vaccp${A + ACC_SLICE});
            $ACC_SLICE *= 2

        __m512 vacc = _mm512_max_ps(vmin, vaccp0);

        vacc = _mm512_min_ps(vmax, vacc);

        _mm512_storeu_ps(output, vacc);
        output += 16;
      }

      if XNN_UNLIKELY(c != 0) {
        assert(c >= 1);
        assert(c <= ${CHANNEL_SUBTILE-1});
        __m512 vaccp0 = _mm512_load_ps(b);
        const __mmask16 vmask = _cvtu32_mask16((uint16_t) ((uint32_t) (UINT32_C(1) << c) - UINT32_C(1)));
        $for K in range(LAST_PASS_TILE):

          const __m512 vi${K}x0 = _mm512_maskz_loadu_ps(vmask, i${K});
          $if K == 0:
            __m512 vk${K}x0 = _mm512_load_ps(w);
          $else:
            __m512 vk${K}x0 = _mm512_load_ps(w + ${K * 16});
          $if 1 <= K < ACCUMULATORS:
            __m512 vaccp${K} = _mm512_mul_ps(vi${K}x0, vk${K}x0);
          $else:
            vaccp${K % ACCUMULATORS} = _mm512_fmadd_ps(vi${K}x0, vk${K}x0, vaccp${K % ACCUMULATORS});

        $if ACCUMULATORS > 1:
          // Add up all accumulators to vaccp0
          $ACC_SLICE = 1
          $while ACC_SLICE < ACCUMULATORS:
            $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
              $if A + ACC_SLICE < ACCUMULATORS:
                vaccp${A} = _mm512_add_ps(vaccp${A}, vaccp${A + ACC_SLICE});
            $ACC_SLICE *= 2

        __m512 vacc = _mm512_max_ps(vmin, vaccp0);
        vacc = _mm512_min_ps(vmax, vacc);

        _mm512_mask_storeu_ps(output, vmask, vacc);
        output += c;
      }

    }
    input = (const float**) ((uintptr_t) input + input_stride);
    output = (float*) ((uintptr_t) output + output_increment);
  } while (--output_width != 0);
}