Specialize 5x5 DepthwiseConv
- Use adwconv for stride == 1
- Improve reuse of weights by only loading once for each
channel.
Change-Id: Idb865668f039450c314dbc5f3046203d9e621240
Bypass-Presubmit-Reason: Flaky test_nexus_boot_robot
diff --git a/crt/kelvin.h b/crt/kelvin.h
index a631593..10cc34f 100644
--- a/crt/kelvin.h
+++ b/crt/kelvin.h
@@ -62,4 +62,16 @@
};
static_assert(sizeof(struct vconv_u8_t) == 4);
+struct vdwconv_u8_t {
+ uint32_t mode:2; // 1:0
+ uint32_t sparsity:2; // 3:2
+ uint32_t regbase:4; // 7:4
+ uint32_t rsvd:4; // 11:8
+ int32_t sbias1:9; // 20:12
+ uint32_t sdata1:1; // 21
+ int32_t sbias2:9; // 30:22
+ uint32_t sdata2:1; // 31
+};
+static_assert(sizeof(struct vdwconv_u8_t) == 4);
+
#endif // CRT_KELVIN_H_
diff --git a/tests/kelvin_isa/vdwconv.cc b/tests/kelvin_isa/vdwconv.cc
index ec155dc..31e391f 100644
--- a/tests/kelvin_isa/vdwconv.cc
+++ b/tests/kelvin_isa/vdwconv.cc
@@ -20,18 +20,6 @@
#include "tests/kelvin_isa/kelvin_test.h"
#include "tests/kelvin_isa/vdwconv_data.h"
-struct vdwconv_u8_t {
- uint32_t mode:2; // 1:0
- uint32_t sparsity:2; // 3:2
- uint32_t regbase:4; // 7:4
- uint32_t rsvd:4; // 11:8
- int32_t sbias1:9; // 20:12
- uint32_t sdata1:1; // 21
- int32_t sbias2:9; // 30:22
- uint32_t sdata2:1; // 31
-};
-static_assert(sizeof(vdwconv_u8_t) == 4);
-
#ifdef TEST_GEN
static int32_t dwconv(const vdwconv_u8_t& cmd, uint8_t ina[3], uint8_t inb[3]) {
int32_t sbias1 = cmd.sbias1;
diff --git a/tflm/opt/depthwise_conv_s8.cc b/tflm/opt/depthwise_conv_s8.cc
index 111fdc1..9e15b3c 100644
--- a/tflm/opt/depthwise_conv_s8.cc
+++ b/tflm/opt/depthwise_conv_s8.cc
@@ -38,8 +38,475 @@
*out1 = *p_in++;
*out2 = *p_in++;
*out3 = *p_in++;
+ }
+}
+
+// special case of input depth = 32n, filter shape of 5x5, stride == 1
+void DepthwiseConvS85x5D32_Stride1(
+ const tflite::DepthwiseParams& params, const int32_t* output_multiplier,
+ const int32_t* output_shift, const tflite::RuntimeShape& input_shape,
+ const int8_t* input_data, const tflite::RuntimeShape& filter_shape,
+ const int8_t* filter_data, const tflite::RuntimeShape& bias_shape,
+ const int32_t* bias_data, const tflite::RuntimeShape& output_shape,
+ int8_t* output_data
+) {
+ const int stride_width = params.stride_width;
+ const int stride_height = params.stride_height;
+ const int pad_width = params.padding_values.width;
+ const int pad_height = params.padding_values.height;
+ const int32_t input_offset = params.input_offset;
+ const int32_t output_offset = params.output_offset;
+ const int32_t output_activation_min = params.quantized_activation_min;
+ const int32_t output_activation_max = params.quantized_activation_max;
+ const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+ const int input_height = input_shape.Dims(1);
+ const int input_width = input_shape.Dims(2);
+ const int input_depth = input_shape.Dims(3);
+ const int output_height = output_shape.Dims(1);
+ const int output_width = output_shape.Dims(2);
+ const int output_depth = output_shape.Dims(3);
+ int32_t swizzled_bias_data[32];
+ int32_t swizzled_shift_multi[32];
+ int32_t swizzled_output_multi[32];
+
+ for (int in_channel = 0; in_channel + 32 <= input_depth; in_channel += 32) {
+ const int output_channel = in_channel;
+ VectorSwizzle(bias_data + output_channel, swizzled_bias_data, 32);
+ VectorSwizzle(output_multiplier + output_channel, swizzled_output_multi, 32);
+ VectorSwizzle(output_shift + output_channel, swizzled_shift_multi, 32);
+
+ union {
+ vdwconv_u8_t dwconv;
+ uint32_t raw;
+ } cmds;
+ cmds.raw = 0;
+ cmds.dwconv.sdata1 = true;
+ cmds.dwconv.sbias1 = input_offset;
+ cmds.dwconv.sdata2 = true;
+ cmds.dwconv.sbias2 = 0;
+ cmds.dwconv.mode = 0;
+ cmds.dwconv.sparsity = 0;
+ cmds.dwconv.regbase = 0;
+
+ // Don't reorder me!
+ const int8_t* p_flt0 = filter_data + in_channel;
+ const int32_t stride = input_depth;
+ vld_b_sp_xx_m(v0, p_flt0, stride);
+ vld_b_sp_xx_m(v4, p_flt0, stride);
+ vld_b_sp_xx_m(v8, p_flt0, stride);
+ vld_b_sp_xx_m(v12, p_flt0, stride);
+ vld_b_sp_xx_m(v16, p_flt0, stride);
+ vld_b_sp_xx_m(v20, p_flt0, stride);
+ vld_b_sp_xx(v24, p_flt0, stride);
+
+ // Extra two registers to get our
+ // total usage to a multiple of 3 for dwconv.
+ vdup_b_x(v25, 0);
+ vdup_b_x(v26, 0);
+
+ for (int batch = 0; batch < batches; ++batch) {
+ const int8_t* p_output = output_data + (batch * output_height * output_width * output_depth) + output_channel;
+ for (int out_y = 0; out_y < output_height; ++out_y) {
+ const int y_offset = out_y * output_width * output_depth;
+ for (int out_x = 0; out_x < output_width; ++out_x) {
+ const int in_x_origin = (out_x * stride_width) - pad_width;
+ const int in_y_origin = (out_y * stride_height) - pad_height;
+
+ bool top_pad = in_y_origin < 0;
+ bool left_pad = in_x_origin < 0;
+ int top_pad_count = top_pad ? 0 - in_y_origin : 0;
+ int left_pad_count = left_pad ? 0 - in_x_origin : 0;
+ bool bottom_pad = (in_y_origin + 4) >= input_height;
+ bool right_pad = (in_x_origin + 4) >= input_width;
+ int bottom_pad_count = std::abs(bottom_pad ? (in_y_origin + 4) - input_height + 1: 0);
+ int right_pad_count = std::abs(right_pad ? (in_x_origin + 4) - input_width + 1 : 0);
+ bool padding_required = top_pad || left_pad || bottom_pad || right_pad;
+ assert(top_pad_count <= pad_height);
+ assert(bottom_pad_count <= pad_height);
+ assert(left_pad_count <= pad_width);
+ assert(right_pad_count <= pad_width);
+ assert(!(left_pad && right_pad));
+ const int8_t* p_in_0 = input_data +
+ (batch * input_height * input_width * input_depth) +
+ (in_y_origin * input_width * input_depth) +
+ (in_x_origin * input_depth) +
+ in_channel;
+ const int8_t* p_in_1 = p_in_0 + (input_width * input_depth);
+ const int8_t* p_in_2 = p_in_1 + (input_width * input_depth);
+ const int8_t* p_in_3 = p_in_2 + (input_width * input_depth);
+ const int8_t* p_in_4 = p_in_3 + (input_width * input_depth);
+ // Extra two registers to get our
+ // total usage to a multiple of 3 for dwconv.
+ vdup_b_x(v52, -input_offset);
+ vdup_b_x(v53, -input_offset);
+ if (!padding_required) {
+ vld_b_sp_xx(v27, p_in_0, input_depth);
+ vld_b_sp_xx_m(v28, p_in_0, input_depth);
+ vld_b_sp_xx_m(v32, p_in_1, input_depth);
+ vld_b_sp_xx(v36, p_in_1, input_depth);
+ vld_b_sp_xx(v37, p_in_2, input_depth);
+ vld_b_sp_xx(v38, p_in_2, input_depth);
+ vld_b_sp_xx(v39, p_in_2, input_depth);
+ vld_b_sp_xx(v40, p_in_2, input_depth);
+ vld_b_sp_xx(v41, p_in_2, input_depth);
+ vld_b_sp_xx(v42, p_in_3, input_depth);
+ vld_b_sp_xx(v43, p_in_3, input_depth);
+ vld_b_sp_xx(v44, p_in_3, input_depth);
+ vld_b_sp_xx(v45, p_in_3, input_depth);
+ vld_b_sp_xx(v46, p_in_3, input_depth);
+ vld_b_sp_xx(v47, p_in_4, input_depth);
+ vld_b_sp_xx_m(v48, p_in_4, input_depth);
+ } else {
+ // Top row
+ if (top_pad_count >= 1) {
+ vdup_b_x(v27, -input_offset);
+ vdup_b_x_m(v28, -input_offset);
+ } else {
+ switch (left_pad_count) {
+ case 2:
+ vdup_b_x(v28, -input_offset);
+ case 1:
+ vdup_b_x(v27, -input_offset);
+ }
+ switch (left_pad_count) {
+ case 0:
+ vld_b_x(v27, p_in_0);
+ case 1:
+ vld_b_x(v28, p_in_0 + input_depth);
+ }
+ vld_b_x(v29, p_in_0 + (2 * input_depth));
+ switch (right_pad_count) {
+ case 2:
+ vdup_b_x(v30, -input_offset);
+ case 1:
+ vdup_b_x(v31, -input_offset);
+ }
+ switch (right_pad_count) {
+ case 0:
+ vld_b_x(v31, p_in_0 + (4 * input_depth));
+ case 1:
+ vld_b_x(v30, p_in_0 + (3 * input_depth));
+ }
+ }
+
+ // 2nd row
+ if (top_pad_count == 2) {
+ vdup_b_x_m(v32, -input_offset);
+ vdup_b_x(v36, -input_offset);
+ } else {
+ switch (left_pad_count) {
+ case 2:
+ vdup_b_x(v33, -input_offset);
+ case 1:
+ vdup_b_x(v32, -input_offset);
+ }
+ switch (left_pad_count) {
+ case 0:
+ vld_b_x(v32, p_in_1);
+ case 1:
+ vld_b_x(v33, p_in_1 + input_depth);
+ }
+ vld_b_x(v34, p_in_1 + (2 * input_depth));
+ switch (right_pad_count) {
+ case 2:
+ vdup_b_x(v35, -input_offset);
+ case 1:
+ vdup_b_x(v36, -input_offset);
+ }
+ switch (right_pad_count) {
+ case 0:
+ vld_b_x(v36, p_in_1 + (4 * input_depth));
+ case 1:
+ vld_b_x(v35, p_in_1 + (3 * input_depth));
+ }
+ }
+
+ // 3rd row
+ switch (left_pad_count) {
+ case 2:
+ vdup_b_x(v38, -input_offset);
+ case 1:
+ vdup_b_x(v37, -input_offset);
+ }
+ switch (left_pad_count) {
+ case 0:
+ vld_b_x(v37, p_in_2);
+ case 1:
+ vld_b_x(v38, p_in_2 + input_depth);
+ }
+ vld_b_x(v39, p_in_2 + (2 * input_depth));
+ switch (right_pad_count) {
+ case 2:
+ vdup_b_x(v40, -input_offset);
+ case 1:
+ vdup_b_x(v41, -input_offset);
+ }
+ switch (right_pad_count) {
+ case 0:
+ vld_b_x(v41, p_in_2 + (4 * input_depth));
+ case 1:
+ vld_b_x(v40, p_in_2 + (3 * input_depth));
+ }
+
+ // 4th row
+ if (bottom_pad_count == 2) {
+ vdup_b_x(v42, -input_offset);
+ vdup_b_x(v43, -input_offset);
+ vdup_b_x(v44, -input_offset);
+ vdup_b_x(v45, -input_offset);
+ vdup_b_x(v46, -input_offset);
+ } else {
+ switch (left_pad_count) {
+ case 2:
+ vdup_b_x(v43, -input_offset);
+ case 1:
+ vdup_b_x(v42, -input_offset);
+ }
+ switch (left_pad_count) {
+ case 0:
+ vld_b_x(v42, p_in_3);
+ case 1:
+ vld_b_x(v43, p_in_3 + input_depth);
+ }
+ switch (right_pad_count) {
+ case 2:
+ vdup_b_x(v45, -input_offset);
+ case 1:
+ vdup_b_x(v46, -input_offset);
+ }
+ vld_b_x(v44, p_in_3 + (2 * input_depth));
+ switch (right_pad_count) {
+ case 0:
+ vld_b_x(v46, p_in_3 + (4 * input_depth));
+ case 1:
+ vld_b_x(v45, p_in_3 + (3 * input_depth));
+ }
+ }
+
+ // 5th row
+ if (bottom_pad_count >= 1) {
+ vdup_b_x(v47, -input_offset);
+ vdup_b_x(v48, -input_offset);
+ vdup_b_x(v49, -input_offset);
+ vdup_b_x(v50, -input_offset);
+ vdup_b_x(v51, -input_offset);
+ } else {
+ switch (left_pad_count) {
+ case 2:
+ vdup_b_x(v48, -input_offset);
+ case 1:
+ vdup_b_x(v47, -input_offset);
+ }
+ switch (left_pad_count) {
+ case 0:
+ vld_b_x(v47, p_in_4);
+ case 1:
+ vld_b_x(v48, p_in_4 + input_depth);
+ }
+ vld_b_x(v49, p_in_4 + (2 * input_depth));
+ switch (right_pad_count) {
+ case 2:
+ vdup_b_x(v50, -input_offset);
+ case 1:
+ vdup_b_x(v51, -input_offset);
+ }
+ switch (right_pad_count) {
+ case 0:
+ vld_b_x(v51, p_in_4 + (4 * input_depth));
+ case 1:
+ vld_b_x(v50, p_in_4 + (3 * input_depth));
+ }
+ }
+ }
+
+ vld_w_x_m(v60, swizzled_bias_data);
+ adwinit_v(v60, v60);
+ adwconv_vxv(v60, v27, cmds, v0);
+ adwconv_vxv(v60, v30, cmds, v3);
+ adwconv_vxv(v60, v33, cmds, v6);
+ adwconv_vxv(v60, v36, cmds, v9);
+ adwconv_vxv(v60, v39, cmds, v12);
+ adwconv_vxv(v60, v42, cmds, v15);
+ adwconv_vxv(v60, v45, cmds, v18);
+ adwconv_vxv(v60, v48, cmds, v21);
+ vdwconv_vxv(v60, v51, cmds, v24);
+
+ vld_w_x_m(v56, swizzled_output_multi);
+ vdmulh_w_rn_vv_m(v60, v60, v56);
+ vld_w_x_m(v56, swizzled_shift_multi);
+ vrsub_w_vx_m(v56, v56, 0);
+ vsha_w_r_vv_m(v60, v60, v56);
+ vadd_w_vx_m(v60, v60, output_offset);
+ vmax_w_vx_m(v60, v60, output_activation_min);
+ vmin_w_vx_m(v60, v60, output_activation_max);
+ vsraqs_b_vx(v60, v60, 0);
+ vst_b_x(v60, p_output + y_offset + (out_x * output_depth));
+ }
+ }
}
}
+}
+
+// special case of input depth = 32n, filter shape of 5x5
+void DepthwiseConvS85x5D32(
+ const tflite::DepthwiseParams& params, const int32_t* output_multiplier,
+ const int32_t* output_shift, const tflite::RuntimeShape& input_shape,
+ const int8_t* input_data, const tflite::RuntimeShape& filter_shape,
+ const int8_t* filter_data, const tflite::RuntimeShape& bias_shape,
+ const int32_t* bias_data, const tflite::RuntimeShape& output_shape,
+ int8_t* output_data
+) {
+ const int stride_width = params.stride_width;
+ const int stride_height = params.stride_height;
+ const int pad_width = params.padding_values.width;
+ const int pad_height = params.padding_values.height;
+ const int32_t input_offset = params.input_offset;
+ const int32_t output_offset = params.output_offset;
+ const int32_t output_activation_min = params.quantized_activation_min;
+ const int32_t output_activation_max = params.quantized_activation_max;
+ const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+ const int input_height = input_shape.Dims(1);
+ const int input_width = input_shape.Dims(2);
+ const int input_depth = input_shape.Dims(3);
+ const int filter_height = filter_shape.Dims(1);
+ const int filter_width = filter_shape.Dims(2);
+ const int output_height = output_shape.Dims(1);
+ const int output_width = output_shape.Dims(2);
+ const int output_depth = output_shape.Dims(3);
+ int32_t swizzled_bias_data[32];
+ int32_t swizzled_shift_multi[32];
+ int32_t swizzled_output_multi[32];
+
+ for (int in_channel = 0; in_channel + 32 <= input_depth; in_channel += 32) {
+ const int output_channel = in_channel;
+ VectorSwizzle(bias_data + output_channel, swizzled_bias_data, 32);
+ VectorSwizzle(output_multiplier + output_channel, swizzled_output_multi, 32);
+ VectorSwizzle(output_shift + output_channel, swizzled_shift_multi, 32);
+
+ vld_w_x_m(v52, swizzled_bias_data);
+ vld_w_x_m(v56, swizzled_output_multi);
+ vld_w_x_m(v60, swizzled_shift_multi);
+ vrsub_w_vx_m(v60, v60, 0);
+
+ // Don't reorder me!
+ const int8_t* p_flt = filter_data + in_channel;
+ vld_b_sp_xx(v6, p_flt, input_depth);
+ vld_b_sp_xx(v7, p_flt, input_depth);
+ vld_b_sp_xx_m(v8, p_flt, input_depth);
+ vld_b_sp_xx_m(v12, p_flt, input_depth);
+ vld_b_sp_xx_m(v16, p_flt, input_depth);
+ vld_b_sp_xx_m(v20, p_flt, input_depth);
+ vld_b_sp_xx_m(v24, p_flt, input_depth);
+ vld_b_sp_xx(v28, p_flt, input_depth);
+ vld_b_sp_xx(v29, p_flt, input_depth);
+ vld_b_sp_xx(v30, p_flt, input_depth);
+
+
+ for (int batch = 0; batch < batches; ++batch) {
+ const int8_t* p_input = input_data + (batch * input_width * input_height * input_depth) + in_channel;
+ const int8_t* p_output = output_data + (batch * output_width * output_height * output_depth) + output_channel;
+ for (int out_y = 0; out_y < output_height; ++out_y) {
+ const int out_y_offset = (out_y * output_width * output_depth);
+ for (int out_x = 0; out_x < output_width; ++out_x) {
+ const int in_x_origin = (out_x * stride_width) - pad_width;
+ const int in_y_origin = (out_y * stride_height) - pad_height;
+
+ // Initialize accumulators w/ bias_data
+ vmv_v_m(v48, v52);
+
+ for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
+ const int in_y = in_y_origin + filter_y;
+ if ((in_y < 0) || (in_y >= input_height)) {
+ continue;
+ }
+ switch (filter_y) {
+ case 0:
+ vaddw_h_vx(v31, v6, 0);
+ vaddw_h_vx(v33, v7, 0);
+ vaddw_h_vx(v35, v8, 0);
+ vaddw_h_vx(v37, v9, 0);
+ vaddw_h_vx(v39, v10, 0);
+ break;
+ case 1:
+ vaddw_h_vx(v31, v11, 0);
+ vaddw_h_vx(v33, v12, 0);
+ vaddw_h_vx(v35, v13, 0);
+ vaddw_h_vx(v37, v14, 0);
+ vaddw_h_vx(v39, v15, 0);
+ break;
+ case 2:
+ vaddw_h_vx(v31, v16, 0);
+ vaddw_h_vx(v33, v17, 0);
+ vaddw_h_vx(v35, v18, 0);
+ vaddw_h_vx(v37, v19, 0);
+ vaddw_h_vx(v39, v20, 0);
+ break;
+ case 3:
+ vaddw_h_vx(v31, v21, 0);
+ vaddw_h_vx(v33, v22, 0);
+ vaddw_h_vx(v35, v23, 0);
+ vaddw_h_vx(v37, v24, 0);
+ vaddw_h_vx(v39, v25, 0);
+ break;
+ case 4:
+ vaddw_h_vx(v31, v26, 0);
+ vaddw_h_vx(v33, v27, 0);
+ vaddw_h_vx(v35, v28, 0);
+ vaddw_h_vx(v37, v29, 0);
+ vaddw_h_vx(v39, v30, 0);
+ break;
+ }
+ const int in_y_offset = in_y * input_width * input_depth;
+ for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
+ const int in_x = in_x_origin + filter_x;
+ if ((in_x < 0) || (in_x >= input_width)) {
+ continue;
+ }
+
+ vld_b_x(v0, p_input + (in_x * input_depth) + in_y_offset);
+
+ vaddw_h_vx(v0, v0, 0);
+ vadd_h_vx(v0, v0, static_cast<int16_t>(input_offset));
+ vadd_h_vx(v1, v1,
+ static_cast<int16_t>(input_offset)); // v0 v1 input
+ switch (filter_x) {
+ case 0:
+ vmulw_w_vv(v2, v1, v32);
+ vmulw_w_vv(v0, v0, v31);
+ break;
+ case 1:
+ vmulw_w_vv(v2, v1, v34);
+ vmulw_w_vv(v0, v0, v33);
+ break;
+ case 2:
+ vmulw_w_vv(v2, v1, v36);
+ vmulw_w_vv(v0, v0, v35);
+ break;
+ case 3:
+ vmulw_w_vv(v2, v1, v38);
+ vmulw_w_vv(v0, v0, v37);
+ break;
+ case 4:
+ vmulw_w_vv(v2, v1, v40);
+ vmulw_w_vv(v0, v0, v39);
+ break;
+ }
+ vadd_w_vv_m(v48, v48, v0);
+ }
+ }
+
+ vdmulh_w_rn_vv_m(v48, v48, v56);
+ vsha_w_r_vv_m(v48, v48, v60);
+ vadd_w_vx_m(v48, v48, output_offset);
+ vmax_w_vx_m(v48, v48, output_activation_min);
+ vmin_w_vx_m(v48, v48, output_activation_max);
+ vsraqs_b_vx(v48, v48, 0);
+ vst_b_x(v48, p_output + out_y_offset + (out_x * output_depth));
+ }
+ }
+ }
+ }
+}
// special case of input depth = 32n
void DepthwiseConvS8D32(
@@ -161,6 +628,8 @@
// TODO(b/141565753): Re-introduce ScopedProfilingLabel on Micro.
const int stride_width = params.stride_width;
const int stride_height = params.stride_height;
+ const int filter_height = filter_shape.Dims(1);
+ const int filter_width = filter_shape.Dims(2);
const int dilation_width_factor = params.dilation_width_factor;
const int dilation_height_factor = params.dilation_height_factor;
const int depth_multiplier = params.depth_multiplier;
@@ -186,7 +655,15 @@
// special case of output depth = 32n
if (output_depth % 32 == 0) {
- fn = DepthwiseConvS8D32;
+ if (filter_width == 5 && filter_height == 5) {
+ if (stride_width <= 1 && stride_height <= 1) {
+ fn = DepthwiseConvS85x5D32_Stride1;
+ } else {
+ fn = DepthwiseConvS85x5D32;
+ }
+ } else {
+ fn = DepthwiseConvS8D32;
+ }
}
fn(params, output_multiplier, output_shift, input_shape, input_data,
