Relax output_depth for ConvS8D4
- ConvS8D4 can now handle output depths that are not a multiple of 8.
Change-Id: Ibfc445e09f392453ddb8c038a5b224b8bcb33c4e
diff --git a/tflm/opt/conv_s8.cc b/tflm/opt/conv_s8.cc
index 2d49dbc..ce56fda 100644
--- a/tflm/opt/conv_s8.cc
+++ b/tflm/opt/conv_s8.cc
@@ -198,7 +198,7 @@
// special case of filter_depth = 4n
if (dilation_width_factor == 1 && dilation_height_factor == 1 &&
stride_width <= 2 && stride_height <= 2 && filter_depth % 4 == 0 &&
- output_depth % 8 == 0 && output_width >= 8 && pad_width <= 1) {
+ output_depth >= 8 && output_width >= 8 && pad_width <= 1) {
fn = kelvin::opt::ConvS8D4;
}
diff --git a/tflm/opt/conv_s8_d4.cc b/tflm/opt/conv_s8_d4.cc
index 0dd3e50..18df3e7 100644
--- a/tflm/opt/conv_s8_d4.cc
+++ b/tflm/opt/conv_s8_d4.cc
@@ -29,11 +29,12 @@
namespace kelvin::opt {
namespace {
-void Filter_8_H_W_M(const int8_t* input, int8_t* output, int H, int W, int M) {
+
+void Filter_N_H_W_M(const int8_t* input, int8_t* output, int N, int H, int W, int M) {
const int8_t(&in)[8][H][W][M] = *(int8_t(*)[8][H][W][M])input;
int8_t(&out)[H][W][M / 4][8][4] = *(int8_t(*)[H][W][M / 4][8][4]) output;
assert(M >= 4);
- for (int zo = 0; zo < 8; ++zo) {
+ for (int zo = 0; zo < N; ++zo) {
for (int ky = 0; ky < H; ++ky) {
for (int kx = 0; kx < W; ++kx) {
for (int zi = 0; zi < M; ++zi) {
@@ -44,28 +45,33 @@
}
}
}
-}
-
-void Swizzle(const int32_t* input, int32_t* output, int N) {
- const int32_t(&in)[N] = *(int32_t(*)[N])input;
- int32_t(&out)[N * 4] = *(int32_t(*)[N * 4]) output;
- // Convert to accumulator swizzle pattern.
- for (int i = 0; i < N / 8; ++i) {
- int32_t* out0 = out + i * 32 + 0;
- int32_t* out1 = out + i * 32 + 16;
- int32_t* out2 = out + i * 32 + 8;
- int32_t* out3 = out + i * 32 + 24;
- for (int j = 0; j < 4; ++j) {
- const int32_t* p_in = in + i * 8;
- for (int k = 0; k < 2; ++k) {
- *out0++ = *p_in++;
- *out1++ = *p_in++;
- *out2++ = *p_in++;
- *out3++ = *p_in++;
+ // Zero out the rest of the output.
+ for (int zo = N; zo < 8; ++zo) {
+ for (int ky = 0; ky < H; ++ky) {
+ for (int kx = 0; kx < W; ++kx) {
+ for (int zi = 0; zi < M; ++zi) {
+ const int zi_hi = zi >> 2; // div4
+ const int zi_lo = zi & 3; // rem4
+ out[ky][kx][zi_hi][zo][zi_lo] = 0;
+ }
}
}
}
}
+
+void Swizzle(const int32_t* input, int32_t* output, int N) {
+ assert(N <= 8);
+ const int32_t(&in)[8] = *(int32_t(*)[8])input;
+ int32_t(&out)[32] = *(int32_t(*)[32]) output;
+ // Convert to accumulator swizzle pattern.
+ memset(out, 0, 32 * sizeof(int32_t));
+ int offsets[] = {0, 16, 8, 24, 1, 17, 9, 25};
+ for (int i = 0; i < N; ++i) {
+ int offset = offsets[i];
+ out[0 + offset] = out[2 + offset] = out[4 + offset] = out[6 + offset] = in[i];
+ }
+}
+
} // namespace
void ConvS8D4(
@@ -137,14 +143,16 @@
int32_t swizzled_mult_data[32];
int32_t swizzled_shift_data[32];
- for (int out_channel = 0; out_channel + 8 <= output_depth; out_channel += 8) {
- Filter_8_H_W_M(filter_data + (out_channel * filter_height * filter_width *
+ int out_channel = 0;
+ do {
+ int out_channels_this_iter = std::min(8, output_depth - out_channel);
+ Filter_N_H_W_M(filter_data + (out_channel * filter_height * filter_width *
filter_input_depth),
- p_swizzled_filter_data, filter_height, filter_width,
+ p_swizzled_filter_data, out_channels_this_iter, filter_height, filter_width,
filter_input_depth);
- Swizzle(bias_data + out_channel, swizzled_bias_data, 8);
- Swizzle(output_multiplier + out_channel, swizzled_mult_data, 8);
- Swizzle(output_shift + out_channel, swizzled_shift_data, 8);
+ Swizzle(bias_data + out_channel, swizzled_bias_data, out_channels_this_iter);
+ Swizzle(output_multiplier + out_channel, swizzled_mult_data, out_channels_this_iter);
+ Swizzle(output_shift + out_channel, swizzled_shift_data, out_channels_this_iter);
vld_w_x_m(v16, swizzled_bias_data);
vld_w_x_m(v20, swizzled_mult_data);
vld_w_x_m(v24, swizzled_shift_data);
@@ -162,7 +170,7 @@
acset_v(v48, v48);
int in_channel = 0;
do {
- int channels_this_iter = std::min(filter_input_depth, 32);
+ 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 =
@@ -320,11 +328,11 @@
vld_b_p_x_m(v8, p_local_filter_start);
vld_b_x_m(v12, p_local_filter_start);
- cmds.conv.stop = (channels_this_iter / 4) - 1;
+ cmds.conv.stop = (in_channels_this_iter / 4) - 1;
aconv_vxv(v48, v0, cmds, v8);
}
}
- in_channel += channels_this_iter;
+ in_channel += in_channels_this_iter;
} while (in_channel < filter_input_depth);
vcget(v48);
vadd_w_vv_m(v48, v48, v16);
@@ -341,44 +349,70 @@
vmax_w_vx_m(v52, v52, output_activation_min);
vsraqs_b_vx(v56, v48, 0);
vsraqs_b_vx(v57, v52, 0);
- 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 < out_xs_this_iter; ++i) {
- if (i > 0) {
- vsliden_b_4_vv(v58, v56, v0);
- vsliden_b_4_vv(v56, v58, v0);
+ 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);
}
- vst_b_l_xx(v56,
- &output_data[tflite::Offset(output_shape, batch, out_y,
- out_x + i, out_channel)],
- 8);
}
- }
- 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 < out_xs_this_iter; ++i) {
- if (i > 4) {
- vsliden_b_4_vv(v58, v57, v0);
- vsliden_b_4_vv(v57, v58, v0);
+ 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);
}
- vst_b_l_xx(v57,
- &output_data[tflite::Offset(output_shape, batch, out_y,
- out_x + i, out_channel)],
- 8);
+ }
+ } 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);
+ }
}
}
out_x += out_xs_this_iter;
} while (out_x < output_width);
}
}
- }
+ out_channel += out_channels_this_iter;
+ } while (out_channel < output_depth);
}
} // namespace kelvin::opt
