From 8fbd71f648b3250a098922172c083a859f5c09d6 Mon Sep 17 00:00:00 2001
From: Manuel Candales <manuelcandales@gmail.com>
Date: Fri, 27 Jun 2025 07:28:00 -0700
Subject: [PATCH] [ET][Portable] Check scalar overflow: op_full

Pull Request resolved: https://github.com/pytorch/executorch/pull/12014


ghstack-source-id: 293086873
@exported-using-ghexport

Differential Revision: [D77386635](https://our.internmc.facebook.com/intern/diff/D77386635/)
---
 kernels/portable/cpu/op_full.cpp |  5 ++++-
 kernels/test/op_full_test.cpp    | 27 +++++++++++++--------------
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/kernels/portable/cpu/op_full.cpp b/kernels/portable/cpu/op_full.cpp
index 83ffcad45a6..b83637f2b91 100644
--- a/kernels/portable/cpu/op_full.cpp
+++ b/kernels/portable/cpu/op_full.cpp
@@ -37,7 +37,10 @@ Tensor& full_out(
   constexpr auto name = "full.out";
 
   ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, name, CTYPE_OUT, [&] {
-    CTYPE_OUT val_casted = utils::scalar_to<CTYPE_OUT>(fill_value);
+    auto opt_val_casted =
+        utils::internal::check_overflow_scalar_cast<CTYPE_OUT>(fill_value);
+    ET_KERNEL_CHECK(ctx, opt_val_casted.has_value(), InvalidArgument, );
+    auto val_casted = opt_val_casted.value();
     auto data_out = out.mutable_data_ptr<CTYPE_OUT>();
     for (const auto i : c10::irange(out.numel())) {
       data_out[i] = val_casted;
diff --git a/kernels/test/op_full_test.cpp b/kernels/test/op_full_test.cpp
index 93129679087..35115dc7ed6 100644
--- a/kernels/test/op_full_test.cpp
+++ b/kernels/test/op_full_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>
@@ -59,6 +60,17 @@ class OpFullOutTest : public OperatorTest {
     op_full_out(aref, 1.0, out);
     EXPECT_TENSOR_EQ(out, tf.ones(size_int32_t));
   }
+
+  template <ScalarType DTYPE>
+  void expect_bad_scalar_value_dies(const Scalar& bad_value) {
+    TensorFactory<DTYPE> tf;
+    std::vector<int32_t> sizes = {2, 2};
+    std::vector<int64_t> sizes_int64_t(sizes.begin(), sizes.end());
+    auto aref = IntArrayRef(sizes_int64_t.data(), sizes_int64_t.size());
+    Tensor out = tf.zeros(sizes);
+
+    ET_EXPECT_KERNEL_FAILURE(context_, op_full_out(aref, bad_value, out));
+  }
 };
 
 #define GENERATE_TEST(_, DTYPE)                  \
@@ -72,20 +84,7 @@ class OpFullOutTest : public OperatorTest {
 
 ET_FORALL_REALHBF16_TYPES(GENERATE_TEST)
 
-TEST_F(OpFullOutTest, ValueOverflow) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel doesn't handle overflow";
-  }
-  TensorFactory<ScalarType::Byte> tf;
-
-  std::vector<int64_t> sizes_int64_t_vec = {2, 3};
-  std::vector<int32_t> sizes_in32_t_vec = {2, 3};
-  auto sizes = IntArrayRef(sizes_int64_t_vec.data(), sizes_int64_t_vec.size());
-
-  Tensor out = tf.zeros(sizes_in32_t_vec);
-
-  op_full_out(sizes, 1000, out);
-}
+GENERATE_SCALAR_OVERFLOW_TESTS(OpFullOutTest)
 
 TEST_F(OpFullOutTest, HalfSupport) {
   TensorFactory<ScalarType::Half> tf;