Skip to content

Commit bfeb43a

Browse files
veblushveblush
authored
Use SafeCast for float-to-integer conversions in CAST kernel (#3660)
* Use SafeCast for float-to-integer conversions in CAST kernel * Updated by review
1 parent da21b09 commit bfeb43a

2 files changed

Lines changed: 101 additions & 57 deletions

File tree

‎tensorflow/lite/micro/kernels/cast.cc‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,11 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16+
#include <algorithm>
17+
#include <type_traits>
18+
1619
#include "tensorflow/lite/c/common.h"
20+
#include "tensorflow/lite/kernels/internal/quantization_util.h"
1721
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
1822
#include "tensorflow/lite/kernels/kernel_util.h"
1923
#include "tensorflow/lite/micro/kernels/kernel_util.h"
@@ -44,10 +48,20 @@ TfLiteStatus CastPrepare(TfLiteContext* context, TfLiteNode* node) {
4448
return kTfLiteOk;
4549
}
4650

51+
template <typename ToT, typename FromT>
52+
ToT CastValue(FromT a) {
53+
if constexpr (std::is_floating_point_v<FromT> && std::is_integral_v<ToT> &&
54+
!std::is_same_v<ToT, bool>) {
55+
return tflite::SafeCast<ToT>(a);
56+
} else {
57+
return static_cast<ToT>(a);
58+
}
59+
}
60+
4761
template <typename FromT, typename ToT>
4862
void copyCast(const FromT* in, ToT* out, int num_elements) {
4963
std::transform(in, in + num_elements, out,
50-
[](FromT a) { return static_cast<ToT>(a); });
64+
[](FromT a) { return CastValue<ToT>(a); });
5165
}
5266

