Changes to optimize max_pool for Kelvin int8 dtype
Change-Id: I9bf37f7a5438d83b93648dacd2cf4eaec86698cb
diff --git a/tests/tflm/BUILD b/tests/tflm/BUILD
index 04063a0..1b6a046 100644
--- a/tests/tflm/BUILD
+++ b/tests/tflm/BUILD
@@ -86,6 +86,22 @@
)
kelvin_test(
+ name = "pooling_test",
+ srcs = [
+ "@tflite-micro//tensorflow/lite/micro/kernels:pooling_test.cc",
+ ],
+ deps = [
+ "//crt",
+ "@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:kernel_runner",
+ "@tflite-micro//tensorflow/lite/micro/testing:micro_test",
+ ],
+)
+
+kelvin_test(
name = "reshape_test",
srcs = [
"@tflite-micro//tensorflow/lite/micro/kernels:reshape_test.cc",
diff --git a/tflm/opt/BUILD b/tflm/opt/BUILD
index f518259..6481a8a 100644
--- a/tflm/opt/BUILD
+++ b/tflm/opt/BUILD
@@ -25,6 +25,7 @@
"leaky_relu_s16.cc",
"leaky_relu_s8.cc",
"memcpy.cc",
+ "max_pool_s8.cc"
],
hdrs = [
"opt.h",
diff --git a/tflm/opt/max_pool_s8.cc b/tflm/opt/max_pool_s8.cc
new file mode 100644
index 0000000..5986746
--- /dev/null
+++ b/tflm/opt/max_pool_s8.cc
@@ -0,0 +1,100 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#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"
+
+namespace kelvin::opt {
+void MaxPoolGeneric(const tflite::PoolParams ¶ms,
+ const tflite::RuntimeShape &input_shape,
+ const int8_t *input_data,
+ const tflite::RuntimeShape &output_shape,
+ int8_t *output_data) {
+ const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+ const int depth = MatchingDim(input_shape, 3, output_shape, 3);
+ const int input_height = input_shape.Dims(1);
+ const int input_width = input_shape.Dims(2);
+ const int output_height = output_shape.Dims(1);
+ const int output_width = output_shape.Dims(2);
+ const int stride_height = params.stride_height;
+ const int stride_width = params.stride_width;
+ for (int batch = 0; batch < batches; ++batch) {
+ for (int out_y = 0; out_y < output_height; ++out_y) {
+ for (int out_x = 0; out_x < output_width; ++out_x) {
+ const int in_x_origin =
+ (out_x * stride_width) - params.padding_values.width;
+ const int in_y_origin =
+ (out_y * stride_height) - params.padding_values.height;
+
+ // Compute the boundaries of the filter region clamped so as to
+ // ensure that the filter window fits in the input array.
+ const int filter_x_start = std::max(0, -in_x_origin);
+ const int filter_x_end =
+ std::min(params.filter_width, input_width - in_x_origin);
+ const int filter_y_start = std::max(0, -in_y_origin);
+ const int filter_y_end =
+ std::min(params.filter_height, input_height - in_y_origin);
+
+ int channel = 0;
+ for (; channel + 32 <= depth; channel += 32) {
+ vdup_b_x(v0, params.quantized_activation_min);
+ for (int filter_y = filter_y_start; filter_y < filter_y_end;
+ ++filter_y) {
+ for (int filter_x = filter_x_start; filter_x < filter_x_end;
+ ++filter_x) {
+ const int in_x = in_x_origin + filter_x;
+ const int in_y = in_y_origin + filter_y;
+ const int8_t *local_input =
+ input_data + Offset(input_shape, batch, in_y, in_x, channel);
+ vld_b_x(v1, local_input);
+ vmax_b_vv(v0, v0, v1);
+ }
+ }
+ vmin_b_vx(v0, v0, params.quantized_activation_max);
+ int8_t *local_output =
+ output_data + Offset(output_shape, batch, out_y, out_x, channel);
+ vst_b_x(v0, local_output);
+ }
+
+ if (channel == depth) {
+ continue;
+ }
+ int remaining_channels = depth - channel;
+ vdup_b_x(v0, params.quantized_activation_min);
+ for (int filter_y = filter_y_start; filter_y < filter_y_end;
+ ++filter_y) {
+ for (int filter_x = filter_x_start; filter_x < filter_x_end;
+ ++filter_x) {
+ const int in_x = in_x_origin + filter_x;
+ const int in_y = in_y_origin + filter_y;
+ const int8_t *local_input =
+ input_data + Offset(input_shape, batch, in_y, in_x, depth - 1);
+ vld_b_l_xx(v1, local_input, remaining_channels);
+ vmax_b_vv(v0, v0, v1);
+ }
+ }
+ vmin_b_vx(v0, v0, params.quantized_activation_max);
+ int8_t *local_output =
+ output_data + Offset(output_shape, batch, out_y, out_x, depth - 1);
+ vst_b_l_xx(v0, local_output, remaining_channels);
+ }
+ }
+ }
+}
+
+} // namespace kelvin::opt
diff --git a/tflm/opt/opt.h b/tflm/opt/opt.h
index e4268c0..9f83bf6 100644
--- a/tflm/opt/opt.h
+++ b/tflm/opt/opt.h
@@ -120,6 +120,10 @@
int channels, int frames, int dilation, const int32_t* output_mult,
const int32_t* output_shift, int32_t output_activation_min,
int32_t output_activation_max, int16_t* output);
+void MaxPoolGeneric(const tflite::PoolParams& params, const tflite::RuntimeShape& input_shape,
+ const int8_t* input_data, const tflite::RuntimeShape& output_shape,
+ int8_t* output_data);
+
} // namespace kelvin::opt
#endif // TFLM_OPT_OPT_H_
