Helper methods for Conv2D kernel - 16x8 helpers, one for 32-bit biases and one for 64-bit biases - 8x8 helper (using aconv) Change-Id: Ife5013ac8eceb51bbc4ca1091f2b465e5b315faf
diff --git a/crt/kelvin.h b/crt/kelvin.h index 9a16617..17163dc 100644 --- a/crt/kelvin.h +++ b/crt/kelvin.h
@@ -7,6 +7,7 @@ #ifndef CRT_KELVIN_H_ #define CRT_KELVIN_H_ +#include <assert.h> #include <math.h> #include <stdint.h> #include <string.h> @@ -39,4 +40,15 @@ #define vm14 v56 #define vm15 v60 +struct vconv_u8_t { + uint32_t mode : 2; // 1:0 + uint32_t start : 5; // 6:2 + uint32_t stop : 5; // 11:7 + 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 vconv_u8_t) == 4); + #endif // CRT_KELVIN_H_
diff --git a/tests/tflm/BUILD b/tests/tflm/BUILD index e58ac93..c0f17f7 100644 --- a/tests/tflm/BUILD +++ b/tests/tflm/BUILD
@@ -20,6 +20,25 @@ ) kelvin_test( + name = "conv_test", + srcs = [ + "@tflite-micro//tensorflow/lite/micro/kernels:conv_test.cc", + ], + hw_test_size = "enormous", + deps = [ + "//crt:crt_header", + "@tflite-micro//tensorflow/lite/c:common", + "@tflite-micro//tensorflow/lite/kernels/internal:tensor", + "@tflite-micro//tensorflow/lite/micro:micro_utils", + "@tflite-micro//tensorflow/lite/micro:test_helpers", + "@tflite-micro//tensorflow/lite/micro/kernels:conv_test_common", + "@tflite-micro//tensorflow/lite/micro/kernels:kernel_runner", + "@tflite-micro//tensorflow/lite/micro/kernels/testdata:conv_test_data", + "@tflite-micro//tensorflow/lite/micro/testing:micro_test", + ], +) + +kelvin_test( name = "leaky_relu_test", srcs = [ "@tflite-micro//tensorflow/lite/micro/kernels:leaky_relu_test.cc",
diff --git a/tflm/opt/BUILD b/tflm/opt/BUILD index a6722df..28419ce 100644 --- a/tflm/opt/BUILD +++ b/tflm/opt/BUILD
@@ -3,6 +3,7 @@ cc_library( name = "opt", srcs = [ + "conv.cc", "elementwise_add_s16.cc", "elementwise_add_s32.cc", "elementwise_add_s8.cc", @@ -17,6 +18,7 @@ target_compatible_with = ["@kelvin_sw//platforms/cpu:kelvin"], deps = [ "//crt:crt", + "@tflite-micro//tensorflow/lite/kernels/internal:common", ], alwayslink = True, )
diff --git a/tflm/opt/conv.cc b/tflm/opt/conv.cc new file mode 100644 index 0000000..92b2a31 --- /dev/null +++ b/tflm/opt/conv.cc
@@ -0,0 +1,375 @@ +// Copyright 2023 Google LLC +// Licensed under the Apache License, Version 2.0, see LICENSE for details. +// SPDX-License-Identifier: Apache-2.0 + +#include <cassert> +#include <memory> + +#include "crt/kelvin.h" +#include "tensorflow/lite/kernels/internal/common.h" +#include "tensorflow/lite/kernels/internal/runtime_shape.h" +#include "tensorflow/lite/kernels/internal/types.h" +#include "tflm/opt/opt.h" +#include "tflm/opt/util.h" + +namespace kelvin::opt { +namespace { +/* clang-format off */ +constexpr const int swizzle[16] = { + 0, 4, 8, 12, + 2, 6, 10, 14, + 1, 5, 9, 13, + 3, 7, 11, 15, +}; +/* clang-format on */ +} // namespace + +void conv_per_channel_b32( + const tflite::ConvParams& params, const int32_t* output_multiplier, + const int32_t* output_shift, const tflite::RuntimeShape& input_shape, + const int16_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, + int16_t* output_data) { + const auto batches = MatchingDim(input_shape, 0, output_shape, 0); + const auto stride_width = params.stride_width; + const auto stride_height = params.stride_height; + const auto dilation_width_factor = params.dilation_width_factor; + const auto dilation_height_factor = params.dilation_height_factor; + const auto pad_width = params.padding_values.width; + const auto pad_height = params.padding_values.height; + const auto input_height = input_shape.Dims(1); + const auto input_width = input_shape.Dims(2); + const auto input_depth = input_shape.Dims(3); + const auto input_offset = params.input_offset; + const auto filter_height = filter_shape.Dims(1); + const auto filter_width = filter_shape.Dims(2); + const auto filter_depth = filter_shape.Dims(3); + const auto output_height = output_shape.Dims(1); + const auto output_width = output_shape.Dims(2); + const auto output_depth = output_shape.Dims(3); + const auto output_offset = params.output_offset; + const auto output_activation_min = params.quantized_activation_min; + const auto output_activation_max = params.quantized_activation_max; + const auto groups = input_depth / filter_depth; + const auto filters_per_group = output_depth / groups; + + for (int batch = 0; batch < batches; ++batch) { + for (int out_y = 0; out_y < output_height; ++out_y) { + const int in_y_origin = out_y * stride_height - pad_height; + for (int out_x = 0; out_x < output_width; ++out_x) { + const int in_x_origin = out_x * stride_width - pad_width; + for (int out_channel = 0; out_channel < output_depth; ++out_channel) { + auto group = out_channel / filters_per_group; + int32_t acc32 = 0; + for (int filter_y = 0; filter_y < filter_height; ++filter_y) { + const int in_y = in_y_origin + dilation_height_factor * filter_y; + for (int filter_x = 0; filter_x < filter_width; ++filter_x) { + const int in_x = in_x_origin + dilation_width_factor * filter_x; + const bool inside = (in_x >= 0) && (in_x < input_width) && + (in_y >= 0) && (in_y < input_height); + if (!inside) { + continue; + } + int in_channel = 0; + do { + int load_count = std::min(filter_depth - in_channel, 16L); + int32_t input_swizzled[16]; + const int16_t* p_input = &input_data[tflite::Offset( + input_shape, batch, in_y, in_x, + in_channel + group * filter_depth)]; + for (int i = 0; i < 16; ++i) { + int swizzle_idx = swizzle[i]; + if (swizzle_idx < load_count) + input_swizzled[i] = *(p_input + swizzle_idx) + input_offset; + else + input_swizzled[i] = 0; + } + vld_w_l_xx(v0, input_swizzled, 4); + vld_w_l_xx(v1, input_swizzled + 4, 4); + vld_w_l_xx(v2, input_swizzled + 8, 4); + vld_w_l_xx(v3, input_swizzled + 12, 4); + vld_b_l_xx(v4, + &filter_data[tflite::Offset(filter_shape, + out_channel, filter_y, + filter_x, in_channel)], + load_count); + vaddw_h_vx(v4, v4, 0); + vaddw_w_vx(v6, v5, 0); + vaddw_w_vx(v4, v4, 0); + + vmul_w_vv_m(vm0, vm0, vm1); + vadd_w_vv(v0, v0, v1); + vadd_w_vv(v0, v0, v2); + vadd_w_vv(v0, v0, v3); + int32_t acc_spill[4]; + vst_w_l_xx(v0, acc_spill, 4); + for (int i = 0; i < 4; ++i) { + acc32 += acc_spill[i]; + } + in_channel += 16; + } while (in_channel + 16 <= filter_depth); + } + } + if (bias_data) { + acc32 = acc32 + bias_data[out_channel]; + } + int32_t acc = tflite::MultiplyByQuantizedMultiplier( + acc32, output_multiplier[out_channel], output_shift[out_channel]); + acc += output_offset; + acc = std::clamp(acc, output_activation_min, output_activation_max); + output_data[tflite::Offset(output_shape, batch, out_y, out_x, + out_channel)] = static_cast<int16_t>(acc); + } + } + } + } +} + +void conv_per_channel_b64( + const tflite::ConvParams& params, const int32_t* output_multiplier, + const int32_t* output_shift, const tflite::RuntimeShape& input_shape, + const int16_t* input_data, const tflite::RuntimeShape& filter_shape, + const int8_t* filter_data, const tflite::RuntimeShape& bias_shape, + const int64_t* bias_data, const tflite::RuntimeShape& output_shape, + int16_t* output_data) { + const auto batches = MatchingDim(input_shape, 0, output_shape, 0); + const auto stride_width = params.stride_width; + const auto stride_height = params.stride_height; + const auto dilation_width_factor = params.dilation_width_factor; + const auto dilation_height_factor = params.dilation_height_factor; + const auto pad_width = params.padding_values.width; + const auto pad_height = params.padding_values.height; + const auto input_height = input_shape.Dims(1); + const auto input_width = input_shape.Dims(2); + const auto input_depth = input_shape.Dims(3); + const auto input_offset = params.input_offset; + const auto filter_height = filter_shape.Dims(1); + const auto filter_width = filter_shape.Dims(2); + const auto filter_depth = filter_shape.Dims(3); + const auto output_height = output_shape.Dims(1); + const auto output_width = output_shape.Dims(2); + const auto output_depth = output_shape.Dims(3); + const auto output_offset = params.output_offset; + const auto output_activation_min = params.quantized_activation_min; + const auto output_activation_max = params.quantized_activation_max; + const auto groups = input_depth / filter_depth; + const auto filters_per_group = output_depth / groups; + for (int batch = 0; batch < batches; ++batch) { + for (int out_y = 0; out_y < output_height; ++out_y) { + const int in_y_origin = out_y * stride_height - pad_height; + for (int out_x = 0; out_x < output_width; ++out_x) { + const int in_x_origin = out_x * stride_width - pad_width; + for (int out_channel = 0; out_channel < output_depth; ++out_channel) { + auto group = out_channel / filters_per_group; + int64_t acc64 = 0; + for (int filter_y = 0; filter_y < filter_height; ++filter_y) { + const int in_y = in_y_origin + dilation_height_factor * filter_y; + for (int filter_x = 0; filter_x < filter_width; ++filter_x) { + const int in_x = in_x_origin + dilation_width_factor * filter_x; + const bool inside = (in_x >= 0) && (in_x < input_width) && + (in_y >= 0) && (in_y < input_height); + if (!inside) { + continue; + } + + int in_channel = 0; + do { + int load_count = std::min(filter_depth - in_channel, 16L); + int32_t input_swizzled[16]; + const int16_t* p_input = &input_data[tflite::Offset( + input_shape, batch, in_y, in_x, + in_channel + group * filter_depth)]; + for (int i = 0; i < 16; ++i) { + int swizzle_idx = swizzle[i]; + if (swizzle_idx < load_count) + input_swizzled[i] = *(p_input + swizzle_idx) + input_offset; + else + input_swizzled[i] = 0; + } + vld_w_l_xx(v0, input_swizzled, 4); + vld_w_l_xx(v1, input_swizzled + 4, 4); + vld_w_l_xx(v2, input_swizzled + 8, 4); + vld_w_l_xx(v3, input_swizzled + 12, 4); + vld_b_l_xx(v4, + &filter_data[tflite::Offset(filter_shape, + out_channel, filter_y, + filter_x, in_channel)], + load_count); + vaddw_h_vx(v4, v4, 0); + vaddw_w_vx(v6, v5, 0); + vaddw_w_vx(v4, v4, 0); + + vmul_w_vv_m(vm0, vm0, vm1); + vadd_w_vv(v0, v0, v1); + vadd_w_vv(v0, v0, v2); + vadd_w_vv(v0, v0, v3); + int32_t acc32[4]; + vst_w_l_xx(v0, acc32, 4); + for (int i = 0; i < 4; ++i) { + acc64 += acc32[i]; + } + in_channel += 16; + } while (in_channel + 16 <= filter_depth); + } + } + if (bias_data) { + acc64 = acc64 + bias_data[out_channel]; + } + int32_t acc = tflite::MultiplyByQuantizedMultiplier( + acc64, output_multiplier[out_channel], output_shift[out_channel]); + acc += output_offset; + acc = std::clamp(acc, output_activation_min, output_activation_max); + output_data[tflite::Offset(output_shape, batch, out_y, out_x, + out_channel)] = static_cast<int16_t>(acc); + } + } + } + } +} + +#define INA0 v0 +#define FLTA0 v8 +#define FLTA1 v9 +#define FLTA2 v10 +#define FLTA3 v11 +#define FLTA4 v12 +#define FLTA5 v13 +#define FLTA6 v14 +#define FLTA7 v15 +#define ACC v48 +#define ACC0 v48 +#define OUT0 v56 +void conv_per_channel_b8( + const tflite::ConvParams& 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 auto batches = MatchingDim(input_shape, 0, output_shape, 0); + const auto stride_width = params.stride_width; + const auto stride_height = params.stride_height; + const auto dilation_width_factor = params.dilation_width_factor; + const auto dilation_height_factor = params.dilation_height_factor; + const auto pad_width = params.padding_values.width; + const auto pad_height = params.padding_values.height; + const auto input_height = input_shape.Dims(1); + const auto input_width = input_shape.Dims(2); + const auto input_depth = input_shape.Dims(3); + const auto input_offset = params.input_offset; + const auto filter_height = filter_shape.Dims(1); + const auto filter_width = filter_shape.Dims(2); + const auto filter_depth = filter_shape.Dims(3); + const auto output_height = output_shape.Dims(1); + const auto output_width = output_shape.Dims(2); + const auto output_depth = output_shape.Dims(3); + const auto output_offset = params.output_offset; + const auto output_activation_min = params.quantized_activation_min; + const auto output_activation_max = params.quantized_activation_max; + const auto groups = input_depth / filter_depth; + const auto filters_per_group = output_depth / groups; + union { + vconv_u8_t conv; + uint32_t raw; + } cmds; + cmds.conv.mode = 0; + cmds.conv.start = 0; + cmds.conv.stop = 7; + cmds.conv.sbias1 = input_offset; + cmds.conv.sdata1 = true; + cmds.conv.sbias2 = 0; + cmds.conv.sdata2 = true; + + // Zero out accumulators. + vdup_b_x(v0, 0); + acset_v(ACC, v0); + vdup_b_x_m(ACC0, 0); + for (int batch = 0; batch < batches; ++batch) { + for (int out_y = 0; out_y < output_height; ++out_y) { + const int in_y_origin = (out_y * stride_height) - pad_height; + for (int out_x = 0; out_x < output_width; /*out_x += 32*/ ++out_x) { + const int in_x_origin = (out_x * stride_width) - pad_width; + for (int out_channel = 0; out_channel < output_depth; ++out_channel) { + auto group = out_channel / filters_per_group; + + for (int filter_y = 0; filter_y < filter_height; ++filter_y) { + const int in_y = in_y_origin + dilation_height_factor * filter_y; + const int in_x = in_x_origin + dilation_width_factor * 0; + + // Zero padding by omitting the areas outside the image. + const bool is_point_inside_image = + (in_x >= 0) && (in_x < input_width) && (in_y >= 0) && + (in_y < input_height); + if (!is_point_inside_image) { + continue; + } + + int q = filter_width * filter_depth; + for (int i = 0; i < q; i += 32) { + int count = std::min(q - i, 32); + count = std::min( + count, static_cast<int>((input_width - in_x) * filter_depth)); + int input_offset = tflite::Offset(input_shape, batch, in_y, in_x, + group * filter_depth) + + i; + vdup_w_x_m(vm0, 0); + vdup_w_x_m(vm1, 0); + vld_b_l_xx(INA0, &input_data[input_offset], count); + int filter_offset = + tflite::Offset(filter_shape, out_channel, filter_y, 0, 0) + i; + vdup_w_x_m(FLTA0, 0); + vdup_w_x_m(FLTA4, 0); + if (count > 0) { + vld_b_l_xx(FLTA0, &filter_data[filter_offset], + std::min(count, 4)); + } + if (count > 4) { + vld_b_l_xx(FLTA1, &filter_data[filter_offset + 4], + std::min(count - 4, 4)); + } + if (count > 8) { + vld_b_l_xx(FLTA2, &filter_data[filter_offset + 8], + std::min(count - 8, 4)); + } + if (count > 12) { + vld_b_l_xx(FLTA3, &filter_data[filter_offset + 12], + std::min(count - 12, 4)); + } + if (count > 16) { + vld_b_l_xx(FLTA4, &filter_data[filter_offset + 16], + std::min(count - 16, 4)); + } + if (count > 20) { + vld_b_l_xx(FLTA5, &filter_data[filter_offset + 20], + std::min(count - 20, 4)); + } + if (count > 24) { + vld_b_l_xx(FLTA6, &filter_data[filter_offset + 24], + std::min(count - 24, 4)); + } + if (count > 28) { + vld_b_l_xx(FLTA7, &filter_data[filter_offset + 28], + std::min(count - 28, 4)); + } + aconv_vxv(ACC, INA0, cmds, FLTA0); + } + } + vcget(ACC); + vadd_w_vx_m(ACC0, ACC0, bias_data[out_channel]); + vsll_w_vx_m(ACC0, ACC0, LEFT_SHIFT(output_shift[out_channel])); + vdmulh_w_r_vx_m(ACC0, ACC0, output_multiplier[out_channel]); + vsha_w_r_vx_m(ACC0, ACC0, RIGHT_SHIFT(output_shift[out_channel])); + vadd_w_vx_m(ACC0, ACC0, output_offset); + vmin_w_vx_m(ACC0, ACC0, output_activation_max); + vmax_w_vx_m(ACC0, ACC0, output_activation_min); + vsraqs_b_vx(OUT0, ACC0, 0); + size_t output_offset = + tflite::Offset(output_shape, batch, out_y, out_x, out_channel); + vst_b_l_xx(OUT0, &output_data[output_offset], 1); + } + } + } + } +} +} // namespace kelvin::opt
diff --git a/tflm/opt/opt.h b/tflm/opt/opt.h index 6009ee3..29741cb 100644 --- a/tflm/opt/opt.h +++ b/tflm/opt/opt.h
@@ -5,6 +5,12 @@ #ifndef TFLM_OPT_OPT_H_ #define TFLM_OPT_OPT_H_ +/* clang-format off */ +#include <cstring> +#include "tensorflow/lite/kernels/internal/runtime_shape.h" +#include "tensorflow/lite/kernels/internal/types.h" +/* clang-format on */ + namespace kelvin::opt { void *memcpy(void *dst, const void *src, size_t n); void elementwise_add_s8(const int8_t* input1, const int8_t* input2, @@ -45,6 +51,27 @@ const int32_t output_shift_alpha, const int32_t output_multiplier_identity, const int32_t output_shift_identity); +void conv_per_channel_b32( + const tflite::ConvParams& params, const int32_t* output_multiplier, + const int32_t* output_shift, const tflite::RuntimeShape& input_shape, + const int16_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, + int16_t* output_data); +void conv_per_channel_b64( + const tflite::ConvParams& params, const int32_t* output_multiplier, + const int32_t* output_shift, const tflite::RuntimeShape& input_shape, + const int16_t* input_data, const tflite::RuntimeShape& filter_shape, + const int8_t* filter_data, const tflite::RuntimeShape& bias_shape, + const int64_t* bias_data, const tflite::RuntimeShape& output_shape, + int16_t* output_data); +void conv_per_channel_b8( + const tflite::ConvParams& 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); } // namespace kelvin::opt #endif // TFLM_OPT_OPT_H_