Skip to content

Commit b702912

Browse files
pytorchbotshoumikhin
pytorchbot
and
shoumikhin
authored
Handle uint types. (#15060)
Summary: . Differential Revision: D84516559 Co-authored-by: Anthony Shoumikhin <[email protected]>
1 parent c34a19c commit b702912

File tree

4 files changed

+175
-10
lines changed

4 files changed

+175
-10
lines changed

‎extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ - (NSString *)description {
271271
ET_CHECK_MSG(false, "Unsupported dtype in description");
272272
}
273273
} ctx;
274-
ET_SWITCH_REALHBBF16_TYPES(
274+
ET_SWITCH_REALHBBF16_AND_UINT_TYPES(
275275
static_cast<ScalarType>(_tensor->scalar_type()),
276276
ctx,
277277
"description",

‎extension/tensor/tensor_ptr.h‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -123,13 +123,14 @@ inline TensorPtr make_tensor_ptr(
123123
}
124124
} ctx;
125125

126-
ET_SWITCH_REALHBBF16_TYPES(type, ctx, "make_tensor_ptr", CTYPE, [&] {
127-
std::transform(
128-
data.begin(),
129-
data.end(),
130-
reinterpret_cast<CTYPE*>(casted_data.data()),
131-
[](const T& val) { return static_cast<CTYPE>(val); });
132-
});
126+
ET_SWITCH_REALHBBF16_AND_UINT_TYPES(
127+
type, ctx, "make_tensor_ptr", CTYPE, [&] {
128+
std::transform(
129+
data.begin(),
130+
data.end(),
131+
reinterpret_cast<CTYPE*>(casted_data.data()),
132+
[](const T& val) { return static_cast<CTYPE>(val); });
133+
});
133134
const auto raw_data_ptr = casted_data.data();
134135
auto data_ptr =
135136
std::make_shared<std::vector<uint8_t>>(std::move(casted_data));

‎extension/tensor/tensor_ptr_maker.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ TensorPtr random_strided(
9696
}
9797
} ctx;
9898

