// 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 = 8 $assert CHANNEL_TILE % CHANNEL_SUBTILE == 0 $CHANNEL_ROUND = 4 $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 #include #include #include #include #include $ISA = {0: "avx", 3: "fma3"}[FMA] 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__${ISA}${"" 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)]) XNN_OOB_READS { assert(channels != 0); assert(output_width != 0); assert(kernel_size > ${FIRST_PASS_TILE}); const __m256 vmin = _mm256_load_ps(params->avx.min); const __m256 vmax = _mm256_load_ps(params->avx.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 = round_up_po2(channels, ${CHANNEL_ROUND}); $if CHANNEL_TILE > 8: for (; c >= ${CHANNEL_TILE}; c -= ${CHANNEL_TILE}) { $for C in range(0, CHANNEL_TILE, 8): $if C == 0: __m256 vacc${ABC[C:C+8]}p0 = _mm256_load_ps(w); $else: __m256 vacc${ABC[C:C+8]}p0 = _mm256_load_ps(w + ${C}); $for K in range(FIRST_PASS_TILE): $for C in range(0, CHANNEL_TILE, 8): $if C == 0: const __m256 vi${K}x${ABC[C:C+8]} = _mm256_loadu_ps(i${K}); $else: const __m256 vi${K}x${ABC[C:C+8]} = _mm256_loadu_ps(i${K} + ${C}); i${K} += ${CHANNEL_TILE}; $for C in range(0, CHANNEL_TILE, 8): const __m256 vk${K}x${ABC[C:C+8]} = _mm256_load_ps(w + ${(K + 1) * CHANNEL_TILE + C}); $for C in range(0, CHANNEL_TILE, 8): $if 1 <= K < ACCUMULATORS: __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]}); $elif FMA == 3: 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}); $else: 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]})); w += ${(FIRST_PASS_TILE + 1) * CHANNEL_TILE}; $if ACCUMULATORS > 1: // Add up all accumulators to vacc${ABC[0:8]}p0 $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, 8): 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}); $ACC_SLICE *= 2 $for C in range(0, CHANNEL_TILE, 8): $if C == 0: _mm256_store_ps(b, vacc${ABC[C:C+8]}p0); $else: _mm256_store_ps(b + ${C}, vacc${ABC[C:C+8]}p0); b += ${CHANNEL_TILE}; } for (; c >= ${CHANNEL_SUBTILE}; c -= ${CHANNEL_SUBTILE}) { __m256 vacc01234567p0 = _mm256_load_ps(w); $for K in range(FIRST_PASS_TILE): const __m256 vi${K}x01234567 = _mm256_loadu_ps(i${K}); i${K} += ${CHANNEL_SUBTILE}; const __m256 vk${K}x01234567 = _mm256_load_ps(w + ${(K + 1) * CHANNEL_SUBTILE}); $if 1 <= K < ACCUMULATORS: __m256 vacc01234567p${K} = _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567); $elif FMA == 3: vacc01234567p${K % ACCUMULATORS} = _mm256_fmadd_ps(vi${K}x01234567, vk${K}x01234567, vacc01234567p${K % ACCUMULATORS}); $else: vacc01234567p${K % ACCUMULATORS} = _mm256_add_ps(vacc01234567p${K % ACCUMULATORS}, _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567)); w += ${(FIRST_PASS_TILE + 1) * CHANNEL_SUBTILE}; $if ACCUMULATORS > 1: // Add up all accumulators to vacc01234567p0 $ACC_SLICE = 1 $while ACC_SLICE < ACCUMULATORS: $for A in range(0, ACCUMULATORS, ACC_SLICE * 2): $if A + ACC_SLICE < ACCUMULATORS: vacc01234567p${A} = _mm256_add_ps(vacc01234567p${A}, vacc01234567p${A + ACC_SLICE}); $ACC_SLICE *= 2 _mm256_store_ps(b, vacc01234567p0); b += ${CHANNEL_SUBTILE}; } if (c != 0) { assert(c >= 1); assert(c <= ${CHANNEL_SUBTILE-1}); const __m256i vmask = _mm256_loadu_si256((const __m256i*) ¶ms->avx.mask_table[7 - c]); __m256 vacc01234567p0 = _mm256_load_ps(w); $for K in range(FIRST_PASS_TILE): const __m256 vi${K}x01234567 = _mm256_maskload_ps(i${K}, vmask); const __m256 vk${K}x01234567 = _mm256_load_ps(w + ${(K + 1) * CHANNEL_SUBTILE}); $if 1 <= K < ACCUMULATORS: __m256 vacc01234567p${K} = _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567); $elif FMA == 3: vacc01234567p${K % ACCUMULATORS} = _mm256_fmadd_ps(vi${K}x01234567, vk${K}x01234567, vacc01234567p${K % ACCUMULATORS}); $else: vacc01234567p${K % ACCUMULATORS} = _mm256_add_ps(vacc01234567p${K % ACCUMULATORS}, _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567)); w += ${(FIRST_PASS_TILE + 1) * CHANNEL_SUBTILE}; $if ACCUMULATORS > 1: // Add up all accumulators to vacc01234567p0 $ACC_SLICE = 1 $while ACC_SLICE < ACCUMULATORS: $for A in range(0, ACCUMULATORS, ACC_SLICE * 2): $if A + ACC_SLICE < ACCUMULATORS: vacc01234567p${A} = _mm256_add_ps(vacc01234567p${A}, vacc01234567p${A + ACC_SLICE}); $ACC_SLICE *= 2 _mm256_store_ps(b, vacc01234567p0); } } // 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 = round_up_po2(channels, ${CHANNEL_ROUND}); $if CHANNEL_TILE > 8: for (; c >= ${CHANNEL_TILE}; c -= ${CHANNEL_TILE}) { $for C in range(0, CHANNEL_TILE, 8): $if C == 0: __m256 vacc${ABC[C:C+8]}p0 = _mm256_load_ps(b); $else: __m256 vacc${ABC[C:C+8]}p0 = _mm256_load_ps(b + ${C}); $for K in range(MIDDLE_PASS_TILE): $for C in range(0, CHANNEL_TILE, 8): $if C == 0: const __m256 vi${K}x${ABC[C:C+8]} = _mm256_loadu_ps(i${K}); $else: const __m256 vi${K}x${ABC[C:C+8]} = _mm256_loadu_ps(i${K} + ${C}); i${K} += ${CHANNEL_TILE}; $for C in range(0, CHANNEL_TILE, 8): $if K == 0 and C == 0: const __m256 vk${K}x${ABC[C:C+8]} = _mm256_load_ps(w); $else: const __m256 vk${K}x${ABC[C:C+8]} = _mm256_load_ps(w + ${K * CHANNEL_TILE + C}); $for C in range(0, CHANNEL_TILE, 8): $if 1 <= K < ACCUMULATORS: __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]}); $elif FMA == 3: 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}); $else: 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]})); w += ${MIDDLE_PASS_TILE * CHANNEL_TILE}; $if ACCUMULATORS > 1: // Add up all accumulators to vacc${ABC[0:8]}p0 $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, 8): 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}); $ACC_SLICE *= 2 $for C in range(0, CHANNEL_TILE, 8): $if C == 0: _mm256_store_ps(b, vacc${ABC[C:C+8]}p0); $else: _mm256_store_ps(b + ${C}, vacc${ABC[C:C+8]}p0); b += ${CHANNEL_TILE}; } for (; c >= 8; c -= 8) { __m256 vacc01234567p0 = _mm256_load_ps(b); $for K in range(MIDDLE_PASS_TILE): const __m256 vi${K}x01234567 = _mm256_loadu_ps(i${K}); i${K} += ${CHANNEL_SUBTILE}; $if K == 0: const __m256 vk${K}x01234567 = _mm256_load_ps(w); $else: const __m256 vk${K}x01234567 = _mm256_load_ps(w + ${K * CHANNEL_SUBTILE}); $if 1 <= K < ACCUMULATORS: __m256 vacc01234567p${K} = _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567); $elif FMA == 3: vacc01234567p${K % ACCUMULATORS} = _mm256_fmadd_ps(vi${K}x01234567, vk${K}x01234567, vacc01234567p${K % ACCUMULATORS}); $else: vacc01234567p${K % ACCUMULATORS} = _mm256_add_ps(vacc01234567p${K % ACCUMULATORS}, _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567)); w += ${MIDDLE_PASS_TILE * CHANNEL_SUBTILE}; $if ACCUMULATORS > 1: // Add up all accumulators to vacc01234567p0 $ACC_SLICE = 1 $while ACC_SLICE < ACCUMULATORS: $for A in range(0, ACCUMULATORS, ACC_SLICE * 2): $if A + ACC_SLICE < ACCUMULATORS: vacc01234567p${A} = _mm256_add_ps(vacc01234567p${A}, vacc01234567p${A + ACC_SLICE}); $ACC_SLICE *= 2 _mm256_store_ps(b, vacc01234567p0); b += ${CHANNEL_SUBTILE}; } if (c != 0) { assert(c >= 1); assert(c <= ${CHANNEL_SUBTILE-1}); const __m256i vmask = _mm256_loadu_si256((const __m256i*) ¶ms->avx.mask_table[7 - c]); __m256 vacc01234567p0 = _mm256_load_ps(b); $for K in range(MIDDLE_PASS_TILE): const __m256 vi${K}x01234567 = _mm256_maskload_ps(i${K}, vmask); $if K == 0: const __m256 vk${K}x01234567 = _mm256_load_ps(w); $else: const __m256 vk${K}x01234567 = _mm256_load_ps(w + ${(K) * CHANNEL_SUBTILE}); $if 1 <= K < ACCUMULATORS: __m256 vacc01234567p${K} = _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567); $elif FMA == 3: vacc01234567p${K % ACCUMULATORS} = _mm256_fmadd_ps(vi${K}x01234567, vk${K}x01234567, vacc01234567p${K % ACCUMULATORS}); $else: vacc01234567p${K % ACCUMULATORS} = _mm256_add_ps(vacc01234567p${K % ACCUMULATORS}, _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567)); w += ${(MIDDLE_PASS_TILE) * CHANNEL_SUBTILE}; $if ACCUMULATORS > 1: // Add up all accumulators to vacc01234567p0 $ACC_SLICE = 1 $while ACC_SLICE < ACCUMULATORS: $for A in range(0, ACCUMULATORS, ACC_SLICE * 2): $if A + ACC_SLICE < ACCUMULATORS: vacc01234567p${A} = _mm256_add_ps(vacc01234567p${A}, vacc01234567p${A + ACC_SLICE}); $ACC_SLICE *= 2 _mm256_store_ps(b, vacc01234567p0); } } // 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 > 8: for (; c >= ${CHANNEL_TILE}; c -= ${CHANNEL_TILE}) { $for C in range(0, CHANNEL_TILE, 8): $if C == 0: __m256 vacc${ABC[C:C+8]}p0 = _mm256_load_ps(b); $else: __m256 vacc${ABC[C:C+8]}p0 = _mm256_load_ps(b + ${C}); b += ${CHANNEL_TILE}; $for K in range(LAST_PASS_TILE): $for C in range(0, CHANNEL_TILE, 8): $if C == 0: const __m256 vi${K}x${ABC[C:C+8]} = _mm256_loadu_ps(i${K}); $else: const __m256 vi${K}x${ABC[C:C+8]} = _mm256_loadu_ps(i${K} + ${C}); i${K} += ${CHANNEL_TILE}; $for C in range(0, CHANNEL_TILE, 8): $if K == 0 and C == 0: __m256 vk${K}x${ABC[C:C+8]} = _mm256_load_ps(w); $else: __m256 vk${K}x${ABC[C:C+8]} = _mm256_load_ps(w + ${K * CHANNEL_TILE + C}); $for C in range(0, CHANNEL_TILE, 8): $if 1 <= K < ACCUMULATORS: __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]}); $elif FMA == 3: 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}); $else: 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]})); w += ${LAST_PASS_TILE * CHANNEL_TILE}; $if ACCUMULATORS > 1: // Add up all accumulators to vacc${ABC[0:8]}p0 $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, 8): 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}); $ACC_SLICE *= 2 $for C in range(0, CHANNEL_TILE, 8): __m256 vacc${ABC[C:C+8]} = _mm256_max_ps(vmin, vacc${ABC[C:C+8]}p0); $for C in range(0, CHANNEL_TILE, 8): vacc${ABC[C:C+8]} = _mm256_min_ps(vmax, vacc${ABC[C:C+8]}); $for C in range(0, CHANNEL_TILE, 8): $if C == 0: _mm256_storeu_ps(output, vacc${ABC[C:C+8]}); $else: _mm256_storeu_ps(output + ${C}, vacc${ABC[C:C+8]}); output += ${CHANNEL_TILE}; } for (; c >= 8; c -= 8) { __m256 vacc01234567p0 = _mm256_load_ps(b); b += 8; $for K in range(LAST_PASS_TILE): const __m256 vi${K}x01234567 = _mm256_loadu_ps(i${K}); i${K} += 8; $if K == 0: __m256 vk${K}x01234567 = _mm256_load_ps(w); $else: __m256 vk${K}x01234567 = _mm256_load_ps(w + ${K * 8}); $if 1 <= K < ACCUMULATORS: __m256 vacc01234567p${K} = _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567); $elif FMA == 3: vacc01234567p${K % ACCUMULATORS} = _mm256_fmadd_ps(vi${K}x01234567, vk${K}x01234567, vacc01234567p${K % ACCUMULATORS}); $else: vacc01234567p${K % ACCUMULATORS} = _mm256_add_ps(vacc01234567p${K % ACCUMULATORS}, _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567)); $if CHANNEL_TILE > 8: w += ${LAST_PASS_TILE * 8}; $else: w += ${LAST_PASS_TILE * CHANNEL_TILE}; $if ACCUMULATORS > 1: // Add up all accumulators to vacc01234567p0 $ACC_SLICE = 1 $while ACC_SLICE < ACCUMULATORS: $for A in range(0, ACCUMULATORS, ACC_SLICE * 2): $if A + ACC_SLICE < ACCUMULATORS: vacc01234567p${A} = _mm256_add_ps(vacc01234567p${A}, vacc01234567p${A + ACC_SLICE}); $ACC_SLICE *= 2 __m256 vacc01234567 = _mm256_max_ps(vmin, vacc01234567p0); vacc01234567 = _mm256_min_ps(vmax, vacc01234567); _mm256_storeu_ps(output, vacc01234567); output += 8; } if XNN_UNLIKELY(c != 0) { assert(c >= 1); assert(c <= ${CHANNEL_SUBTILE-1}); __m256 vacc01234567p0 = _mm256_load_ps(b); const __m256i vmask = _mm256_loadu_si256((const __m256i*) ¶ms->avx.mask_table[7 - c]); $for K in range(LAST_PASS_TILE): const __m256 vi${K}x01234567 = _mm256_maskload_ps(i${K}, vmask); $if K == 0: __m256 vk${K}x01234567 = _mm256_load_ps(w); $else: __m256 vk${K}x01234567 = _mm256_load_ps(w + ${K * 8}); $if 1 <= K < ACCUMULATORS: __m256 vacc01234567p${K} = _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567); $elif FMA == 3: vacc01234567p${K % ACCUMULATORS} = _mm256_fmadd_ps(vi${K}x01234567, vk${K}x01234567, vacc01234567p${K % ACCUMULATORS}); $else: vacc01234567p${K % ACCUMULATORS} = _mm256_add_ps(vacc01234567p${K % ACCUMULATORS}, _mm256_mul_ps(vi${K}x01234567, vk${K}x01234567)); $if ACCUMULATORS > 1: // Add up all accumulators to vacc${ABC[0:8]}p0 $ACC_SLICE = 1 $while ACC_SLICE < ACCUMULATORS: $for A in range(0, ACCUMULATORS, ACC_SLICE * 2): $if A + ACC_SLICE < ACCUMULATORS: vacc01234567p${A} = _mm256_add_ps(vacc01234567p${A}, vacc01234567p${A + ACC_SLICE}); $ACC_SLICE *= 2 __m256 vacc01234567 = _mm256_max_ps(vmin, vacc01234567p0); vacc01234567 = _mm256_min_ps(vmax, vacc01234567); __m128 vacc0123 = _mm256_castps256_ps128(vacc01234567); if (c & 4) { _mm_storeu_ps(output, vacc0123); vacc0123 = _mm256_extractf128_ps(vacc01234567, 1); output += 4; } if (c & 2) { _mm_storel_pi((__m64*) output, vacc0123); vacc0123 = _mm_movehl_ps(vacc0123, vacc0123); output += 2; } if (c & 1) { _mm_store_ss(output, vacc0123); output += 1; } } } input = (const float**) ((uintptr_t) input + input_stride); output = (float*) ((uintptr_t) output + output_increment); } while (--output_width != 0); }