Support functions for LeakyRelu
- Add int8 and int16 helper methods for implementing a LeakyRelu kernel.
Change-Id: Id017f614fa6638273eb63f50e1196ca6da44ca51
diff --git a/tests/tflm/BUILD b/tests/tflm/BUILD
index 71f70ab..05c196b 100644
--- a/tests/tflm/BUILD
+++ b/tests/tflm/BUILD
@@ -1,4 +1,5 @@
load("//build_tools/bazel:kelvin.bzl", "kelvin_test")
+
package(default_visibility = ["//visibility:public"])
kelvin_test(
@@ -10,10 +11,26 @@
"//crt:crt_header",
"@tflite-micro//tensorflow/lite/c:common",
"@tflite-micro//tensorflow/lite/kernels/internal:tensor",
- "@tflite-micro//tensorflow/lite/micro/kernels:kernel_runner",
- "@tflite-micro//tensorflow/lite/micro/testing:micro_test",
"@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 = "leaky_relu_test",
+ srcs = [
+ "@tflite-micro//tensorflow/lite/micro/kernels:leaky_relu_test.cc",
+ ],
+ 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:kernel_runner",
+ "@tflite-micro//tensorflow/lite/micro/testing:micro_test",
],
)
@@ -26,9 +43,9 @@
"//crt:crt_header",
"@tflite-micro//tensorflow/lite/c:common",
"@tflite-micro//tensorflow/lite/kernels/internal:tensor",
- "@tflite-micro//tensorflow/lite/micro/kernels:kernel_runner",
- "@tflite-micro//tensorflow/lite/micro/testing:micro_test",
"@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",
],
)
diff --git a/tflm/opt/BUILD b/tflm/opt/BUILD
index 2ba3900..8da2593 100644
--- a/tflm/opt/BUILD
+++ b/tflm/opt/BUILD
@@ -3,18 +3,20 @@
cc_library(
name = "opt",
srcs = [
- "elementwise_add_s8.cc",
"elementwise_add_s16.cc",
"elementwise_add_s32.cc",
+ "elementwise_add_s8.cc",
+ "leaky_relu_s16.cc",
+ "leaky_relu_s8.cc",
"memcpy.cc",
],
hdrs = [
"opt.h",
"util.h",
],
+ target_compatible_with = ["@kelvin_sw//platforms/cpu:kelvin"],
deps = [
"//crt:crt_header",
],
alwayslink = True,
- target_compatible_with = ["@kelvin_sw//platforms/cpu:kelvin"],
)
diff --git a/tflm/opt/leaky_relu_s16.cc b/tflm/opt/leaky_relu_s16.cc
new file mode 100644
index 0000000..e3ac66d
--- /dev/null
+++ b/tflm/opt/leaky_relu_s16.cc
@@ -0,0 +1,73 @@
+// Copyright 2023 Google LLC
+// Licensed under the Apache License, Version 2.0, see LICENSE for details.
+// SPDX-License-Identifier: Apache-2.0
+
+#include <limits>
+
+#include "crt/kelvin.h"
+#include "tflm/opt/opt.h"
+#include "tflm/opt/util.h"
+
+namespace kelvin::opt {
+void leaky_relu_s16(const int16_t* input, int16_t* output,
+ const int32_t block_size, const int32_t input_zero_point,
+ const int32_t output_zero_point,
+ const int32_t output_multiplier_alpha,
+ const int32_t output_shift_alpha,
+ const int32_t output_multiplier_identity,
+ const int32_t output_shift_identity) {
+ constexpr int32_t quantized_output_min = std::numeric_limits<int16_t>::min();
+ constexpr int32_t quantized_output_max = std::numeric_limits<int16_t>::max();
+ int32_t right_shift_identity = std::min(output_shift_identity, 0L);
+ int32_t left_shift_identity = std::max(output_shift_identity, 0L);
+ int32_t right_shift_alpha = std::min(output_shift_alpha, 0L);
+ int32_t left_shift_alpha = std::max(output_shift_alpha, 0L);
+ int blocks = block_size;
+ int vl;
+ getmaxvl_h(vl);
+ while (blocks) {
+ int count = std::min(blocks, vl);
+
+ // Load data from the input, and widen.
+ vld_h_lp_xx(v0, input, count);
+ vaddw_w_vx(v0, v0, 0);
+
+ // Subtract out the provided offset from the inputs.
+ vsub_w_vx_m(vm0, vm0, input_zero_point);
+
+ // Compute the Relu on all inputs, as if they were >=0.
+ vsll_w_vx_m(vm2, vm0, left_shift_identity);
+ vdmulh_w_r_vx_m(vm2, vm2, output_multiplier_identity);
+ vsha_w_vx_m(vm2, vm2, RIGHT_SHIFT(right_shift_identity));
+ vadd_w_vx_m(vm2, vm2, output_zero_point);
+ vmax_w_vx_m(vm2, vm2, quantized_output_min);
+ vmin_w_vx_m(vm2, vm2, quantized_output_max);
+
+ // Compute the Relu on all inputs, as if they were <0.
+ vsll_w_vx_m(vm1, vm0, left_shift_alpha);
+ vdmulh_w_r_vx_m(vm1, vm1, output_multiplier_alpha);
+ vsha_w_vx_m(vm1, vm1, RIGHT_SHIFT(right_shift_alpha));
+ vadd_w_vx_m(vm1, vm1, output_zero_point);
+ vmax_w_vx_m(vm1, vm1, quantized_output_min);
+ vmin_w_vx_m(vm1, vm1, quantized_output_max);
+
+ // Compute a boolean vector for inputs >=0.
+ vge_w_vx_m(vm3, vm0, 0);
+ // Compute a boolean vector for inputs <0.
+ vlt_w_vx_m(vm0, vm0, 0);
+ // Multiply the `identity` results by the >=0 vector.
+ vmul_w_vv_m(vm2, vm2, vm3);
+ // Multiply the `alpha` results by the <0 vector.
+ vmul_w_vv_m(vm0, vm1, vm0);
+ // Sum the two resulting vectors.
+ vadd_w_vv_m(vm0, vm0, vm2);
+
+ // Narrow/swizzle, and store to output.
+ vsrans_h_vx(v0, v0, 0);
+ vst_h_lp_xx(v0, output, count);
+
+ blocks -= count;
+ }
+}
+
+} // namespace kelvin::opt
diff --git a/tflm/opt/leaky_relu_s8.cc b/tflm/opt/leaky_relu_s8.cc
new file mode 100644
index 0000000..15b9218
--- /dev/null
+++ b/tflm/opt/leaky_relu_s8.cc
@@ -0,0 +1,76 @@
+// Copyright 2023 Google LLC
+// Licensed under the Apache License, Version 2.0, see LICENSE for details.
+// SPDX-License-Identifier: Apache-2.0
+
+#include <limits>
+
+#include "crt/kelvin.h"
+#include "tflm/opt/opt.h"
+#include "tflm/opt/util.h"
+
+namespace kelvin::opt {
+
+void leaky_relu_s8(const int8_t* input, int8_t* output,
+ const int32_t block_size, const int32_t input_zero_point,
+ const int32_t output_zero_point,
+ const int32_t output_multiplier_alpha,
+ const int32_t output_shift_alpha,
+ const int32_t output_multiplier_identity,
+ const int32_t output_shift_identity) {
+ constexpr int32_t quantized_output_min = std::numeric_limits<int16_t>::min();
+ constexpr int32_t quantized_output_max = std::numeric_limits<int16_t>::max();
+ int32_t right_shift_identity = std::min(output_shift_identity, 0L);
+ int32_t left_shift_identity = std::max(output_shift_identity, 0L);
+ int32_t right_shift_alpha = std::min(output_shift_alpha, 0L);
+ int32_t left_shift_alpha = std::max(output_shift_alpha, 0L);
+ int blocks = block_size;
+ int vl;
+ getmaxvl_b(vl);
+ while (blocks) {
+ int count = std::min(blocks, vl);
+
+ // Load data from the input, and widen (now we can use vm0).
+ vld_b_lp_xx(v0, input, count);
+ vaddw_h_vx(v0, v0, 0);
+ vaddw_w_vx(v2, v1, 0);
+ vaddw_w_vx(v0, v0, 0);
+
+ // Subtract out the provided offset from the inputs.
+ vsub_w_vx_m(vm0, vm0, input_zero_point);
+
+ // Compute the Relu on all inputs, as if they were >=0.
+ vsll_w_vx_m(vm2, vm0, left_shift_identity);
+ vdmulh_w_r_vx_m(vm2, vm2, output_multiplier_identity);
+ vsha_w_vx_m(vm2, vm2, RIGHT_SHIFT(right_shift_identity));
+ vadd_w_vx_m(vm2, vm2, output_zero_point);
+ vmax_w_vx_m(vm2, vm2, quantized_output_min);
+ vmin_w_vx_m(vm2, vm2, quantized_output_max);
+
+ // Compute the Relu on all inputs, as if they were <0.
+ vsll_w_vx_m(vm1, vm0, left_shift_alpha);
+ vdmulh_w_r_vx_m(vm1, vm1, output_multiplier_alpha);
+ vsha_w_vx_m(vm1, vm1, RIGHT_SHIFT(right_shift_alpha));
+ vadd_w_vx_m(vm1, vm1, output_zero_point);
+ vmax_w_vx_m(vm1, vm1, quantized_output_min);
+ vmin_w_vx_m(vm1, vm1, quantized_output_max);
+
+ // Compute a boolean vector for inputs >=0.
+ vge_w_vx_m(vm3, vm0, 0);
+ // Compute a boolean vector for inputs <0.
+ vlt_w_vx_m(vm0, vm0, 0);
+ // Multiply the `identity` results by the >=0 vector.
+ vmul_w_vv_m(vm2, vm2, vm3);
+ // Multiply the `alpha` results by the <0 vector.
+ vmul_w_vv_m(vm0, vm1, vm0);
+ // Sum the two resulting vectors.
+ vadd_w_vv_m(vm0, vm0, vm2);
+
+ // Narrow/swizzle, and store to output.
+ vsraqs_b_vx(v0, v0, 0);
+ vst_b_lp_xx(v0, output, count);
+
+ blocks -= count;
+ }
+}
+
+} // namespace kelvin::opt
diff --git a/tflm/opt/opt.h b/tflm/opt/opt.h
index 12075ab..6009ee3 100644
--- a/tflm/opt/opt.h
+++ b/tflm/opt/opt.h
@@ -31,6 +31,20 @@
int32_t* output, const int32_t output_activation_min,
const int32_t output_activation_max,
const int32_t block_size);
+void leaky_relu_s8(const int8_t* input, int8_t* output,
+ const int32_t block_size, const int32_t input_zero_point,
+ const int32_t output_zero_point,
+ const int32_t output_multiplier_alpha,
+ const int32_t output_shift_alpha,
+ const int32_t output_multiplier_identity,
+ const int32_t output_shift_identity);
+void leaky_relu_s16(const int16_t* input, int16_t* output,
+ const int32_t block_size, const int32_t input_zero_point,
+ const int32_t output_zero_point,
+ const int32_t output_multiplier_alpha,
+ const int32_t output_shift_alpha,
+ const int32_t output_multiplier_identity,
+ const int32_t output_shift_identity);
} // namespace kelvin::opt
#endif // TFLM_OPT_OPT_H_
