From 5dcf5c6cdd4cb5e4d5cddbd16ae7274e3af32d9a Mon Sep 17 00:00:00 2001 From: lucylq Date: Tue, 11 Nov 2025 15:17:40 -0800 Subject: [PATCH] Fix write-heap-buffer-overflow in copy_out Differential Revision: D80885980 Pull Request resolved: https://github.com/pytorch/executorch/pull/15584 (cherry picked from commit 652cc5a8d35f5bdfecfe41b5f9944c2985582c5e) --- kernels/portable/cpu/op_copy.cpp | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/kernels/portable/cpu/op_copy.cpp b/kernels/portable/cpu/op_copy.cpp index 968231fc42e..8164d1ebb02 100644 --- a/kernels/portable/cpu/op_copy.cpp +++ b/kernels/portable/cpu/op_copy.cpp @@ -49,7 +49,8 @@ Tensor& copy_out( // Use direct copy fast path if broadcast is not needed and tensors are // non-empty if (internal::sizes_match_ignoring_leading_1s(out.sizes(), src.sizes()) && - src.numel() > 0) { + src.numel() > 0 && out.nbytes() >= src.nbytes() && + tensors_have_same_dtype(src, out)) { std::memcpy(out.mutable_data_ptr(), src.const_data_ptr(), src.nbytes()); } else { ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, op_name, CTYPE, [&]() { @@ -91,8 +92,9 @@ Tensor& copy_( // Use direct copy fast path if broadcast is not needed and tensors are // non-empty if (internal::sizes_match_ignoring_leading_1s(in.sizes(), src.sizes()) && - src.numel() > 0) { - std::memcpy(in.mutable_data_ptr(), src.const_data_ptr(), in.nbytes()); + src.numel() > 0 && in.nbytes() >= src.nbytes() && + tensors_have_same_dtype(src, in)) { + std::memcpy(in.mutable_data_ptr(), src.const_data_ptr(), src.nbytes()); } else { ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, op_name, CTYPE, [&]() { utils::apply_bitensor_elementwise_fn<