diff --git a/backends/arm/runtime/VGFBackend.cpp b/backends/arm/runtime/VGFBackend.cpp index cfec706b8b9..035b8b37b10 100644 --- a/backends/arm/runtime/VGFBackend.cpp +++ b/backends/arm/runtime/VGFBackend.cpp @@ -5,10 +5,12 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include using namespace std; +#include #include #include #include @@ -190,8 +192,18 @@ class VGFBackend final : public ::executorch::runtime::BackendInterface { if (!io->is_input) continue; - size_t io_size = accumulate( - io->size.begin(), io->size.end(), io->elt_size, std::multiplies<>()); + size_t io_size = io->elt_size; + for (int64_t dim : io->size) { + ET_CHECK_OR_RETURN_ERROR( + dim >= 0, + InvalidArgument, + "Negative dimension in IO size: %" PRId64, + dim); + ET_CHECK_OR_RETURN_ERROR( + !c10::mul_overflows(io_size, static_cast(dim), &io_size), + InvalidArgument, + "Overflow computing IO buffer size"); + } void* data; if (!repr->map_io(io, &data)) { @@ -225,8 +237,18 @@ class VGFBackend final : public ::executorch::runtime::BackendInterface { if (io->is_input) continue; - size_t io_size = accumulate( - io->size.begin(), io->size.end(), io->elt_size, std::multiplies<>()); + size_t io_size = io->elt_size; + for (int64_t dim : io->size) { + ET_CHECK_OR_RETURN_ERROR( + dim >= 0, + InvalidArgument, + "Negative dimension in IO size: %" PRId64, + dim); + ET_CHECK_OR_RETURN_ERROR( + !c10::mul_overflows(io_size, static_cast(dim), &io_size), + InvalidArgument, + "Overflow computing IO buffer size"); + } void* data; if (!repr->map_io(io, &data)) {