5367
template <typename FromT>
@@ -69,6 +83,9 @@ TfLiteStatus copyToTensor(TfLiteContext* context, const FromT* in,
6983
case kTfLiteFloat32:
7084
copyCast(in, tflite::micro::GetTensorData<float>(out), num_elements);
7185
break;
86+
case kTfLiteBool:
87+
copyCast(in, tflite::micro::GetTensorData<bool>(out), num_elements);
88+
break;
7289
default:
7390
// Unsupported type.
7491
MicroPrintf("Output type %s (%d) not supported.",

‎tensorflow/lite/micro/kernels/cast_test.cc‎

Lines changed: 83 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@ See the License for the specific language governing permissions and
1313
limitations under the License.
1414
==============================================================================*/
1515

16+
#include <cstddef>
17+
#include <cstdint>
18+
#include <limits>
19+
1620
#include "tensorflow/lite/c/builtin_op_data.h"
1721
#include "tensorflow/lite/c/common.h"
1822
#include "tensorflow/lite/micro/kernels/kernel_runner.h"
@@ -23,18 +27,17 @@ namespace tflite {
2327
namespace testing {
2428
namespace {
2529

26-
template <typename inputT, typename outputT>
27-
void TestCast(int* input_dims_data, const inputT* input_data,
28-
const outputT* expected_output_data, outputT* output_data) {
29-
TfLiteIntArray* input_dims = IntArrayFromInts(input_dims_data);
30-
TfLiteIntArray* output_dims = IntArrayFromInts(input_dims_data);
31-
const int output_dims_count = ElementCount(*output_dims);
32-
constexpr int inputs_size = 1;
33-
constexpr int outputs_size = 1;
34-
constexpr int tensors_size = inputs_size + outputs_size;
30+
template <typename InputT, typename OutputT, size_t N>
31+
void TestCast(const InputT (&input)[N], const OutputT (&golden)[N]) {
32+
OutputT output_data[N];
33+
34+
int dims_data[] = {1, static_cast<int>(N)};
35+
TfLiteIntArray* dims = IntArrayFromInts(dims_data);
36+
37+
constexpr int tensors_size = 2;
3538
TfLiteTensor tensors[tensors_size] = {
36-
CreateTensor(input_data, input_dims),
37-
CreateTensor(output_data, output_dims),
39+
CreateTensor(input, dims),
40+
CreateTensor(output_data, dims),
3841
};
3942

4043
int inputs_array_data[] = {1, 0};
@@ -44,95 +47,119 @@ void TestCast(int* input_dims_data, const inputT* input_data,
4447

4548
const TFLMRegistration registration = Register_CAST();
4649
micro::KernelRunner runner(registration, tensors, tensors_size, inputs_array,
47-
outputs_array,
48-
/*builtin_data=*/nullptr);
50+
outputs_array, /*builtin_data=*/nullptr);
4951

5052
EXPECT_EQ(kTfLiteOk, runner.InitAndPrepare());
5153
EXPECT_EQ(kTfLiteOk, runner.Invoke());
5254

53-
for (int i = 0; i < output_dims_count; ++i) {
54-
EXPECT_EQ(expected_output_data[i], output_data[i]);
55+
for (size_t i = 0; i < N; ++i) {
56+
EXPECT_EQ(golden[i], output_data[i]);
5557
}
5658
}
5759

60+
template <typename IntT>
61+
void TestFloatToInt(float pos_overflow, float neg_overflow) {
62+
constexpr bool is_signed = std::numeric_limits<IntT>::is_signed;
63+
const float input[] = {100.f,
64+
1.0f,
65+
0.f,
66+
0.4f,
67+
1.999f,
68+
1.1f,
69+
-1.0f,
70+
-100.f,
71+
pos_overflow,
72+
neg_overflow,
73+
std::numeric_limits<float>::infinity(),
74+
-std::numeric_limits<float>::infinity(),
75+
std::numeric_limits<float>::quiet_NaN()};
76+
const IntT golden[] = {100,
77+
1,
78+
0,
79+
0,
80+
1,
81+
1,
82+
is_signed ? static_cast<IntT>(-1) : IntT{0},
83+
is_signed ? static_cast<IntT>(-100) : IntT{0},
84+
std::numeric_limits<IntT>::max(),
85+
std::numeric_limits<IntT>::min(),
86+
std::numeric_limits<IntT>::max(),
87+
std::numeric_limits<IntT>::min(),
88+
0};
89+
TestCast(input, golden);
90+
}
91+
5892
} // namespace
5993
} // namespace testing
6094
} // namespace tflite
6195

6296
TEST(CastTest, CastFloatToInt8) {
63-
int8_t output_data[6];
64-
int input_dims[] = {2, 3, 2};
65-
66-
// TODO(b/178391195): Test negative and out-of-range numbers.
67-
const float input_values[] = {100.f, 1.0f, 0.f, 0.4f, 1.999f, 1.1f};
68-
const int8_t golden[] = {100, 1, 0, 0, 1, 1};
69-
tflite::testing::TestCast(input_dims, input_values, golden, output_data);
97+
tflite::testing::TestFloatToInt<int8_t>(200.f, -200.f);
7098
}
7199

72100
TEST(CastTest, CastFloatToInt16) {
73-
int16_t output_data[6];
74-
int input_dims[] = {2, 3, 2};
101+
tflite::testing::TestFloatToInt<int16_t>(40000.f, -40000.f);
102+
}
103+
104+
TEST(CastTest, CastFloatToInt32) {
105+
tflite::testing::TestFloatToInt<int32_t>(1e15f, -1e15f);
106+
}
107+
108+
TEST(CastTest, CastFloatToUInt32) {
109+
tflite::testing::TestFloatToInt<uint32_t>(1e15f, -1e15f);
110+
}
75111

76-
// TODO(b/178391195): Test negative and out-of-range numbers.
77-
const float input_values[] = {100.f, 1.0f, 0.f, 0.4f, 1.999f, 1.1f};
78-
const int16_t golden[] = {100, 1, 0, 0, 1, 1};
79-
tflite::testing::TestCast(input_dims, input_values, golden, output_data);
112+
TEST(CastTest, CastFloatToBool) {
113+
const float input[] = {1.0f,
114+
0.0f,
115+
-1.0f,
116+
0.001f,
117+
std::numeric_limits<float>::infinity(),
118+
std::numeric_limits<float>::quiet_NaN()};
119+
const bool golden[] = {true, false, true, true, true, true};
120+
tflite::testing::TestCast(input, golden);
80121
}
81122

82123
TEST(CastTest, CastInt8ToFloat) {
83-
float output_data[6];
84-
int input_dims[] = {2, 3, 2};
85-
const int8_t input_values[] = {123, 0, 1, 2, 3, 4};
124+
const int8_t input[] = {123, 0, 1, 2, 3, 4};
86125
const float golden[] = {123.f, 0.f, 1.f, 2.f, 3.f, 4.f};
87-
tflite::testing::TestCast(input_dims, input_values, golden, output_data);
126+
tflite::testing::TestCast(input, golden);
88127
}
89128

90129
TEST(CastTest, CastInt16ToFloat) {
91-
float output_data[6];
92-
int input_dims[] = {2, 3, 2};
93-
const int16_t input_values[] = {123, 0, 1, 2, 3, 4};
130+
const int16_t input[] = {123, 0, 1, 2, 3, 4};
94131
const float golden[] = {123.f, 0.f, 1.f, 2.f, 3.f, 4.f};
95-
tflite::testing::TestCast(input_dims, input_values, golden, output_data);
132+
tflite::testing::TestCast(input, golden);
96133
}
97134

98135
TEST(CastTest, CastInt16ToInt32) {
99-
int32_t output_data[6];
100-
int input_dims[] = {2, 3, 2};
101-
const int16_t input_values[] = {123, 0, 1, 2, 3, 4};
136+
const int16_t input[] = {123, 0, 1, 2, 3, 4};
102137
const int32_t golden[] = {123, 0, 1, 2, 3, 4};
103-
tflite::testing::TestCast(input_dims, input_values, golden, output_data);
138+
tflite::testing::TestCast(input, golden);
104139
}
105140

106141
TEST(CastTest, CastInt32ToInt16) {
107-
int16_t output_data[6];
108-
int input_dims[] = {2, 3, 2};
109-
const int32_t input_values[] = {123, 0, 1, 2, 3, 4};
142+
const int32_t input[] = {123, 0, 1, 2, 3, 4};
110143
const int16_t golden[] = {123, 0, 1, 2, 3, 4};
111-
tflite::testing::TestCast(input_dims, input_values, golden, output_data);
144+
tflite::testing::TestCast(input, golden);
112145
}
113146

114147
TEST(CastTest, CastUInt32ToInt32) {
115-
int32_t output_data[6];
116-
int input_dims[] = {2, 2, 3};
117-
const uint32_t input_values[] = {100, 200, 300, 400, 500, 600};
148+
const uint32_t input[] = {100, 200, 300, 400, 500, 600};
118149
const int32_t golden[] = {100, 200, 300, 400, 500, 600};
119-
tflite::testing::TestCast(input_dims, input_values, golden, output_data);
150+
tflite::testing::TestCast(input, golden);
120151
}
121152

122153
TEST(CastTest, CastInt32ToUInt32) {
123-
uint32_t output_data[6];
124-
int input_dims[] = {2, 2, 3};
125-
const int32_t input_values[] = {100, 200, 300, 400, 500, 600};
154+
const int32_t input[] = {100, 200, 300, 400, 500, 600};
126155
const uint32_t golden[] = {100, 200, 300, 400, 500, 600};
127-
tflite::testing::TestCast(input_dims, input_values, golden, output_data);
156+
tflite::testing::TestCast(input, golden);
128157
}
129158

130159
TEST(CastTest, CastBoolToFloat) {
131-
float output_data[6];
132-
int input_dims[] = {2, 2, 3};
133-
const bool input_values[] = {true, true, false, true, false, true};
160+
const bool input[] = {true, true, false, true, false, true};
134161
const float golden[] = {1.f, 1.0f, 0.f, 1.0f, 0.0f, 1.0f};
135-
tflite::testing::TestCast(input_dims, input_values, golden, output_data);
162+
tflite::testing::TestCast(input, golden);
136163
}
137164

138165
TF_LITE_MICRO_TESTS_MAIN

0 commit comments

Comments
 (0)