diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index 4e219b07a66..cdda246887c 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -1540,6 +1540,25 @@ Error Method::execute() { "Cannot execute until method has been initialized."); ET_LOG(Debug, "Executing method: %s.", method_meta().name()); + // Validate that outputs are set. + for (size_t i = 0; i < method_meta().num_outputs(); i++) { + auto& output = mutable_value(get_output_index(i)); + if (!output.isTensor()) { + continue; + } + + auto tensor_meta = method_meta().output_tensor_meta(i); + if (tensor_meta->is_memory_planned()) { + continue; + } + + auto& t = output.toTensor(); + if (t.const_data_ptr() == nullptr) { + ET_LOG(Error, "Output tensor at index %" ET_PRIsize_t " is not set.", i); + return Error::InvalidState; + } + } + // Chains are executed sequentially today, but future async designs may // branch and run many in parallel or out of order. for (step_state_.chain_idx = 0; step_state_.chain_idx < n_chains_; diff --git a/runtime/executor/test/method_test.cpp b/runtime/executor/test/method_test.cpp index f597746e0fd..8cd5f26c8e6 100644 --- a/runtime/executor/test/method_test.cpp +++ b/runtime/executor/test/method_test.cpp @@ -390,3 +390,36 @@ TEST_F(MethodTest, OptionalTensorListDeserialization) { EXPECT_EQ(outputs.toTensor().size(2), 10); } */ + +TEST_F(MethodTest, UnsetOutputTest) { + // Validate that methods with non-memory planned outputs return an error when + // the output data pointer is not set. It should not crash. + + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["cat"]->load_method("forward", &mmm.get()); + ASSERT_EQ(method.error(), Error::Ok); + + constexpr int buffer_size = 8; + float buffer[buffer_size]; + for (int i = 0; i < buffer_size; ++i) { + buffer[i] = 0.f; + } + int32_t sizes[2] = {2, 4}; + uint8_t dim_order[2] = {0, 1}; + int32_t strides[2] = {4, 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, + 2, + sizes, + buffer, + dim_order, + strides); + + auto input_err = + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0); + ASSERT_EQ(input_err, Error::Ok); + + // Execute the method once. It should return an error (not crash). + auto execute_error = method->execute(); + ASSERT_EQ(execute_error, Error::InvalidState); +}