diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm index 3a2b640b7d7..fd3cd3b1134 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm @@ -271,7 +271,7 @@ - (NSString *)description { ET_CHECK_MSG(false, "Unsupported dtype in description"); } } ctx; - ET_SWITCH_REALHBBF16_TYPES( + ET_SWITCH_REALHBBF16_AND_UINT_TYPES( static_cast(_tensor->scalar_type()), ctx, "description", diff --git a/extension/tensor/tensor_ptr.h b/extension/tensor/tensor_ptr.h index 900252109d3..92893b48158 100644 --- a/extension/tensor/tensor_ptr.h +++ b/extension/tensor/tensor_ptr.h @@ -123,13 +123,14 @@ inline TensorPtr make_tensor_ptr( } } ctx; - ET_SWITCH_REALHBBF16_TYPES(type, ctx, "make_tensor_ptr", CTYPE, [&] { - std::transform( - data.begin(), - data.end(), - reinterpret_cast(casted_data.data()), - [](const T& val) { return static_cast(val); }); - }); + ET_SWITCH_REALHBBF16_AND_UINT_TYPES( + type, ctx, "make_tensor_ptr", CTYPE, [&] { + std::transform( + data.begin(), + data.end(), + reinterpret_cast(casted_data.data()), + [](const T& val) { return static_cast(val); }); + }); const auto raw_data_ptr = casted_data.data(); auto data_ptr = std::make_shared>(std::move(casted_data)); diff --git a/extension/tensor/tensor_ptr_maker.cpp b/extension/tensor/tensor_ptr_maker.cpp index 511b0ebe582..b71dfab8eeb 100644 --- a/extension/tensor/tensor_ptr_maker.cpp +++ b/extension/tensor/tensor_ptr_maker.cpp @@ -96,7 +96,7 @@ TensorPtr random_strided( } } ctx; - ET_SWITCH_REALHBBF16_TYPES(type, ctx, "random_strided", CTYPE, [&] { + ET_SWITCH_REALHBBF16_AND_UINT_TYPES(type, ctx, "random_strided", CTYPE, [&] { std::generate_n(tensor->mutable_data_ptr(), tensor->numel(), [&]() { return static_cast(distribution(gen)); }); @@ -138,7 +138,7 @@ TensorPtr full_strided( } } ctx; - ET_SWITCH_REALHBBF16_TYPES(type, ctx, "full_strided", CTYPE, [&] { + ET_SWITCH_REALHBBF16_AND_UINT_TYPES(type, ctx, "full_strided", CTYPE, [&] { CTYPE value; ET_EXTRACT_SCALAR(fill_value, value); std::fill( diff --git a/extension/tensor/test/tensor_ptr_test.cpp b/extension/tensor/test/tensor_ptr_test.cpp index 9156a0c4b10..4a765625934 100644 --- a/extension/tensor/test/tensor_ptr_test.cpp +++ b/extension/tensor/test/tensor_ptr_test.cpp @@ -1073,3 +1073,167 @@ TEST_F(TensorPtrTest, TensorDataCastingInvalidCast) { }, ""); } + +TEST_F(TensorPtrTest, TensorDataOnlyUInt16Type) { + std::vector data = {1u, 65535u, 42u, 0u}; + auto tensor = make_tensor_ptr(std::move(data)); + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 4); + EXPECT_EQ(tensor->strides()[0], 1); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt16); + auto ptr = tensor->const_data_ptr(); + EXPECT_EQ(ptr[0], 1u); + EXPECT_EQ(ptr[1], 65535u); + EXPECT_EQ(ptr[2], 42u); + EXPECT_EQ(ptr[3], 0u); +} + +TEST_F(TensorPtrTest, TensorDataOnlyUInt32Type) { + std::vector data = {0u, 123u, 4000000000u}; + auto tensor = make_tensor_ptr(std::move(data)); + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 3); + EXPECT_EQ(tensor->strides()[0], 1); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt32); + auto ptr = tensor->const_data_ptr(); + EXPECT_EQ(ptr[0], 0u); + EXPECT_EQ(ptr[1], 123u); + EXPECT_EQ(ptr[2], 4000000000u); +} + +TEST_F(TensorPtrTest, TensorDataOnlyUInt64Type) { + std::vector data = {0ull, 1ull, 9000000000000000000ull}; + auto tensor = make_tensor_ptr(std::move(data)); + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 3); + EXPECT_EQ(tensor->strides()[0], 1); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt64); + auto ptr = tensor->const_data_ptr(); + EXPECT_EQ(ptr[0], 0ull); + EXPECT_EQ(ptr[1], 1ull); + EXPECT_EQ(ptr[2], 9000000000000000000ull); +} + +TEST_F(TensorPtrTest, TensorUint8dataUInt32Type) { + std::vector values = {1u, 4000000000u, 123u}; + const auto* bytes = reinterpret_cast(values.data()); + std::vector raw(bytes, bytes + values.size() * sizeof(uint32_t)); + auto tensor = make_tensor_ptr( + {3}, std::move(raw), executorch::aten::ScalarType::UInt32); + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 3); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt32); + auto ptr = tensor->const_data_ptr(); + EXPECT_EQ(ptr[0], 1u); + EXPECT_EQ(ptr[1], 4000000000u); + EXPECT_EQ(ptr[2], 123u); +} + +TEST_F(TensorPtrTest, TensorUint8dataUInt64Type) { + std::vector values = {0ull, 42ull, 9000000000000000000ull}; + const auto* bytes = reinterpret_cast(values.data()); + std::vector raw(bytes, bytes + values.size() * sizeof(uint64_t)); + auto tensor = make_tensor_ptr( + {3}, std::move(raw), executorch::aten::ScalarType::UInt64); + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 3); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt64); + auto ptr = tensor->const_data_ptr(); + EXPECT_EQ(ptr[0], 0ull); + EXPECT_EQ(ptr[1], 42ull); + EXPECT_EQ(ptr[2], 9000000000000000000ull); +} + +TEST_F(TensorPtrTest, TensorUint8dataSizeMismatchUInt32ExpectDeath) { + std::vector data( + 3 * executorch::aten::elementSize(executorch::aten::ScalarType::UInt32) - + 1); + ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({3}, std::move(data)); }, ""); +} + +TEST_F(TensorPtrTest, TensorUint8dataSizeMismatchUInt64ExpectDeath) { + std::vector data( + 2 * executorch::aten::elementSize(executorch::aten::ScalarType::UInt64) + + 1); + ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({2}, std::move(data)); }, ""); +} + +TEST_F(TensorPtrTest, TensorDataCastingFromInt32ToUInt16) { + std::vector data = {-1, 65535, 65536, -65536}; + auto tensor = + make_tensor_ptr(std::move(data), executorch::aten::ScalarType::UInt16); + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 4); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt16); + auto ptr = tensor->const_data_ptr(); + EXPECT_EQ(ptr[0], static_cast(-1)); + EXPECT_EQ(ptr[1], static_cast(65535)); + EXPECT_EQ(ptr[2], static_cast(65536)); + EXPECT_EQ(ptr[3], static_cast(-65536)); +} + +TEST_F(TensorPtrTest, TensorDataCastingFromUInt32ToFloat) { + std::vector data = {0u, 123u, 4000000000u}; + auto tensor = + make_tensor_ptr(std::move(data), executorch::aten::ScalarType::Float); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::Float); + auto ptr = tensor->const_data_ptr(); + EXPECT_FLOAT_EQ(ptr[0], 0.0f); + EXPECT_FLOAT_EQ(ptr[1], 123.0f); + EXPECT_FLOAT_EQ(ptr[2], 4000000000.0f); +} + +TEST_F(TensorPtrTest, TensorDataCastingFromFloatToUInt32) { + std::vector data = {1.0f, 2.0f}; + auto tensor = + make_tensor_ptr(std::move(data), executorch::aten::ScalarType::UInt32); + + EXPECT_EQ(tensor->dim(), 1); + EXPECT_EQ(tensor->size(0), 2); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt32); + + auto ptr = tensor->const_data_ptr(); + EXPECT_EQ(ptr[0], 1u); + EXPECT_EQ(ptr[1], 2u); +} + +TEST_F(TensorPtrTest, MakeTensorPtrFromExistingTensorUInt32) { + std::vector data = {10u, 20u, 30u, 40u}; + auto tensor = make_tensor_ptr({2, 2}, data); + auto alias = make_tensor_ptr(tensor); + EXPECT_EQ(alias->dim(), 2); + EXPECT_EQ(alias->size(0), 2); + EXPECT_EQ(alias->size(1), 2); + EXPECT_EQ(alias->scalar_type(), executorch::aten::ScalarType::UInt32); + EXPECT_EQ( + alias->const_data_ptr(), tensor->const_data_ptr()); +} + +TEST_F(TensorPtrTest, CloneTensorPtrFromExistingTensorUInt32) { + std::vector data = {10u, 20u, 30u, 40u}; + auto tensor = make_tensor_ptr({2, 2}, std::move(data)); + auto cloned = clone_tensor_ptr(tensor); + EXPECT_EQ(cloned->dim(), 2); + EXPECT_EQ(cloned->size(0), 2); + EXPECT_EQ(cloned->size(1), 2); + EXPECT_EQ(cloned->scalar_type(), executorch::aten::ScalarType::UInt32); + EXPECT_NE( + cloned->const_data_ptr(), tensor->const_data_ptr()); + auto ptr = cloned->const_data_ptr(); + EXPECT_EQ(ptr[0], 10u); + EXPECT_EQ(ptr[3], 40u); +} + +TEST_F(TensorPtrTest, Tensor2DUInt16OwningData) { + std::vector data = {1u, 2u, 3u, 4u, 5u, 6u}; + auto tensor = make_tensor_ptr({2, 3}, std::move(data)); + EXPECT_EQ(tensor->dim(), 2); + EXPECT_EQ(tensor->size(0), 2); + EXPECT_EQ(tensor->size(1), 3); + EXPECT_EQ(tensor->strides()[0], 3); + EXPECT_EQ(tensor->strides()[1], 1); + EXPECT_EQ(tensor->scalar_type(), executorch::aten::ScalarType::UInt16); + auto ptr = tensor->const_data_ptr(); + EXPECT_EQ(ptr[0], 1u); + EXPECT_EQ(ptr[5], 6u); +}