Catch invalid scalar type when parsing tensors (#1518)

Summary:
Pull Request resolved: #1518

Fail non-fatally when encountering an unknown/unhandled `ScalarType` in a `.pte` file. As part of this:
- Move the "types not supported yet" logic out of `scalar_type_util` and into `tensor_parser`, since that decision is an aspect of the runtime and not a fundamental aspect of `ScalarType`.
- Remove the now-duplicate `sizeof_scalar_type` function, which is the same as the exsting `elementSize` function. Before this diff, `sizeof_scalar_type` did the "unsupported" checks that have now moved.
- Add an `isValid()` function to let users of `ScalarType` know whether a given enum value is legit. This makes it possible to avoid the fatal error when calling `elementSize` on a bad value.
- Add unit tests for the new `isValid()`.

Reviewed By: larryliu0820

Differential Revision: D52451738

fbshipit-source-id: 88e47d7a3c688e3e2a68cc86114935c7ff6a73b5

Author: dbort

runtime/core/exec_aten/util/scalar_type_util.h

@@ -253,6 +253,16 @@ ET_FORALL_SCALAR_TYPES(SPECIALIZE_CppTypeToScalarType)
 // Utility functions to retrieve metadata for a given ScalarType
 //
 
+/**
+ * Returns true if the parameter is one of the values covered by
+ * ET_FORALL_SCALAR_TYPES.
+ */
+inline bool isValid(exec_aten::ScalarType type) {
+  return static_cast<int8_t>(type) >= 0 &&
+      type < exec_aten::ScalarType::NumOptions &&
+      type != exec_aten::ScalarType::Undefined;
+}
+
 /**
  * Returns the name of a ScalarType as a C string.
  *
@@ -541,38 +551,6 @@ inline exec_aten::ScalarType promoteTypes(
   return _promoteTypesLookup[static_cast<int>(a)][static_cast<int>(b)];
 }
 
-/**
- * Return the size of corresponding ctype given ScalarType.
- */
-inline size_t sizeof_scalar_type(exec_aten::ScalarType type) {
-  // Reject types that are not yet supported or are out of bounds.
-  ET_CHECK_MSG(
-      type != exec_aten::ScalarType::Half &&
-          type != exec_aten::ScalarType::ComplexHalf &&
-          type != exec_aten::ScalarType::ComplexFloat &&
-          type != exec_aten::ScalarType::ComplexDouble &&
-          type != exec_aten::ScalarType::BFloat16 &&
-          type != exec_aten::ScalarType::Undefined,
-      "Invalid or unsupported ScalarType %" PRId8,
-      static_cast<int8_t>(type));
-
-  size_t type_size = 0;
-#define SCALAR_TYPE_SIZE(ctype, dtype) \
-  case exec_aten::ScalarType::dtype:   \
-    type_size = sizeof(ctype);         \
-    break;
-
-  switch (type) {
-    ET_FORALL_SCALAR_TYPES(SCALAR_TYPE_SIZE)
-    default:
-      ET_CHECK_MSG(
-          false, "Invalid input ScalarType %" PRId8, static_cast<int8_t>(type));
-  }
-#undef SCALAR_TYPE_SIZE
-
-  return type_size;
-}
-
 //
 // Helper macros for switch case macros (see below)
 //

runtime/core/exec_aten/util/test/scalar_type_util_test.cpp

@@ -70,6 +70,26 @@ TEST(ScalarTypeUtilTest, ElementSize) {
   }
 }
 
