Skip to content

Commit 9ccc24c

Browse files
shoumikhinpytorchbot
shoumikhin
authored and
pytorchbot
committed
Handle uint types. (#15055)
Summary: . Differential Revision: D84516559 (cherry picked from commit f19882b)
1 parent 61866dc commit 9ccc24c

4 files changed

Lines changed: 175 additions & 10 deletions

File tree

‎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
@@ -825,3 +825,167 @@ TEST_F(TensorPtrTest, TensorDataCastingInvalidCast) {
825825
},
826826
"");
827827
}
828+
829+
TEST_F(TensorPtrTest, TensorDataOnlyUInt16Type) {
830+
std::vector<uint16_t> data = {1u, 65535u, 42u, 0u};
831+
auto tensor = make_tensor_ptr(std::move(data));
832+
EXPECT_EQ(tensor->dim(), 1);
833+
EXPECT_EQ(tensor->size(0), 4);
834+
EXPECT_EQ(tensor->strides()[0], 1);
835+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt16);
836+
auto ptr = tensor->const_data_ptr<uint16_t>();
837+
EXPECT_EQ(ptr[0], 1u);
838+
EXPECT_EQ(ptr[1], 65535u);
839+
EXPECT_EQ(ptr[2], 42u);
840+
EXPECT_EQ(ptr[3], 0u);
841+
}
842+
843+
TEST_F(TensorPtrTest, TensorDataOnlyUInt32Type) {
844+
std::vector<uint32_t> data = {0u, 123u, 4000000000u};
845+
auto tensor = make_tensor_ptr(std::move(data));
846+
EXPECT_EQ(tensor->dim(), 1);
847+
EXPECT_EQ(tensor->size(0), 3);
848+
EXPECT_EQ(tensor->strides()[0], 1);
849+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt32);
850+
auto ptr = tensor->const_data_ptr<uint32_t>();
851+
EXPECT_EQ(ptr[0], 0u);
852+
EXPECT_EQ(ptr[1], 123u);
853+
EXPECT_EQ(ptr[2], 4000000000u);
854+
}
855+
856+
TEST_F(TensorPtrTest, TensorDataOnlyUInt64Type) {
857+
std::vector<uint64_t> data = {0ull, 1ull, 9000000000000000000ull};
858+
auto tensor = make_tensor_ptr(std::move(data));
859+
EXPECT_EQ(tensor->dim(), 1);
860+
EXPECT_EQ(tensor->size(0), 3);
861+
EXPECT_EQ(tensor->strides()[0], 1);
862+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt64);
863+
auto ptr = tensor->const_data_ptr<uint64_t>();
864+
EXPECT_EQ(ptr[0], 0ull);
865+
EXPECT_EQ(ptr[1], 1ull);
866+
EXPECT_EQ(ptr[2], 9000000000000000000ull);
867+
}
868+
869+
TEST_F(TensorPtrTest, TensorUint8dataUInt32Type) {
870+
std::vector<uint32_t> values = {1u, 4000000000u, 123u};
871+
const auto* bytes = reinterpret_cast<const uint8_t*>(values.data());
872+
std::vector<uint8_t> raw(bytes, bytes + values.size() * sizeof(uint32_t));
873+
auto tensor = make_tensor_ptr(
874+
{3}, std::move(raw), executorch::aten::ScalarType::UInt32);
875+
EXPECT_EQ(tensor->dim(), 1);
876+
EXPECT_EQ(tensor->size(0), 3);
877+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt32);
878+
auto ptr = tensor->const_data_ptr<uint32_t>();
879+
EXPECT_EQ(ptr[0], 1u);
880+
EXPECT_EQ(ptr[1], 4000000000u);
881+
EXPECT_EQ(ptr[2], 123u);
882+
}
883+
884+
TEST_F(TensorPtrTest, TensorUint8dataUInt64Type) {
885+
std::vector<uint64_t> values = {0ull, 42ull, 9000000000000000000ull};
886+
const auto* bytes = reinterpret_cast<const uint8_t*>(values.data());
887+
std::vector<uint8_t> raw(bytes, bytes + values.size() * sizeof(uint64_t));
888+
auto tensor = make_tensor_ptr(
889+
{3}, std::move(raw), executorch::aten::ScalarType::UInt64);
890+
EXPECT_EQ(tensor->dim(), 1);
891+
EXPECT_EQ(tensor->size(0), 3);
892+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt64);
893+
auto ptr = tensor->const_data_ptr<uint64_t>();
894+
EXPECT_EQ(ptr[0], 0ull);
895+
EXPECT_EQ(ptr[1], 42ull);
896+
EXPECT_EQ(ptr[2], 9000000000000000000ull);
897+
}
898+
899+
TEST_F(TensorPtrTest, TensorUint8dataSizeMismatchUInt32ExpectDeath) {
900+
std::vector<uint8_t> data(
901+
3 * executorch::aten::elementSize(executorch::aten::ScalarType::UInt32) -
902+
1);
903+
ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({3}, std::move(data)); }, "");
904+
}
905+
906+
TEST_F(TensorPtrTest, TensorUint8dataSizeMismatchUInt64ExpectDeath) {
907+
std::vector<uint8_t> data(
908+
2 * executorch::aten::elementSize(executorch::aten::ScalarType::UInt64) +
909+
1);
910+
ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({2}, std::move(data)); }, "");
911+
}
912+
913+
TEST_F(TensorPtrTest, TensorDataCastingFromInt32ToUInt16) {
914+
std::vector<int32_t> data = {-1, 65535, 65536, -65536};
915+
auto tensor =
916+
make_tensor_ptr(std::move(data), executorch::aten::ScalarType::UInt16);
917+
EXPECT_EQ(tensor->dim(), 1);
918+
EXPECT_EQ(tensor->size(0), 4);
919+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt16);
920+
auto ptr = tensor->const_data_ptr<uint16_t>();
921+
EXPECT_EQ(ptr[0], static_cast<uint16_t>(-1));
922+
EXPECT_EQ(ptr[1], static_cast<uint16_t>(65535));
923+
EXPECT_EQ(ptr[2], static_cast<uint16_t>(65536));
924+
EXPECT_EQ(ptr[3], static_cast<uint16_t>(-65536));
925+
}
926+
927+
TEST_F(TensorPtrTest, TensorDataCastingFromUInt32ToFloat) {
928+
std::vector<uint32_t> data = {0u, 123u, 4000000000u};
929+
auto tensor =
930+
make_tensor_ptr(std::move(data), executorch::aten::ScalarType::Float);
931+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float);
932+
auto ptr = tensor->const_data_ptr<float>();
933+
EXPECT_FLOAT_EQ(ptr[0], 0.0f);
934+
EXPECT_FLOAT_EQ(ptr[1], 123.0f);
935+
EXPECT_FLOAT_EQ(ptr[2], 4000000000.0f);
936+
}
937+
938+
TEST_F(TensorPtrTest, TensorDataCastingFromFloatToUInt32) {
939+
std::vector<float> data = {1.0f, 2.0f};
940+
auto tensor =
941+
make_tensor_ptr(std::move(data), executorch::aten::ScalarType::UInt32);
942+
943+
EXPECT_EQ(tensor->dim(), 1);
944+
EXPECT_EQ(tensor->size(0), 2);
945+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt32);
946+
947+
auto ptr = tensor->const_data_ptr<uint32_t>();
948+
EXPECT_EQ(ptr[0], 1u);
949+
EXPECT_EQ(ptr[1], 2u);
950+
}
951+
952+
TEST_F(TensorPtrTest, MakeTensorPtrFromExistingTensorUInt32) {
953+
std::vector<uint32_t> data = {10u, 20u, 30u, 40u};
954+
auto tensor = make_tensor_ptr({2, 2}, data);
955+
auto alias = make_tensor_ptr(tensor);
956+
EXPECT_EQ(alias->dim(), 2);
957+
EXPECT_EQ(alias->size(0), 2);
958+
EXPECT_EQ(alias->size(1), 2);
959+
EXPECT_EQ(alias->scalar_type(), executorch::aten::ScalarType::UInt32);
960+
EXPECT_EQ(
961+
alias->const_data_ptr<uint32_t>(), tensor->const_data_ptr<uint32_t>());
962+
}
963+
964+
TEST_F(TensorPtrTest, CloneTensorPtrFromExistingTensorUInt32) {
965+
std::vector<uint32_t> data = {10u, 20u, 30u, 40u};
966+
auto tensor = make_tensor_ptr({2, 2}, std::move(data));
967+
auto cloned = clone_tensor_ptr(tensor);
968+
EXPECT_EQ(cloned->dim(), 2);
969+
EXPECT_EQ(cloned->size(0), 2);
970+
EXPECT_EQ(cloned->size(1), 2);
971+
EXPECT_EQ(cloned->scalar_type(), executorch::aten::ScalarType::UInt32);
972+
EXPECT_NE(
973+
cloned->const_data_ptr<uint32_t>(), tensor->const_data_ptr<uint32_t>());
974+
auto ptr = cloned->const_data_ptr<uint32_t>();
975+
EXPECT_EQ(ptr[0], 10u);
976+
EXPECT_EQ(ptr[3], 40u);
977+
}
978+
979+
TEST_F(TensorPtrTest, Tensor2DUInt16OwningData) {
980+
std::vector<uint16_t> data = {1u, 2u, 3u, 4u, 5u, 6u};
981+
auto tensor = make_tensor_ptr({2, 3}, std::move(data));
982+
EXPECT_EQ(tensor->dim(), 2);
983+
EXPECT_EQ(tensor->size(0), 2);
984+
EXPECT_EQ(tensor->size(1), 3);
985+
EXPECT_EQ(tensor->strides()[0], 3);
986+
EXPECT_EQ(tensor->strides()[1], 1);
987+
EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt16);
988+
auto ptr = tensor->const_data_ptr<uint16_t>();
989+
EXPECT_EQ(ptr[0], 1u);
990+
EXPECT_EQ(ptr[5], 6u);
991+
}

0 commit comments

Comments
 (0)