diff --git a/runtime/executor/method_meta.cpp b/runtime/executor/method_meta.cpp index c284a0d82fb..2f45660b290 100644 --- a/runtime/executor/method_meta.cpp +++ b/runtime/executor/method_meta.cpp @@ -52,7 +52,7 @@ Result get_tag( } } -size_t calculate_nbytes( +Result calculate_nbytes( Span sizes, executorch::aten::ScalarType scalar_type) { size_t n = 1; @@ -61,7 +61,13 @@ size_t calculate_nbytes( prev_n = n; n *= sizes[i]; // Check for overflow - ET_CHECK(sizes[i] == 0 || n / sizes[i] == prev_n); + ET_CHECK_OR_RETURN_ERROR( + sizes[i] == 0 || n / sizes[i] == prev_n, + InvalidArgument, + "Invalid size[%zu]: %d. Potentially overflowed, expect to be 0 or prev_n: %zu", + i, + sizes[i], + prev_n); } size_t elem_size = executorch::runtime::elementSize(scalar_type); @@ -70,25 +76,47 @@ size_t calculate_nbytes( n = n * elem_size; // Check for overflow - ET_CHECK(elem_size == 0 || n / elem_size == prev_n); + ET_CHECK_OR_RETURN_ERROR( + elem_size == 0 || n / elem_size == prev_n, + InvalidArgument, + "Invalid elem_size: %zu. Potentially overflowed, expect to be 0 or prev_n: %zu", + elem_size, + prev_n); return n; } } // namespace +/*static*/ Result TensorInfo::create( + Span sizes, + Span dim_order, + executorch::aten::ScalarType scalar_type, + const bool is_memory_planned, + std::string_view name) { + auto nbytes = calculate_nbytes(sizes, scalar_type); + ET_CHECK_OR_RETURN_ERROR( + nbytes.ok(), + InvalidArgument, + "Failed to calculate nbytes for TensorInfo"); + + return TensorInfo( + sizes, dim_order, scalar_type, is_memory_planned, name, nbytes.get()); +} + TensorInfo::TensorInfo( Span sizes, Span dim_order, executorch::aten::ScalarType scalar_type, const bool is_memory_planned, - std::string_view name) + std::string_view name, + size_t nbytes) : sizes_(sizes), dim_order_(dim_order), name_(name), scalar_type_(scalar_type), is_memory_planned_(is_memory_planned), - nbytes_(calculate_nbytes(sizes_, scalar_type_)) {} + nbytes_(nbytes) {} Span TensorInfo::sizes() const { return sizes_; @@ -160,7 +188,7 @@ Result MethodMeta::input_tensor_meta(size_t index) const { auto input_index = s_plan_->inputs()->Get(index); // input_index was already validated by input_tag(). auto tensor_value = s_plan_->values()->Get(input_index)->val_as_Tensor(); - return TensorInfo( + return TensorInfo::create( Span( tensor_value->sizes()->data(), tensor_value->sizes()->size()), Span( @@ -212,7 +240,7 @@ Result MethodMeta::output_tensor_meta(size_t index) const { // output_index was already validated by output_tag(). auto tensor_value = s_plan_->values()->Get(output_index)->val_as_Tensor(); - return TensorInfo( + return TensorInfo::create( Span( tensor_value->sizes()->data(), tensor_value->sizes()->size()), Span( @@ -255,7 +283,7 @@ Result MethodMeta::attribute_tensor_meta(size_t index) const { auto t_name = tensor_value->extra_tensor_info()->fully_qualified_name(); // Count constant returns as memory planned - return TensorInfo( + return TensorInfo::create( Span( tensor_value->sizes()->data(), tensor_value->sizes()->size()), Span( diff --git a/runtime/executor/method_meta.h b/runtime/executor/method_meta.h index 1b3be75ef17..79fd05c28ee 100644 --- a/runtime/executor/method_meta.h +++ b/runtime/executor/method_meta.h @@ -77,13 +77,32 @@ class TensorInfo final { friend class MethodMeta; friend class testing::TensorInfoTestFriend; - TensorInfo( + /** + * Create a TensorInfo instance. + * + * @param[in] sizes The sizes of the tensor. + * @param[in] dim_order The dim order of the tensor. + * @param[in] scalar_type The scalar type of the tensor. + * @param[in] is_memory_planned Whether the tensor's memory was planned. + * @param[in] name The fully qualified name of the tensor. + * @returns A Result containing the TensorInfo on success, or an error on + * failure. + */ + static Result create( Span sizes, Span dim_order, executorch::aten::ScalarType scalar_type, const bool is_memory_planned, std::string_view name); + TensorInfo( + Span sizes, + Span dim_order, + executorch::aten::ScalarType scalar_type, + const bool is_memory_planned, + std::string_view name, + size_t nbytes); + /** * The sizes of the tensor. * diff --git a/runtime/executor/test/method_meta_test.cpp b/runtime/executor/test/method_meta_test.cpp index f5a07e352aa..e4ef2e72a85 100644 --- a/runtime/executor/test/method_meta_test.cpp +++ b/runtime/executor/test/method_meta_test.cpp @@ -39,12 +39,13 @@ class TensorInfoTestFriend final { executorch::aten::ScalarType scalar_type, const bool is_memory_planned, executorch::aten::string_view name) { - return TensorInfo( - Span(sizes.data(), sizes.size()), - Span(dim_order.data(), dim_order.size()), - scalar_type, - is_memory_planned, - name); + return TensorInfo::create( + Span(sizes.data(), sizes.size()), + Span(dim_order.data(), dim_order.size()), + scalar_type, + is_memory_planned, + name) + .get(); } }; } // namespace testing