8x16 Depthwise convolution for Kelvin. Change-Id: Ibb1beba0aad3d0f816c44982d4ec83e338f0d084
diff --git a/tflm/opt/BUILD b/tflm/opt/BUILD index 28419ce..f3f6e7d 100644 --- a/tflm/opt/BUILD +++ b/tflm/opt/BUILD
@@ -4,6 +4,7 @@ name = "opt", srcs = [ "conv.cc", + "depthwise_conv_s16.cc", "elementwise_add_s16.cc", "elementwise_add_s32.cc", "elementwise_add_s8.cc",
diff --git a/tflm/opt/depthwise_conv_s16.cc b/tflm/opt/depthwise_conv_s16.cc new file mode 100644 index 0000000..e05ef13 --- /dev/null +++ b/tflm/opt/depthwise_conv_s16.cc
@@ -0,0 +1,108 @@ +// Copyright 2023 Google LLC +// Licensed under the Apache License, Version 2.0, see LICENSE for details. +// SPDX-License-Identifier: Apache-2.0 + +#include <algorithm> + +#include "crt/kelvin.h" +#include "tflm/opt/opt.h" +#include "tensorflow/lite/kernels/internal/common.h" + +namespace kelvin::opt { + +void DepthwiseConv2DKelvinS16K3x1(const int16_t* activations, + const int8_t* weights, + const int64_t* biases, + int channels, int frames, + const int32_t* output_mult, + const int32_t* output_shift, + int32_t output_activation_min, + int32_t output_activation_max, + int16_t* output) { + for (int c = 0; c + 32 <= channels; c += 32) { + // Load weights and interleave into correct order [v58-v63]. + // Because there are more activations than weights, interleave weights. + const int8_t* local_weights0 = weights + c; + vld_b_p_xx(v0, local_weights0, channels); + vaddw_h_vx(v48, v0, 0); + vzip_h_vv(v58, v48, v49); + + vld_b_p_xx(v1, local_weights0, channels); + vaddw_h_vx(v50, v1, 0); + vzip_h_vv(v60, v50, v51); + + vld_b_x(v2, local_weights0); + vaddw_h_vx(v52, v2, 0); + vzip_h_vv(v62, v52, v53); + + // Assume biases fit in 32-bit. This assumption is verified offline. + // Load biases and swizzle [v52-v55]. + int32_t local_biases[32]; + for (int j = 0; j < 32; j++) { + local_biases[j] = static_cast<int32_t>(biases[c + j]); + } + vld_w_x_m(v4, local_biases); + vzip_w_vv(v52, v4, v5); + vzip_w_vv(v54, v6, v7); + + // Accumulators will be [v48 - v51]. + const int16_t* local_activations0 = activations + c; + const int16_t* local_activations1 = local_activations0 + 16; + int16_t* local_output = output + c; + + // Registers [v0-v5 will be for loading activations] + // Preload for valid padding: + vld_h_p_xx(v0, local_activations0, channels); + vld_h_p_xx(v1, local_activations1, channels); + vld_h_p_xx(v2, local_activations0, channels); + vld_h_p_xx(v3, local_activations1, channels); + int frames_left = frames - 2; + + const int32_t* local_output_mult = output_mult + c; + const int32_t* local_output_shift = output_shift + c; + + int32_t accumulators[32]; + while (frames_left > 0) { + vld_h_p_xx(v4, local_activations0, channels); + vld_h_p_xx(v5, local_activations1, channels); + vmulw_w_vv(v48, v58, v0); // Clobber accumulator + vmulw_w_vv(v50, v59, v1); // Clobber accumulator + vadd_w_vv_m(v48, v48, v52); // Add bias. + vmulw_w_vv(v40, v60, v2); + vmulw_w_vv(v42, v61, v3); + vadd_w_vv_m(v48, v48, v40); + vmulw_w_vv(v44, v62, v4); + vmulw_w_vv(v46, v63, v5); + vadd_w_vv_m(v48, v48, v44); + + vzip_w_vv(v48, v48, v49); // Swizzle accumulators + vzip_w_vv(v50, v50, v51); + + vst_w_x_m(v48, accumulators); // Store accumulators + + // Output pipeline in scalar, to preserve bit accuracy with the ARM CPU + // implementation. + for (int i = 0; i < 32; i++) { + int32_t result = tflite::MultiplyByQuantizedMultiplier( + static_cast<int64_t>(accumulators[i]), local_output_mult[i], + local_output_shift[i]); + + local_output[i] = static_cast<int16_t>( + std::clamp(result, output_activation_min, output_activation_max)); + } + + // Slide registers + vmvp_vv(v0, v2, v3); + vmvp_vv(v2, v4, v5); + + local_output += channels; + frames_left--; + } + } + // TODO(derekjchow): Handle channels % 32 cases. + // Break it down into: + // - one loop looking for 16 byte stripes + // - one final loop handling remainder +} + +} // namespace kelvin::opt \ No newline at end of file
diff --git a/tflm/opt/opt.h b/tflm/opt/opt.h index 29741cb..f12596c 100644 --- a/tflm/opt/opt.h +++ b/tflm/opt/opt.h
@@ -72,6 +72,11 @@ const int8_t* filter_data, const tflite::RuntimeShape& bias_shape, const int32_t* bias_data, const tflite::RuntimeShape& output_shape, int8_t* output_data); +void DepthwiseConv2DKelvinS16K3x1( + const int16_t* activations, const int8_t* weights, const int64_t* biases, + int channels, int frames, const int32_t* output_mult, + const int32_t* output_shift, int32_t output_activation_min, + int32_t output_activation_max, int16_t* output); } // namespace kelvin::opt #endif // TFLM_OPT_OPT_H_