99-
ET_SWITCH_REALHBBF16_TYPES(type, ctx, "random_strided", CTYPE, [&] {
99+
ET_SWITCH_REALHBBF16_AND_UINT_TYPES(type, ctx, "random_strided", CTYPE, [&] {
100100
std::generate_n(tensor->mutable_data_ptr<CTYPE>(), tensor->numel(), [&]() {
101101
return static_cast<CTYPE>(distribution(gen));
102102
});
@@ -138,7 +138,7 @@ TensorPtr full_strided(
138138
}
139139
} ctx;
140140

141-
ET_SWITCH_REALHBBF16_TYPES(type, ctx, "full_strided", CTYPE, [&] {
141+
ET_SWITCH_REALHBBF16_AND_UINT_TYPES(type, ctx, "full_strided", CTYPE, [&] {
142142
CTYPE value;
143143
ET_EXTRACT_SCALAR(fill_value, value);
144144
std::fill(

‎extension/tensor/test/tensor_ptr_test.cpp‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -893,3 +893,167 @@ TEST_F(TensorPtrTest, TensorDataCastingInvalidCast) {
893893
},
894894
"");
895895
}
896+
897+
TEST_F(TensorPtrTest, TensorDataOnlyUInt16Type) {
898+
std::vector<uint16_t> data = {1u, 65535u, 42u, 0u};
899+
auto tensor = make_tensor_ptr(std::move(data));
900+
EXPECT_EQ(tensor->dim(), 1);
901+
EXPECT_EQ(tensor->size(0), 4);
902+
EXPECT_EQ(tensor->strides()[0], 1);
903+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt16);
904+
auto ptr = tensor->const_data_ptr<uint16_t>();
905+
EXPECT_EQ(ptr[0], 1u);
906+
EXPECT_EQ(ptr[1], 65535u);
907+
EXPECT_EQ(ptr[2], 42u);
908+
EXPECT_EQ(ptr[3], 0u);
909+
}
910+
911+
TEST_F(TensorPtrTest, TensorDataOnlyUInt32Type) {
912+
std::vector<uint32_t> data = {0u, 123u, 4000000000u};
913+
auto tensor = make_tensor_ptr(std::move(data));
914+
EXPECT_EQ(tensor->dim(), 1);
915+
EXPECT_EQ(tensor->size(0), 3);
916+
EXPECT_EQ(tensor->strides()[0], 1);
917+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt32);
918+
auto ptr = tensor->const_data_ptr<uint32_t>();
919+
EXPECT_EQ(ptr[0], 0u);
920+
EXPECT_EQ(ptr[1], 123u);
921+
EXPECT_EQ(ptr[2], 4000000000u);
922+
}
923+
924+
TEST_F(TensorPtrTest, TensorDataOnlyUInt64Type) {
925+
std::vector<uint64_t> data = {0ull, 1ull, 9000000000000000000ull};
926+
auto tensor = make_tensor_ptr(std::move(data));
927+
EXPECT_EQ(tensor->dim(), 1);
928+
EXPECT_EQ(tensor->size(0), 3);
929+
EXPECT_EQ(tensor->strides()[0], 1);
930+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt64);
931+
auto ptr = tensor->const_data_ptr<uint64_t>();
932+
EXPECT_EQ(ptr[0], 0ull);
933+
EXPECT_EQ(ptr[1], 1ull);
934+
EXPECT_EQ(ptr[2], 9000000000000000000ull);
935+
}
936+
937+
TEST_F(TensorPtrTest, TensorUint8dataUInt32Type) {
938+
std::vector<uint32_t> values = {1u, 4000000000u, 123u};
939+
const auto* bytes = reinterpret_cast<const uint8_t*>(values.data());
940+
std::vector<uint8_t> raw(bytes, bytes + values.size() * sizeof(uint32_t));
941+
auto tensor = make_tensor_ptr(
942+
{3}, std::move(raw), executorch::aten::ScalarType::UInt32);
943+
EXPECT_EQ(tensor->dim(), 1);
944+
EXPECT_EQ(tensor->size(0), 3);
945+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt32);
946+
auto ptr = tensor->const_data_ptr<uint32_t>();
947+
EXPECT_EQ(ptr[0], 1u);
948+
EXPECT_EQ(ptr[1], 4000000000u);
949+
EXPECT_EQ(ptr[2], 123u);
950+
}
951+
952+
TEST_F(TensorPtrTest, TensorUint8dataUInt64Type) {
953+
std::vector<uint64_t> values = {0ull, 42ull, 9000000000000000000ull};
954+
const auto* bytes = reinterpret_cast<const uint8_t*>(values.data());
955+
std::vector<uint8_t> raw(bytes, bytes + values.size() * sizeof(uint64_t));
956+
auto tensor = make_tensor_ptr(
957+
{3}, std::move(raw), executorch::aten::ScalarType::UInt64);
958+
EXPECT_EQ(tensor->dim(), 1);
959+
EXPECT_EQ(tensor->size(0), 3);
960+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt64);
961+
auto ptr = tensor->const_data_ptr<uint64_t>();
962+
EXPECT_EQ(ptr[0], 0ull);
963+
EXPECT_EQ(ptr[1], 42ull);
964+
EXPECT_EQ(ptr[2], 9000000000000000000ull);
965+
}
966+
967+
TEST_F(TensorPtrTest, TensorUint8dataSizeMismatchUInt32ExpectDeath) {
968+
std::vector<uint8_t> data(
969+
3 * executorch::aten::elementSize(executorch::aten::ScalarType::UInt32) -
970+
1);
971+
ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({3}, std::move(data)); }, "");
972+
}
973+
974+
TEST_F(TensorPtrTest, TensorUint8dataSizeMismatchUInt64ExpectDeath) {
975+
std::vector<uint8_t> data(
976+
2 * executorch::aten::elementSize(executorch::aten::ScalarType::UInt64) +
977+
1);
978+
ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({2}, std::move(data)); }, "");
979+
}
980+
981+
TEST_F(TensorPtrTest, TensorDataCastingFromInt32ToUInt16) {
982+
std::vector<int32_t> data = {-1, 65535, 65536, -65536};
983+
auto tensor =
984+
make_tensor_ptr(std::move(data), executorch::aten::ScalarType::UInt16);
985+
EXPECT_EQ(tensor->dim(), 1);
986+
EXPECT_EQ(tensor->size(0), 4);
987+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt16);
988+
auto ptr = tensor->const_data_ptr<uint16_t>();
989+
EXPECT_EQ(ptr[0], static_cast<uint16_t>(-1));
990+
EXPECT_EQ(ptr[1], static_cast<uint16_t>(65535));
991+
EXPECT_EQ(ptr[2], static_cast<uint16_t>(65536));
992+
EXPECT_EQ(ptr[3], static_cast<uint16_t>(-65536));
993+
}
994+
995+
TEST_F(TensorPtrTest, TensorDataCastingFromUInt32ToFloat) {
996+
std::vector<uint32_t> data = {0u, 123u, 4000000000u};
997+
auto tensor =
998+
make_tensor_ptr(std::move(data), executorch::aten::ScalarType::Float);
999+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float);
1000+
auto ptr = tensor->const_data_ptr<float>();
1001+
EXPECT_FLOAT_EQ(ptr[0], 0.0f);
1002+
EXPECT_FLOAT_EQ(ptr[1], 123.0f);
1003+
EXPECT_FLOAT_EQ(ptr[2], 4000000000.0f);
1004+
}
1005+
1006+
TEST_F(TensorPtrTest, TensorDataCastingFromFloatToUInt32) {
1007+
std::vector<float> data = {1.0f, 2.0f};
1008+
auto tensor =
1009+
make_tensor_ptr(std::move(data), executorch::aten::ScalarType::UInt32);
1010+
1011+
EXPECT_EQ(tensor->dim(), 1);
1012+
EXPECT_EQ(tensor->size(0), 2);
1013+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt32);
1014+
1015+
auto ptr = tensor->const_data_ptr<uint32_t>();
1016+
EXPECT_EQ(ptr[0], 1u);
1017+
EXPECT_EQ(ptr[1], 2u);
1018+
}
1019+
1020+
TEST_F(TensorPtrTest, MakeTensorPtrFromExistingTensorUInt32) {
1021+
std::vector<uint32_t> data = {10u, 20u, 30u, 40u};
1022+
auto tensor = make_tensor_ptr({2, 2}, data);
1023+
auto alias = make_tensor_ptr(tensor);
1024+
EXPECT_EQ(alias->dim(), 2);
1025+
EXPECT_EQ(alias->size(0), 2);
1026+
EXPECT_EQ(alias->size(1), 2);
1027+
EXPECT_EQ(alias->scalar_type(), executorch::aten::ScalarType::UInt32);
1028+
EXPECT_EQ(
1029+
alias->const_data_ptr<uint32_t>(), tensor->const_data_ptr<uint32_t>());
1030+
}
1031+
1032+
TEST_F(TensorPtrTest, CloneTensorPtrFromExistingTensorUInt32) {
1033+
std::vector<uint32_t> data = {10u, 20u, 30u, 40u};
1034+
auto tensor = make_tensor_ptr({2, 2}, std::move(data));
1035+
auto cloned = clone_tensor_ptr(tensor);
1036+
EXPECT_EQ(cloned->dim(), 2);
1037+
EXPECT_EQ(cloned->size(0), 2);
1038+
EXPECT_EQ(cloned->size(1), 2);
1039+
EXPECT_EQ(cloned->scalar_type(), executorch::aten::ScalarType::UInt32);
1040+
EXPECT_NE(
1041+
cloned->const_data_ptr<uint32_t>(), tensor->const_data_ptr<uint32_t>());
1042+
auto ptr = cloned->const_data_ptr<uint32_t>();
1043+
EXPECT_EQ(ptr[0], 10u);
1044+
EXPECT_EQ(ptr[3], 40u);
1045+
}
1046+
1047+
TEST_F(TensorPtrTest, Tensor2DUInt16OwningData) {
1048+
std::vector<uint16_t> data = {1u, 2u, 3u, 4u, 5u, 6u};
1049+
auto tensor = make_tensor_ptr({2, 3}, std::move(data));
1050+
EXPECT_EQ(tensor->dim(), 2);
1051+
EXPECT_EQ(tensor->size(0), 2);
1052+
EXPECT_EQ(tensor->size(1), 3);
1053+
EXPECT_EQ(tensor->strides()[0], 3);
1054+
EXPECT_EQ(tensor->strides()[1], 1);
1055+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt16);
1056+
auto ptr = tensor->const_data_ptr<uint16_t>();
1057+
EXPECT_EQ(ptr[0], 1u);
1058+
EXPECT_EQ(ptr[5], 6u);
1059+
}

0 commit comments

Comments
 (0)