Skip to content

Commit 90237aa

Browse files
HyukjinKwonHyukjinKwon
committed
[C++] Support half-float tensors in equality comparison
1 parent d2315fe commit 90237aa

2 files changed

Lines changed: 79 additions & 17 deletions

File tree

‎cpp/src/arrow/compare.cc‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1249,8 +1249,9 @@ bool StridedFloatTensorContentEquals(const int dim_index, int64_t left_offset,
12491249
int64_t right_offset, const Tensor& left,
12501250
const Tensor& right, const EqualOptions& opts) {
12511251
using c_type = typename DataType::c_type;
1252-
static_assert(std::is_floating_point<c_type>::value,
1253-
"DataType must be a floating point type");
1252+
static_assert(std::is_floating_point<c_type>::value ||
1253+
std::is_same<DataType, HalfFloatType>::value,
1254+
"DataType must be a floating point type or half-float");
12541255

12551256
const auto n = left.shape()[dim_index];
12561257
const auto left_stride = left.strides()[dim_index];
@@ -1309,8 +1310,9 @@ bool TensorEquals(const Tensor& left, const Tensor& right, const EqualOptions& o
13091310
}
13101311

13111312
switch (left.type_id()) {
1312-
// TODO: Support half-float tensors
1313-
// case Type::HALF_FLOAT:
1313+
case Type::HALF_FLOAT:
1314+
return FloatTensorEquals<HalfFloatType>(left, right, opts);
1315+
13141316
case Type::FLOAT:
13151317
return FloatTensorEquals<FloatType>(left, right, opts);
13161318

@@ -1347,8 +1349,9 @@ bool FloatSparseTensorDataEquals(const typename DataType::c_type* left_data,
13471349
const typename DataType::c_type* right_data,
13481350
const int64_t length, const EqualOptions& opts) {
13491351
using c_type = typename DataType::c_type;
1350-
static_assert(std::is_floating_point<c_type>::value,
1351-
"DataType must be a floating point type");
1352+
static_assert(std::is_floating_point<c_type>::value ||
1353+
std::is_same<DataType, HalfFloatType>::value,
1354+
"DataType must be a floating point type or half-float");
13521355
if (opts.nans_equal()) {
13531356
if (left_data == right_data) {
13541357
return true;
@@ -1395,8 +1398,11 @@ struct SparseTensorEqualsImpl<SparseIndexType, SparseIndexType> {
13951398
const uint8_t* left_data = left.data()->data();
13961399
const uint8_t* right_data = right.data()->data();
13971400
switch (left.type()->id()) {
1398-
// TODO: Support half-float tensors
1399-
// case Type::HALF_FLOAT:
1401+
case Type::HALF_FLOAT:
1402+
return FloatSparseTensorDataEquals<HalfFloatType>(
1403+
reinterpret_cast<const uint16_t*>(left_data),
1404+
reinterpret_cast<const uint16_t*>(right_data), length, opts);
1405+
14001406
case Type::FLOAT:
14011407
return FloatSparseTensorDataEquals<FloatType>(
14021408
reinterpret_cast<const float*>(left_data),

‎cpp/src/arrow/tensor_test.cc‎

Lines changed: 65 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
#include "arrow/tensor.h"
3232
#include "arrow/testing/gtest_util.h"
3333
#include "arrow/type.h"
34+
#include "arrow/util/float16.h"
3435

3536
namespace arrow {
3637

@@ -505,8 +506,18 @@ TYPED_TEST_P(TestFloatTensor, Equals) {
505506

506507
std::vector<int64_t> shape = {4, 4};
507508

508-
std::vector<c_data_type> c_values = {1, 2, 3, 4, 5, 6, 7, 8,
509+
constexpr int c_vals_as_int_arr[] = {1, 2, 3, 4, 5, 6, 7, 8,
509510
9, 10, 11, 12, 13, 14, 15, 16};
511+
std::vector<c_data_type> c_values;
512+
if constexpr (std::is_same_v<DataType, HalfFloatType>) {
513+
for (int i : c_vals_as_int_arr) {
514+
c_values.push_back(util::Float16::FromFloat(static_cast<float>(i)).bits());
515+
}
516+
} else {
517+
for (int i : c_vals_as_int_arr) {
518+
c_values.push_back(static_cast<c_data_type>(i));
519+
}
520+
}
510521
std::vector<int64_t> c_strides = {unit_size * shape[1], unit_size};
511522
Tensor tc1(TypeTraits<DataType>::type_singleton(), Buffer::Wrap(c_values), shape,
512523
c_strides);
@@ -515,8 +526,18 @@ TYPED_TEST_P(TestFloatTensor, Equals) {
515526
Tensor tc2(TypeTraits<DataType>::type_singleton(), Buffer::Wrap(c_values_2), shape,
516527
c_strides);
517528

518-
std::vector<c_data_type> f_values = {1, 5, 9, 13, 2, 6, 10, 14,
529+
constexpr int f_vals_as_int_arr[] = {1, 5, 9, 13, 2, 6, 10, 14,
519530
3, 7, 11, 15, 4, 8, 12, 16};
531+
std::vector<c_data_type> f_values;
532+
if constexpr (std::is_same_v<DataType, HalfFloatType>) {
533+
for (int i : f_vals_as_int_arr) {
534+
f_values.push_back(util::Float16::FromFloat(static_cast<float>(i)).bits());
535+
}
536+
} else {
537+
for (int i : f_vals_as_int_arr) {
538+
f_values.push_back(static_cast<c_data_type>(i));
539+
}
540+
}
520541
Tensor tc3(TypeTraits<DataType>::type_singleton(), Buffer::Wrap(f_values), shape,
521542
c_strides);
522543

@@ -531,9 +552,19 @@ TYPED_TEST_P(TestFloatTensor, Equals) {
531552
Tensor tf3(TypeTraits<DataType>::type_singleton(), Buffer::Wrap(c_values), shape,
532553
f_strides);
533554

534-
std::vector<c_data_type> nc_values = {1, 0, 5, 0, 9, 0, 13, 0, 2, 0, 6,
555+
constexpr int nc_vals_as_int_arr[] = {1, 0, 5, 0, 9, 0, 13, 0, 2, 0, 6,
535556
0, 10, 0, 14, 0, 3, 0, 7, 0, 11, 0,
536557
15, 0, 4, 0, 8, 0, 12, 0, 16, 0};
558+
std::vector<c_data_type> nc_values;
559+
if constexpr (std::is_same_v<DataType, HalfFloatType>) {
560+
for (int i : nc_vals_as_int_arr) {
561+
nc_values.push_back(util::Float16::FromFloat(static_cast<float>(i)).bits());
562+
}
563+
} else {
564+
for (int i : nc_vals_as_int_arr) {
565+
nc_values.push_back(static_cast<c_data_type>(i));
566+
}
567+
}
537568
std::vector<int64_t> nc_strides = {unit_size * 2, unit_size * 2 * shape[0]};
538569
Tensor tnc(TypeTraits<DataType>::type_singleton(), Buffer::Wrap(nc_values), shape,
539570
nc_strides);
@@ -569,29 +600,54 @@ TYPED_TEST_P(TestFloatTensor, Equals) {
569600
EXPECT_FALSE(tf3.Equals(tnc));
570601

571602
// signed zeros
572-
c_values[0] = -0.0;
573-
c_values_2[0] = 0.0;
603+
if constexpr (std::is_same_v<DataType, HalfFloatType>) {
604+
c_values[0] = 0x8000; // -0.0
605+
c_values_2[0] = 0x0000; // +0.0
606+
} else {
607+
c_values[0] = -0.0;
608+
c_values_2[0] = 0.0;
609+
}
574610
EXPECT_TRUE(tc1.Equals(tc2));
575611
EXPECT_FALSE(tc1.Equals(tc2, EqualOptions().signed_zeros_equal(false)));
576612

577613
// tensors with NaNs
578-
const c_data_type nan_value = static_cast<c_data_type>(NAN);
614+
c_data_type nan_value;
615+
if constexpr (std::is_same_v<DataType, HalfFloatType>) {
616+
nan_value = 0x7e00; // NaN
617+
} else {
618+
nan_value = static_cast<c_data_type>(NAN);
619+
}
620+
579621
c_values[0] = nan_value;
580-
EXPECT_TRUE(std::isnan(tc1.Value<DataType>({0, 0})));
622+
if constexpr (std::is_same_v<DataType, HalfFloatType>) {
623+
EXPECT_TRUE(util::Float16::FromBits(tc1.Value<DataType>({0, 0})).is_nan());
624+
} else {
625+
EXPECT_TRUE(std::isnan(tc1.Value<DataType>({0, 0})));
626+
}
581627
EXPECT_FALSE(tc1.Equals(tc1)); // same object
582628
EXPECT_TRUE(tc1.Equals(tc1, EqualOptions().nans_equal(true))); // same object
583-
EXPECT_FALSE(std::isnan(tc2.Value<DataType>({0, 0})));
629+
630+
if constexpr (std::is_same_v<DataType, HalfFloatType>) {
631+
EXPECT_FALSE(util::Float16::FromBits(tc2.Value<DataType>({0, 0})).is_nan());
632+
} else {
633+
EXPECT_FALSE(std::isnan(tc2.Value<DataType>({0, 0})));
634+
}
584635
EXPECT_FALSE(tc1.Equals(tc2)); // different memory
585636
EXPECT_FALSE(tc1.Equals(tc2, EqualOptions().nans_equal(true))); // different memory
586637

587638
c_values_2[0] = nan_value;
588-
EXPECT_TRUE(std::isnan(tc2.Value<DataType>({0, 0})));
639+
if constexpr (std::is_same_v<DataType, HalfFloatType>) {
640+
EXPECT_TRUE(util::Float16::FromBits(tc2.Value<DataType>({0, 0})).is_nan());
641+
} else {
642+
EXPECT_TRUE(std::isnan(tc2.Value<DataType>({0, 0})));
643+
}
589644
EXPECT_FALSE(tc1.Equals(tc2)); // different memory
590645
EXPECT_TRUE(tc1.Equals(tc2, EqualOptions().nans_equal(true))); // different memory
591646
}
592647

593648
REGISTER_TYPED_TEST_SUITE_P(TestFloatTensor, Equals);
594649

650+
INSTANTIATE_TYPED_TEST_SUITE_P(Float16, TestFloatTensor, HalfFloatType);
595651
INSTANTIATE_TYPED_TEST_SUITE_P(Float32, TestFloatTensor, FloatType);
596652
INSTANTIATE_TYPED_TEST_SUITE_P(Float64, TestFloatTensor, DoubleType);
597653

0 commit comments

Comments
 (0)