Improve ConvS8D4
- Remove tflite::Offset
- Lift some invariants out of loops
- Loop unroll pragmas
- Remove conditional branching from output stage
Change-Id: Iabf35537ef3cdba3da4375d8b7473fecd1fa38f2
diff --git a/tflm/opt/conv_s8_d4.cc b/tflm/opt/conv_s8_d4.cc
index 18df3e7..54c20bd 100644
--- a/tflm/opt/conv_s8_d4.cc
+++ b/tflm/opt/conv_s8_d4.cc
@@ -27,6 +27,9 @@
#include "tensorflow/lite/kernels/internal/runtime_shape.h"
#include "tensorflow/lite/kernels/internal/types.h"
+#define unlikely(x) (__builtin_expect(false || (x), false))
+#define likely(x) (__builtin_expect(false || (x), true))
+
namespace kelvin::opt {
namespace {
@@ -159,8 +162,12 @@
vrsub_w_vx_m(v24, v24, 0);
for (int batch = 0; batch < batches; ++batch) {
+ const int8_t* p_output =
+ output_data + (batch * output_height * output_width * output_depth) +
+ out_channel;
for (int out_y = 0; out_y < output_height; ++out_y) {
const int in_y_origin = (out_y * stride_height) - pad_height;
+ const int out_y_offset = (out_y * output_width * output_depth);
int out_x = 0;
do {
int out_xs_this_iter = std::min(8, output_width - out_x);
@@ -171,55 +178,59 @@
int in_channel = 0;
do {
int in_channels_this_iter = std::min(filter_input_depth, 32);
- for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
- const int in_y = in_y_origin + dilation_height_factor * filter_y;
- const bool is_row_inside_input =
- (in_y >= 0) && (in_y < input_height);
- if (!is_row_inside_input) {
- continue;
+ // Calculate first valid filter_y
+ int filter_y = 0;
+ {
+ int in_y = in_y_origin;
+ while (in_y < 0) {
+ ++filter_y;
+ in_y += (dilation_height_factor);
}
+ }
+ for (; filter_y < filter_height; ++filter_y) {
+ const int y_filter_offset =
+ (filter_y * filter_width * 8 * input_depth);
+ const int in_y = in_y_origin + dilation_height_factor * filter_y;
+ if (in_y >= input_height) {
+ break;
+ }
+ const int8_t* p_in =
+ input_data + in_channel + (in_y * input_width * input_depth) +
+ (batch * input_height * input_width * input_depth);
+ int in_x[8];
+#pragma GCC unroll 8
+ for (int i = 0; i < 8; ++i) {
+ in_x[i] = ((out_x + i) * stride_width) - pad_width;
+ }
for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
- int in_x[8];
- bool right_pad = false;
+ const int8_t* p_in_x[8];
int first_right_pad = -1;
+
+#pragma GCC unroll 8
for (int i = 0; i < 8; ++i) {
- const int in_x_origin =
- ((out_x + i) * stride_width) - pad_width;
- in_x[i] = in_x_origin + dilation_width_factor * filter_x;
+ p_in_x[i] = p_in + (in_x[i] * input_depth);
}
- bool left_pad = (in_x[0] < 0);
+
+#pragma GCC unroll 8
for (int i = 7; i >= 0; --i) {
if (in_x[i] < input_width) {
break;
}
- right_pad = true;
first_right_pad = i;
}
+ bool left_pad = (in_x[0] < 0);
+ bool right_pad = (first_right_pad != -1);
- if (left_pad) {
+ int stride = input_depth * stride_width;
+
+ if (unlikely(left_pad)) {
vdup_b_x(v0, -input_offset);
- vld_b_s_xx(
- v1,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[1], in_channel)],
- input_depth * stride_width);
- vld_b_s_xx(
- v2,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[2], in_channel)],
- input_depth * stride_width);
- vld_b_s_xx(
- v3,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[3], in_channel)],
- input_depth * stride_width);
- vld_b_s_xx_m(
- v4,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[4], in_channel)],
- input_depth * stride_width);
- } else if (right_pad) {
+ vld_b_s_xx(v1, p_in_x[1], stride);
+ vld_b_s_xx(v2, p_in_x[2], stride);
+ vld_b_s_xx(v3, p_in_x[3], stride);
+ vld_b_s_xx_m(v4, p_in_x[4], stride);
+ } else if (unlikely(right_pad)) {
int first_pad = std::min(first_right_pad, out_xs_this_iter);
switch (first_pad) {
case 0:
@@ -241,88 +252,36 @@
}
switch (8 - first_pad) { // rest (stripmines?)
case 0:
- vld_b_s_xx(
- v7,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[7], in_channel)],
- input_depth * stride_width);
+ vld_b_s_xx(v7, p_in_x[7], stride);
case 1:
- vld_b_s_xx(
- v6,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[6], in_channel)],
- input_depth * stride_width);
+ vld_b_s_xx(v6, p_in_x[6], stride);
case 2:
- vld_b_s_xx(
- v5,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[5], in_channel)],
- input_depth * stride_width);
+ vld_b_s_xx(v5, p_in_x[5], stride);
case 3:
- vld_b_s_xx(
- v4,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[4], in_channel)],
- input_depth * stride_width);
+ vld_b_s_xx(v4, p_in_x[4], stride);
case 4:
- vld_b_s_xx(
- v3,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[3], in_channel)],
- input_depth * stride_width);
+ vld_b_s_xx(v3, p_in_x[3], stride);
case 5:
- vld_b_s_xx(
- v2,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[2], in_channel)],
- input_depth * stride_width);
+ vld_b_s_xx(v2, p_in_x[2], stride);
case 6:
- vld_b_s_xx(
- v1,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[1], in_channel)],
- input_depth * stride_width);
+ vld_b_s_xx(v1, p_in_x[1], stride);
case 7:
- vld_b_s_xx(
- v0,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[0], in_channel)],
- input_depth * stride_width);
+ vld_b_s_xx(v0, p_in_x[0], stride);
}
- } else if (!left_pad && !right_pad) {
+ } else if (likely(!left_pad && !right_pad)) {
// Inputs
- vld_b_s_xx_m(
- v0,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[0], in_channel)],
- input_depth * stride_width);
- vld_b_s_xx_m(
- v4,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[4], in_channel)],
- input_depth * stride_width);
+ vld_b_s_xx_m(v0, p_in_x[0], stride);
+ vld_b_s_xx_m(v4, p_in_x[4], stride);
} else {
- vdup_b_x(v0, -input_offset);
- vdup_b_x(v7, -input_offset);
- vld_b_s_xx_m(
- v1,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[1], in_channel)],
- input_depth * stride_width);
- vld_b_s_xx(
- v5,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[5], in_channel)],
- input_depth * stride_width);
- vld_b_s_xx(
- v6,
- &input_data[tflite::Offset(input_shape, batch, in_y,
- in_x[6], in_channel)],
- input_depth * stride_width);
+ vdup_b_x(v0, neg_input_offset);
+ vdup_b_x(v7, neg_input_offset);
+ vld_b_s_xx_m(v1, p_in_x[1], stride);
+ vld_b_s_xx(v5, p_in_x[5], stride);
+ vld_b_s_xx(v6, p_in_x[6], stride);
}
- size_t local_filter_offset =
- (filter_y * filter_width * 8 * input_depth) +
- (filter_x * 8 * input_depth) + (in_channel * 8);
+ size_t local_filter_offset = y_filter_offset +
+ (filter_x * 8 * input_depth) +
+ (in_channel * 8);
int8_t* p_local_filter_start =
p_swizzled_filter_data + local_filter_offset;
vld_b_p_x_m(v8, p_local_filter_start);
@@ -330,6 +289,11 @@
cmds.conv.stop = (in_channels_this_iter / 4) - 1;
aconv_vxv(v48, v0, cmds, v8);
+
+#pragma GCC unroll 8
+ for (int i = 0; i < 8; ++i) {
+ in_x[i] += dilation_width_factor;
+ }
}
}
in_channel += in_channels_this_iter;
@@ -349,64 +313,38 @@
vmax_w_vx_m(v52, v52, output_activation_min);
vsraqs_b_vx(v56, v48, 0);
vsraqs_b_vx(v57, v52, 0);
- if (out_channels_this_iter == 8) {
- if (out_xs_this_iter >= 4) {
- vstq_b_s_xx(v56,
- &output_data[tflite::Offset(output_shape, batch, out_y,
- out_x, out_channel)],
- output_depth);
- } else {
- for (int i = 0; i < std::min(4, out_xs_this_iter); ++i) {
- if (i > 0) {
- vsliden_b_4_vv(v58, v56, v0);
- vsliden_b_4_vv(v56, v58, v0);
- }
- vst_b_l_xx(v56,
- &output_data[tflite::Offset(output_shape, batch, out_y,
- out_x + i, out_channel)],
- out_channels_this_iter);
- }
- }
- if (out_xs_this_iter == 8) {
- vstq_b_s_xx(v57,
- &output_data[tflite::Offset(output_shape, batch, out_y,
- out_x + 4, out_channel)],
- output_depth);
- } else if (out_xs_this_iter > 4) {
- for (int i = 4; i < std::min(8, out_xs_this_iter); ++i) {
- if (i > 4) {
- vsliden_b_4_vv(v58, v57, v0);
- vsliden_b_4_vv(v57, v58, v0);
- }
- vst_b_l_xx(v57,
- &output_data[tflite::Offset(output_shape, batch, out_y,
- out_x + i, out_channel)],
- out_channels_this_iter);
- }
- }
- } else {
- for (int i = 0; i < std::min(4, out_xs_this_iter); ++i) {
- if (i > 0) {
- vsliden_b_4_vv(v58, v56, v0);
- vsliden_b_4_vv(v56, v58, v0);
- }
- vst_b_l_xx(v56,
- &output_data[tflite::Offset(output_shape, batch, out_y,
- out_x + i, out_channel)],
- out_channels_this_iter);
- }
- if (out_xs_this_iter > 4) {
- for (int i = 4; i < std::min(8, out_xs_this_iter); ++i) {
- if (i > 4) {
- vsliden_b_4_vv(v58, v57, v0);
- vsliden_b_4_vv(v57, v58, v0);
- }
- vst_b_l_xx(v57,
- &output_data[tflite::Offset(output_shape, batch, out_y,
- out_x + i, out_channel)],
- out_channels_this_iter);
- }
- }
+
+ const int8_t* p_out_x[8];
+#pragma GCC unroll 8
+ for (int i = 0; i < 8; ++i) {
+ p_out_x[i] = p_output + out_y_offset + ((out_x + i) * output_depth);
+ }
+
+ vslidep_h_4_vv(v58, v57, v57); // x7
+ vslidep_h_4_vv(v59, v58, v58); // x6
+ vslidep_h_4_vv(v60, v59, v59); // x5
+ vslidep_h_4_vv(v61, v60, v60); // x4
+ vslidep_h_4_vv(v62, v56, v56); // x3
+ vslidep_h_4_vv(v63, v62, v62); // x2
+ vslidep_h_4_vv(v57, v63, v63); // x1
+ vslidep_h_4_vv(v56, v57, v57); // x0
+ switch (out_xs_this_iter) {
+ case 8:
+ vst_b_l_xx(v58, p_out_x[7], out_channels_this_iter);
+ case 7:
+ vst_b_l_xx(v59, p_out_x[6], out_channels_this_iter);
+ case 6:
+ vst_b_l_xx(v60, p_out_x[5], out_channels_this_iter);
+ case 5:
+ vst_b_l_xx(v61, p_out_x[4], out_channels_this_iter);
+ case 4:
+ vst_b_l_xx(v62, p_out_x[3], out_channels_this_iter);
+ case 3:
+ vst_b_l_xx(v63, p_out_x[2], out_channels_this_iter);
+ case 2:
+ vst_b_l_xx(v57, p_out_x[1], out_channels_this_iter);
+ case 1:
+ vst_b_l_xx(v56, p_out_x[0], out_channels_this_iter);
}
out_x += out_xs_this_iter;
} while (out_x < output_width);