+TEST(ScalarTypeUtilTest, IsValidTrue) {
+  // Some valid types.
+  EXPECT_TRUE(torch::executor::isValid(ScalarType::Byte));
+  EXPECT_TRUE(torch::executor::isValid(ScalarType::Float));
+  EXPECT_TRUE(torch::executor::isValid(ScalarType::ComplexFloat));
+  EXPECT_TRUE(torch::executor::isValid(ScalarType::Bits16));
+}
+
+TEST(ScalarTypeUtilTest, IsValidFalse) {
+  // Undefined, which is sort of a special case since it's not part of the
+  // iteration macros but is still a part of the enum.
+  EXPECT_FALSE(torch::executor::isValid(ScalarType::Undefined));
+
+  // Some out-of-range types, also demonstrating that NumOptions is not really a
+  // scalar type.
+  EXPECT_FALSE(torch::executor::isValid(ScalarType::NumOptions));
+  EXPECT_FALSE(torch::executor::isValid(static_cast<ScalarType>(127)));
+  EXPECT_FALSE(torch::executor::isValid(static_cast<ScalarType>(-1)));
+}
+
 TEST(ScalarTypeUtilTest, UnknownTypeElementSizeDies) {
   // Undefined, which is sort of a special case since it's not part of the
   // iteration macros but is still a part of the enum.

runtime/core/portable_type/tensor_impl.cpp

@@ -47,12 +47,12 @@ TensorImpl::TensorImpl(
       data_(data),
       dim_(dim),
       numel_(compute_numel(sizes, dim)),
-      capacity_(numel_ * sizeof_scalar_type(type)),
+      capacity_(numel_ * elementSize(type)),
       type_(type),
       shape_dynamism_(dynamism) {}
 
 size_t TensorImpl::nbytes() const {
-  return numel_ * sizeof_scalar_type(type_);
+  return numel_ * elementSize(type_);
 }
 
 ssize_t TensorImpl::size(ssize_t dim) const {
@@ -78,7 +78,7 @@ ScalarType TensorImpl::scalar_type() const {
 
 // Return the size of one element of the tensor
 ssize_t TensorImpl::element_size() const {
-  return sizeof_scalar_type(type_);
+  return elementSize(type_);
 }
 
 const ArrayRef<TensorImpl::SizesType> TensorImpl::sizes() const {
@@ -145,7 +145,7 @@ Error TensorImpl::internal_resize_contiguous(ArrayRef<SizesType> new_sizes) {
 
   // Upper bounded tensors can be reshaped but not beyond upper bound
   if (shape_dynamism_ == TensorShapeDynamism::DYNAMIC_BOUND) {
-    auto new_nbytes = new_numel * sizeof_scalar_type(type_);
+    auto new_nbytes = new_numel * elementSize(type_);
     ET_CHECK_OR_RETURN_ERROR(
         new_nbytes <= capacity_,
         NotSupported,

runtime/core/portable_type/test/executor_tensor_test.cpp

@@ -15,14 +15,17 @@ namespace executor {
 
 TEST(TensorTest, InvalidScalarType) {
   TensorImpl::SizesType sizes[1] = {1};
-  // A type that executorch doesn't support yet.
-  ET_EXPECT_DEATH({ TensorImpl x(ScalarType::BFloat16, 1, sizes); }, "");
 
-  // The literal Undefined type.
+  // Undefined, which is sort of a special case since it's not part of the
+  // iteration macros but is still a part of the enum.
   ET_EXPECT_DEATH({ TensorImpl y(ScalarType::Undefined, 1, sizes); }, "");
 
-  // An int value that doesn't map to a valid enum value
+  // Some out-of-range types, also demonstrating that NumOptions is not really a
+  // scalar type.
   ET_EXPECT_DEATH({ TensorImpl y(ScalarType::NumOptions, 1, sizes); }, "");
+  ET_EXPECT_DEATH(
+      { TensorImpl y(static_cast<ScalarType>(127), 1, sizes); }, "");
+  ET_EXPECT_DEATH({ TensorImpl y(static_cast<ScalarType>(-1), 1, sizes); }, "");
 }
 
 TEST(TensorTest, SetData) {

runtime/executor/method_meta.cpp

@@ -59,7 +59,7 @@ size_t calculate_nbytes(
   for (ssize_t i = 0; i < sizes.size(); i++) {
     n *= sizes[i];
   }
-  return n * sizeof_scalar_type(scalar_type);
+  return n * torch::executor::elementSize(scalar_type);
 }
 
 } // namespace

runtime/executor/tensor_parser_aten.cpp

@@ -9,6 +9,7 @@
 #include <executorch/runtime/executor/tensor_parser.h>
 
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
+#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/program.h>
 #include <executorch/runtime/platform/profiler.h>
@@ -43,6 +44,11 @@ Result<at::Tensor> parseTensor(
 
   // get metadata
   at::ScalarType type = static_cast<at::ScalarType>(s_tensor->scalar_type());
+  ET_CHECK_OR_RETURN_ERROR(
+      isValid(type),
+      InvalidProgram,
+      "Invalid ScalarType %" PRId8,
+      static_cast<int8_t>(type));
   auto options = at::CPU(type).options();
 
   // convert int32 in serialization to int64 for aten

runtime/executor/tensor_parser_portable.cpp

@@ -11,6 +11,7 @@
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
+#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/program.h>
 #include <executorch/runtime/platform/profiler.h>
@@ -33,6 +34,19 @@ Result<torch::executor::Tensor> parseTensor(
       "Non-zero storage offset %" PRId32 " not supported",
       s_tensor->storage_offset());
 
+  ScalarType scalar_type = static_cast<ScalarType>(s_tensor->scalar_type());
+  ET_CHECK_OR_RETURN_ERROR(
+      isValid(scalar_type) &&
+          // Types that do not yet have deserialization support.
+          scalar_type != exec_aten::ScalarType::Half &&
+          scalar_type != exec_aten::ScalarType::ComplexHalf &&
+          scalar_type != exec_aten::ScalarType::ComplexFloat &&
+          scalar_type != exec_aten::ScalarType::ComplexDouble &&
+          scalar_type != exec_aten::ScalarType::BFloat16,
+      InvalidProgram,
+      "Invalid or unsupported ScalarType %" PRId8,
+      static_cast<int8_t>(scalar_type));
+
   TensorShapeDynamism dynamism =
       static_cast<TensorShapeDynamism>(s_tensor->shape_dynamism());
   // TODO(T133200526): Remove this check once fully dynamic shapes are
@@ -90,7 +104,7 @@ Result<torch::executor::Tensor> parseTensor(
   // Placement new on the allocated memory space. Note that we create this first
   // with null data so we can find its expected size before getting its memory.
   new (tensor_impl) torch::executor::TensorImpl(
-      static_cast<ScalarType>(s_tensor->scalar_type()),
+      scalar_type,
       dim,
       sizes,
       /*data=*/nullptr,