blob: 659892ffa6391be382fd08f61872a6aed65f9ba7 [file] [log] [blame]
// Copyright 2020 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 CHANNEL_TILE % 4 == 0
$assert KERNEL_TILE >= 2
$assert ACCUMULATORS >= 1
$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
#include <assert.h>
#include <wasm_simd128.h>
#include <xnnpack/dwconv.h>
$assert ACTIVATION in ["LINEAR", "RELU", "MINMAX"]
$ACTIVATION_SUFFIX = {"LINEAR": ""}.get(ACTIVATION, "_" + ACTIVATION.lower())
$ARCH_SUFFIX = "" if ACTIVATION in ["LINEAR", "RELU"] else "_x86" if X86 else "_arm"
$PARAMS = {"LINEAR": "xnn_f32_default_params", "RELU": "xnn_f32_relu_params", "MINMAX": "xnn_f32_minmax_params"}[ACTIVATION]
void xnn_f32_dwconv${ACTIVATION_SUFFIX}_ukernel_up${CHANNEL_TILE}x${KERNEL_TILE}__wasmsimd${ARCH_SUFFIX}${"" if ACCUMULATORS == 1 else "_acc%d" % ACCUMULATORS}(
size_t channels,
size_t output_width,
const float** input,
const float* weights,
float* output,
size_t input_stride,
size_t output_increment,
size_t input_offset,
const float* zero,
const union ${PARAMS} params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS
{
assert(channels != 0);
assert(output_width != 0);
$if ACTIVATION == "MINMAX":
const v128_t vmin = wasm_v128_load64_splat(params->wasmsimd.min);
const v128_t vmax = wasm_v128_load64_splat(params->wasmsimd.max);
$elif ACTIVATION == "RELU":
const v128_t vzero = wasm_i32x4_const_splat(0);
do {
$for K in range(KERNEL_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 = (const float**) ((uintptr_t) input + input_stride);
size_t c = channels;
const float* w = weights;
for (; c >= ${CHANNEL_TILE}; c -= ${CHANNEL_TILE}) {
v128_t vacc${ABC[0:4]}p0 = wasm_v128_load(w);
$for C in range(4, CHANNEL_TILE, 4):
v128_t vacc${ABC[C:C+4]}p0 = wasm_v128_load(w + ${C});
$for K in range(KERNEL_TILE):
const v128_t vi${K}x${ABC[0:4]} = wasm_v128_load(i${K});
$for C in range(4, CHANNEL_TILE, 4):
const v128_t vi${K}x${ABC[C:C+4]} = wasm_v128_load(i${K} + ${C});
i${K} += ${CHANNEL_TILE};
$for C in range(0, CHANNEL_TILE, 4):
const v128_t vk${K}x${ABC[C:C+4]} = wasm_v128_load(w + ${(K + 1) * CHANNEL_TILE + C});
$for C in range(0, CHANNEL_TILE, 4):
$if 1 <= K < ACCUMULATORS:
v128_t vacc${ABC[C:C+4]}p${K} = wasm_f32x4_mul(vi${K}x${ABC[C:C+4]}, vk${K}x${ABC[C:C+4]});
$else:
vacc${ABC[C:C+4]}p${K % ACCUMULATORS} = wasm_f32x4_add(vacc${ABC[C:C+4]}p${K % ACCUMULATORS}, wasm_f32x4_mul(vi${K}x${ABC[C:C+4]}, vk${K}x${ABC[C:C+4]}));
w += ${(KERNEL_TILE + 1) * CHANNEL_TILE};
$if ACCUMULATORS > 1:
// Add up all accumulators to vacc${ABC[0:CHANNEL_TILE]}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, 4):
vacc${ABC[C:C+4]}p${A} = wasm_f32x4_add(vacc${ABC[C:C+4]}p${A}, vacc${ABC[C:C+4]}p${A + ACC_SLICE});
$ACC_SLICE *= 2
$if ACTIVATION == "MINMAX":
$for C in range(0, CHANNEL_TILE, 4):
$if X86:
v128_t vacc${ABC[C:C+4]} = wasm_f32x4_pmax(vmin, vacc${ABC[C:C+4]}p0);
$else:
v128_t vacc${ABC[C:C+4]} = wasm_f32x4_max(vacc${ABC[C:C+4]}p0, vmin);
$for C in range(0, CHANNEL_TILE, 4):
$if X86:
vacc${ABC[C:C+4]} = wasm_f32x4_pmin(vmax, vacc${ABC[C:C+4]});
$else:
vacc${ABC[C:C+4]} = wasm_f32x4_min(vacc${ABC[C:C+4]}, vmax);
$elif ACTIVATION == "RELU":
$for C in range(0, CHANNEL_TILE, 4):
const v128_t vacc${ABC[C:C+4]} = wasm_i32x4_max(vacc${ABC[C:C+4]}p0, vzero);
$elif ACTIVATION == "LINEAR":
$for C in range(0, CHANNEL_TILE, 4):
const v128_t vacc${ABC[C:C+4]} = vacc${ABC[C:C+4]}p0;
wasm_v128_store(output, vacc${ABC[0:4]});
$for C in range(4, CHANNEL_TILE, 4):
wasm_v128_store(output + ${C}, vacc${ABC[C:C+4]});
output += ${CHANNEL_TILE};
}
$if CHANNEL_TILE > 4:
for (; c >= 4; c -= 4) {
v128_t vacc0123p0 = wasm_v128_load(w);
$for K in range(KERNEL_TILE):
const v128_t vi${K}x0123 = wasm_v128_load(i${K});
i${K} += 4;
const v128_t vk${K}x0123 = wasm_v128_load(w + ${(K + 1) * CHANNEL_TILE});
$if 1 <= K < ACCUMULATORS:
v128_t vacc0123p${K} = wasm_f32x4_mul(vi${K}x0123, vk${K}x0123);
$else:
vacc0123p${K % ACCUMULATORS} = wasm_f32x4_add(vacc0123p${K % ACCUMULATORS}, wasm_f32x4_mul(vi${K}x0123, vk${K}x0123));
w += 4;
$if ACCUMULATORS > 1:
// Add up all accumulators to vacc${ABC[0:CHANNEL_TILE]}p0
$ACC_SLICE = 1
$while ACC_SLICE < ACCUMULATORS:
$for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
$if A + ACC_SLICE < ACCUMULATORS:
vacc0123p${A} = wasm_f32x4_add(vacc0123p${A}, vacc0123p${A + ACC_SLICE});
$ACC_SLICE *= 2
$if ACTIVATION == "MINMAX":
$if X86:
v128_t vacc0123 = wasm_f32x4_pmax(vmin, vacc0123p0);
vacc0123 = wasm_f32x4_pmin(vmax, vacc0123);
$else:
v128_t vacc0123 = wasm_f32x4_max(vacc0123p0, vmin);
vacc0123 = wasm_f32x4_min(vacc0123, vmax);
$elif ACTIVATION == "RELU":
const v128_t vacc0123 = wasm_i32x4_max(vacc0123p0, vzero);
$elif ACTIVATION == "LINEAR":
const v128_t vacc0123 = vacc0123p0;
wasm_v128_store(output, vacc0123);
output += 4;
}
if XNN_UNLIKELY(c != 0) {
v128_t vacc0123p0 = wasm_v128_load(w);
$for K in range(KERNEL_TILE):
const v128_t vi${K}x0123 = wasm_v128_load(i${K});
const v128_t vk${K}x0123 = wasm_v128_load(w + ${(K+1) * CHANNEL_TILE});
$if 1 <= K < ACCUMULATORS:
v128_t vacc0123p${K} = wasm_f32x4_mul(vi${K}x0123, vk${K}x0123);
$else:
vacc0123p${K % ACCUMULATORS} = wasm_f32x4_add(vacc0123p${K % ACCUMULATORS}, wasm_f32x4_mul(vi${K}x0123, vk${K}x0123));
$if ACCUMULATORS > 1:
// Add up all accumulators to vacc${ABC[0:CHANNEL_TILE]}p0
$ACC_SLICE = 1
$while ACC_SLICE < ACCUMULATORS:
$for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
$if A + ACC_SLICE < ACCUMULATORS:
vacc0123p${A} = wasm_f32x4_add(vacc0123p${A}, vacc0123p${A + ACC_SLICE});
$ACC_SLICE *= 2
$if ACTIVATION == "MINMAX":
$if X86:
v128_t vacc0123 = wasm_f32x4_pmax(vmin, vacc0123p0);
vacc0123 = wasm_f32x4_pmin(vmax, vacc0123);
$else:
v128_t vacc0123 = wasm_f32x4_max(vacc0123p0, vmin);
vacc0123 = wasm_f32x4_min(vacc0123, vmax);
$elif ACTIVATION == "RELU":
v128_t vacc0123 = wasm_i32x4_max(vacc0123p0, vzero);
$elif ACTIVATION == "LINEAR":
v128_t vacc0123 = vacc0123p0;
if (c & 2) {
*((double*) output) = wasm_f64x2_extract_lane(vacc0123, 0);
vacc0123 = wasm_v32x4_shuffle(vacc0123, vacc0123, 2, 3, 2, 3);
output += 2;
}
if (c & 1) {
*output = wasm_f32x4_extract_lane(vacc0123, 0);
output += 1;
}
}
output = (float*) ((uintptr_t) output + output_increment);
} while (--output_width != 0);
}