diff --git a/backends/aoti/slim/c10/core/ScalarType.h b/backends/aoti/slim/c10/core/ScalarType.h index 28391f012d5..3acf2bebc48 100644 --- a/backends/aoti/slim/c10/core/ScalarType.h +++ b/backends/aoti/slim/c10/core/ScalarType.h @@ -133,6 +133,24 @@ inline bool isBoolType(ScalarType t) { return t == ScalarType::Bool; } +/// Checks if the scalar type is a valid/supported type. +/// @param t The scalar type to check. +/// @return true if the scalar type is valid, false otherwise. +inline bool isValidScalarType(ScalarType t) { + switch (t) { + case ScalarType::Char: + case ScalarType::Short: + case ScalarType::Int: + case ScalarType::Long: + case ScalarType::Float: + case ScalarType::Bool: + case ScalarType::BFloat16: + return true; + default: + return false; + } +} + inline std::ostream& operator<<(std::ostream& stream, ScalarType scalar_type) { return stream << toString(scalar_type); } diff --git a/backends/aoti/slim/core/slim_tensor.h b/backends/aoti/slim/core/slim_tensor.h index b32186152e2..391ae2248b4 100644 --- a/backends/aoti/slim/core/slim_tensor.h +++ b/backends/aoti/slim/core/slim_tensor.h @@ -52,6 +52,7 @@ class SlimTensor { storage_offset_(storage_offset), dtype_(dtype) { set_sizes_and_strides(sizes, strides); + check_supportive(); } /** @@ -65,6 +66,7 @@ class SlimTensor { is_contiguous_(true) { sizes_and_strides_.set_sizes({0}); sizes_and_strides_.set_strides({1}); + check_supportive(); } // Default copy/move operations @@ -556,6 +558,13 @@ class SlimTensor { static_cast(numel_)); } + void check_supportive() const { + ET_CHECK_MSG( + c10::isValidScalarType(dtype_), + "invalid dtype %d", + static_cast(dtype_)); + } + Storage storage_; int64_t storage_offset_{0}; c10::SizesAndStrides sizes_and_strides_; diff --git a/backends/aoti/slim/core/storage.h b/backends/aoti/slim/core/storage.h index be4b374b333..bd227dbb43a 100644 --- a/backends/aoti/slim/core/storage.h +++ b/backends/aoti/slim/core/storage.h @@ -240,7 +240,11 @@ class MaybeOwningStorage { data_(data), capacity_(nbytes), deleter_(detail::noop), - is_owning_(false) {} + is_owning_(false) { + if (!device.is_cuda() && !device.is_cpu()) { + ET_CHECK_MSG(false, "Unsupported device type: %s", device.str().c_str()); + } + } /// Default constructor is deleted - storage must have a device. MaybeOwningStorage() = delete; diff --git a/backends/aoti/slim/core/test/test_slimtensor_basic.cpp b/backends/aoti/slim/core/test/test_slimtensor_basic.cpp index 0337cca544a..29ebe5875b4 100644 --- a/backends/aoti/slim/core/test/test_slimtensor_basic.cpp +++ b/backends/aoti/slim/core/test/test_slimtensor_basic.cpp @@ -489,4 +489,38 @@ TEST(SlimTensorBasicTest, DataPtrWithOffset) { EXPECT_EQ(data, static_cast(base) + 5 * sizeof(float)); } +// ============================================================================= +// Dtype and Device Type Validation Tests +// ============================================================================= + +TEST(SlimTensorValidationTest, InvalidDtypeUndefined) { + std::vector sizes = {2, 3}; + std::vector strides = {3, 1}; + size_t nbytes = 6 * sizeof(float); + Storage storage = make_cpu_storage(nbytes); + + EXPECT_DEATH( + SlimTensor( + std::move(storage), + makeArrayRef(sizes), + makeArrayRef(strides), + c10::ScalarType::Undefined), + ""); +} + +TEST(SlimTensorValidationTest, InvalidDtypeDouble) { + std::vector sizes = {2, 3}; + std::vector strides = {3, 1}; + size_t nbytes = 6 * sizeof(double); + Storage storage = make_cpu_storage(nbytes); + + EXPECT_DEATH( + SlimTensor( + std::move(storage), + makeArrayRef(sizes), + makeArrayRef(strides), + static_cast(7)), // Double = 7 + ""); +} + } // namespace executorch::backends::aoti::slim diff --git a/backends/aoti/slim/factory/test/test_empty.cpp b/backends/aoti/slim/factory/test/test_empty.cpp index 39a39ac0f76..f4d0e31da9c 100644 --- a/backends/aoti/slim/factory/test/test_empty.cpp +++ b/backends/aoti/slim/factory/test/test_empty.cpp @@ -233,6 +233,65 @@ TEST(EmptyTest, CanWriteAndReadData) { } } +// ============================================================================= +// Dtype and Device Type Validation Tests +// ============================================================================= + +TEST(EmptyStridedTest, InvalidDtypeUndefined) { + std::vector sizes = {2, 3}; + std::vector strides = {3, 1}; + + EXPECT_DEATH( + empty_strided( + makeArrayRef(sizes), + makeArrayRef(strides), + c10::ScalarType::Undefined), + ""); +} + +TEST(EmptyStridedTest, InvalidDtypeDouble) { + std::vector sizes = {2, 3}; + std::vector strides = {3, 1}; + + EXPECT_DEATH( + empty_strided( + makeArrayRef(sizes), + makeArrayRef(strides), + static_cast(7)), // Double = 7 + ""); +} + +TEST(EmptyStridedTest, InvalidDeviceType) { + std::vector sizes = {2, 3}; + std::vector strides = {3, 1}; + + c10::Device invalid_device(static_cast(100), 0); + + EXPECT_DEATH( + empty_strided( + makeArrayRef(sizes), + makeArrayRef(strides), + c10::ScalarType::Float, + invalid_device), + ""); +} + +TEST(EmptyTest, InvalidDtypeUndefined) { + EXPECT_DEATH(empty({2, 3}, c10::ScalarType::Undefined), ""); +} + +TEST(EmptyTest, InvalidDtypeDouble) { + EXPECT_DEATH( + empty({2, 3}, static_cast(7)), // Double = 7 + ""); +} + +TEST(EmptyTest, InvalidDeviceType) { + c10::Device invalid_device(static_cast(100), 0); + + EXPECT_DEATH(empty({2, 3}, c10::ScalarType::Float, invalid_device), ""); +} + #ifdef CUDA_AVAILABLE // ============================================================================= diff --git a/backends/aoti/slim/factory/test/test_from_blob.cpp b/backends/aoti/slim/factory/test/test_from_blob.cpp index a0c8f3cc495..d066d7ba79c 100644 --- a/backends/aoti/slim/factory/test/test_from_blob.cpp +++ b/backends/aoti/slim/factory/test/test_from_blob.cpp @@ -315,6 +315,37 @@ TEST(FromBlobTest, WithArrayRef) { EXPECT_TRUE(tensor.is_contiguous()); } +// ============================================================================= +// Dtype and Device Type Validation Tests +// ============================================================================= + +TEST(FromBlobTest, InvalidDtypeUndefined) { + constexpr size_t kNumFloats = 6; + float external_data[kNumFloats]; + + EXPECT_DEATH( + from_blob(external_data, {2, 3}, c10::ScalarType::Undefined), ""); +} + +TEST(FromBlobTest, InvalidDtypeDouble) { + constexpr size_t kNumFloats = 6; + float external_data[kNumFloats]; + + EXPECT_DEATH( + from_blob(external_data, {2, 3}, static_cast(7)), ""); +} + +TEST(FromBlobTest, InvalidDeviceType) { + constexpr size_t kNumFloats = 6; + float external_data[kNumFloats]; + + c10::Device invalid_device(static_cast(100), 0); + + EXPECT_DEATH( + from_blob(external_data, {2, 3}, c10::ScalarType::Float, invalid_device), + ""); +} + // ============================================================================= // CUDA from_blob Tests // Tests are skipped at runtime if CUDA hardware is not available.