From ae6cf8973ec4ddabd4c121d2263de27215c71b91 Mon Sep 17 00:00:00 2001
From: Manuel Candales <manuelcandales@gmail.com>
Date: Thu, 26 Jun 2025 11:03:44 -0700
Subject: [PATCH] [ET][Portable] Check scalar overflow: op_full_like

Differential Revision: [D77386637](https://our.internmc.facebook.com/intern/diff/D77386637/)

[ghstack-poisoned]
---
 kernels/portable/cpu/op_full_like.cpp | 22 +++++++++-------------
 kernels/test/op_full_like_test.cpp    | 15 +++++++++++++++
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/kernels/portable/cpu/op_full_like.cpp b/kernels/portable/cpu/op_full_like.cpp
index 7671cd61ea9..0e263fb9c10 100644
--- a/kernels/portable/cpu/op_full_like.cpp
+++ b/kernels/portable/cpu/op_full_like.cpp
@@ -48,23 +48,19 @@ Tensor& full_like_out(
       out,
       "Failed to resize output tensor.");
 
-  ScalarType val_type = utils::get_scalar_dtype(fill_value);
   ScalarType out_type = out.scalar_type();
 
   constexpr auto name = "scalar_tensor.out";
 
-  ET_SWITCH_REALB_TYPES(val_type, ctx, name, CTYPE_VAL, [&] {
-    CTYPE_VAL val;
-    ET_KERNEL_CHECK(
-        ctx, utils::extract_scalar(fill_value, &val), InvalidArgument, );
-
-    ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, name, CTYPE_OUT, [&] {
-      CTYPE_OUT val_casted = static_cast<CTYPE_OUT>(val);
-      auto data_out = out.mutable_data_ptr<CTYPE_OUT>();
-      for (const auto i : c10::irange(out.numel())) {
-        data_out[i] = val_casted;
-      }
-    });
+  ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, name, CTYPE_OUT, [&] {
+    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;
+    }
   });
 
   return out;
diff --git a/kernels/test/op_full_like_test.cpp b/kernels/test/op_full_like_test.cpp
index c0b9dcc4107..6e7692f5347 100644
--- a/kernels/test/op_full_like_test.cpp
+++ b/kernels/test/op_full_like_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>
@@ -65,6 +66,18 @@ class OpFullLikeTest : public OperatorTest {
     ET_EXPECT_KERNEL_FAILURE(
         context_, op_full_like_out(in, value, memory_format, out));
   }
+
+  template <ScalarType DTYPE>
+  void expect_bad_scalar_value_dies(const Scalar& bad_value) {
+    TensorFactory<DTYPE> tf;
+    const std::vector<int32_t> sizes = {2, 2};
+    Tensor in = tf.zeros(sizes);
+    Tensor out = tf.zeros(sizes);
+    optional<MemoryFormat> memory_format;
+
+    ET_EXPECT_KERNEL_FAILURE(
+        context_, op_full_like_out(in, bad_value, memory_format, out));
+  }
 };
 
 template <>
@@ -209,3 +222,5 @@ TEST_F(OpFullLikeTest, HalfSupport) {
   op_full_like_out(in, INFINITY, memory_format, out);
   EXPECT_TENSOR_CLOSE(out, tf.full({2, 3}, INFINITY));
 }
+
+GENERATE_SCALAR_OVERFLOW_TESTS(OpFullLikeTest)