[ET][Portable] Check scalar overflow: op_scalar_tensor (#12082)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #12025 by
@manuelcandales
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/manuelcandales/123/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/manuelcandales/123/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/manuelcandales/122/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/manuelcandales/123/orig
@diff-train-skip-merge

---------

Co-authored-by: Manuel Candales <[email protected]>

Author: pytorchbot

kernels/portable/cpu/op_scalar_tensor.cpp

@@ -24,17 +24,11 @@ scalar_tensor_out(KernelRuntimeContext& ctx, const Scalar& s, Tensor& out) {
 
   constexpr auto name = "scalar_tensor.out";
 
-  if (s.isFloatingPoint() &&
-      executorch::runtime::isIntegralType(out_type, false)) {
-    ET_SWITCH_INT_TYPES(out_type, ctx, name, CTYPE, [&]() {
-      out.mutable_data_ptr<CTYPE>()[0] =
-          static_cast<CTYPE>(utils::scalar_to<int64_t>(s));
-    });
-  } else {
-    ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, name, CTYPE, [&]() {
-      out.mutable_data_ptr<CTYPE>()[0] = utils::scalar_to<CTYPE>(s);
-    });
-  }
+  ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, name, CTYPE, [&]() {
+    auto opt_val_casted = utils::internal::check_overflow_scalar_cast<CTYPE>(s);
+    ET_KERNEL_CHECK(ctx, opt_val_casted.has_value(), InvalidArgument, );
+    out.mutable_data_ptr<CTYPE>()[0] = opt_val_casted.value();
+  });
 
   return out;
 }

kernels/test/op_scalar_tensor_test.cpp

@@ -7,6 +7,7 @@
  */
 
 #include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
+#include <executorch/kernels/test/ScalarOverflowTestMacros.h>
 #include <executorch/kernels/test/TestUtil.h>
 #include <executorch/kernels/test/supported_features.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
@@ -71,6 +72,14 @@ class OpScalarTensorOutTest : public OperatorTest {
 
     ET_EXPECT_KERNEL_FAILURE(context_, op_scalar_tensor_out(value, out));
   }
+
+  template <ScalarType DTYPE>
+  void expect_bad_scalar_value_dies(const Scalar& bad_value) {
+    TensorFactory<DTYPE> tf;
+    Tensor out = tf.zeros({});
+
+    ET_EXPECT_KERNEL_FAILURE(context_, op_scalar_tensor_out(bad_value, out));
+  }
 };
 
 #define GENERATE_TEST_0D(ctype, dtype)                      \
@@ -131,3 +140,5 @@ TEST_F(OpScalarTensorOutTest, HalfSupport) {
   op_scalar_tensor_out(INFINITY, out);
   EXPECT_TENSOR_CLOSE(out, tf.make({}, {INFINITY}));
 }
+
+GENERATE_SCALAR_OVERFLOW_TESTS(OpScalarTensorOutTest)