Refactor write_tensor_or_raise_error to Return Error Codes Instead of Raising Errors

Differential Revision: D75111482

Pull Request resolved: pytorch/executorch#11051

Author: Juntian777

devtools/etdump/etdump_flatcc.cpp

@@ -394,9 +394,12 @@ Result<bool> ETDumpGen::log_intermediate_output_delegate_helper(
 
   // Check the type of `output` then call the corresponding logging functions
   if constexpr (std::is_same<T, Tensor>::value) {
-    long offset = write_tensor_or_raise_error(output);
+    Result<long> offset = write_tensor_or_return_error(output);
+    ET_CHECK_MSG(
+        offset.ok(),
+        "write_tensor_or_return_error() failed to write tensor to debug buffer");
     Result<etdump_Tensor_ref_t> tensor_ref =
-        add_tensor_entry(builder_, output, offset);
+        add_tensor_entry(builder_, output, offset.get());
     if (!tensor_ref.ok()) {
       return tensor_ref.error();
     }
@@ -408,9 +411,12 @@ Result<bool> ETDumpGen::log_intermediate_output_delegate_helper(
   } else if constexpr (std::is_same<T, ArrayRef<Tensor>>::value) {
     etdump_Tensor_vec_start(builder_);
     for (size_t i = 0; i < output.size(); ++i) {
-      long offset = write_tensor_or_raise_error(output[i]);
+      Result<long> offset = write_tensor_or_return_error(output[i]);
+      ET_CHECK_MSG(
+          offset.ok(),
+          "write_tensor_or_return_error() failed to write tensor to debug buffer");
       Result<etdump_Tensor_ref_t> tensor_ref =
-          add_tensor_entry(builder_, output[i], offset);
+          add_tensor_entry(builder_, output[i], offset.get());
       if (!tensor_ref.ok()) {
         return tensor_ref.error();
       }
@@ -566,9 +572,12 @@ Result<bool> ETDumpGen::log_evalue(
   switch (evalue.tag) {
     case Tag::Tensor: {
       executorch::aten::Tensor tensor = evalue.toTensor();
-      long offset = write_tensor_or_raise_error(tensor);
+      Result<long> offset = write_tensor_or_return_error(tensor);
+      ET_CHECK_MSG(
+          offset.ok(),
+          "write_tensor_or_return_error() failed to write tensor to debug buffer");
       Result<etdump_Tensor_ref_t> tensor_ref =
-          add_tensor_entry(builder_, tensor, offset);
+          add_tensor_entry(builder_, tensor, offset.get());
       if (!tensor_ref.ok()) {
         return tensor_ref.error();
       }
@@ -591,9 +600,12 @@ Result<bool> ETDumpGen::log_evalue(
           evalue.toTensorList();
       etdump_Tensor_vec_start(builder_);
       for (size_t i = 0; i < tensors.size(); ++i) {
-        long offset = write_tensor_or_raise_error(tensors[i]);
+        Result<long> offset = write_tensor_or_return_error(tensors[i]);
+        ET_CHECK_MSG(
+            offset.ok(),
+            "write_tensor_or_return_error() failed to write tensor to debug buffer");
         Result<etdump_Tensor_ref_t> tensor_ref =
-            add_tensor_entry(builder_, tensors[i], offset);
+            add_tensor_entry(builder_, tensors[i], offset.get());
         if (!tensor_ref.ok()) {
           return tensor_ref.error();
         }
@@ -689,7 +701,7 @@ void ETDumpGen::set_delegation_intermediate_output_filter(
   filter_ = filter;
 }
 
-long ETDumpGen::write_tensor_or_raise_error(Tensor tensor) {
+Result<long> ETDumpGen::write_tensor_or_return_error(Tensor tensor) {
   // Previously, the function copy_tensor_to_debug_buffer returned 0xFF..F when
   // given an empty tensor, which is an invalid offset for most buffers. In our
   // data sink, we will return the current debug_buffer_offset for better
@@ -702,14 +714,14 @@ long ETDumpGen::write_tensor_or_raise_error(Tensor tensor) {
     return static_cast<size_t>(-1);
   }
 
-  ET_CHECK_MSG(
-      data_sink_, "Must set data sink before writing tensor-like data");
+  if (!data_sink_) {
+    return Error::InvalidArgument;
+  }
   Result<size_t> ret =
       data_sink_->write(tensor.const_data_ptr(), tensor.nbytes());
-  ET_CHECK_MSG(
-      ret.ok(),
-      "Failed to write tensor with error 0x%" PRIx32,
-      static_cast<uint32_t>(ret.error()));
+  if (!ret.ok()) {
+    return ret.error();
+  }
   return static_cast<long>(ret.get());
 }
 

devtools/etdump/etdump_flatcc.h

@@ -184,7 +184,7 @@ class ETDumpGen : public ::executorch::runtime::EventTracer {
       DelegateDebugIntId delegate_debug_index,
       const T& output);
 
-  long write_tensor_or_raise_error(executorch::aten::Tensor tensor);
+  Result<long> write_tensor_or_return_error(executorch::aten::Tensor tensor);
 
   struct flatcc_builder* builder_;
   size_t num_blocks_ = 0;

devtools/etdump/tests/etdump_test.cpp

@@ -76,7 +76,7 @@ class ProfilerETDumpTest : public ::testing::Test {
     ET_EXPECT_DEATH(
         gen->log_intermediate_output_delegate(
             "test_event_tensor", kUnsetDelegateDebugIntId, tf.ones({3, 2})),
-        "Must set data sink before writing tensor-like data");
+        "failed to write tensor to debug buffer");
   }
 
   void check_log_with_filter(
@@ -301,7 +301,7 @@ TEST_F(ProfilerETDumpTest, DebugEvent) {
       if (j == 0) {
         ET_EXPECT_DEATH(
             etdump_gen[i]->log_evalue(evalue_tensor),
-            "Must set data sink before writing tensor-like data");
+            "failed to write tensor to debug buffer");
 
         // Set debug buffer with span
         etdump_gen[i]->set_debug_buffer(buffer);
@@ -315,7 +315,7 @@ TEST_F(ProfilerETDumpTest, DebugEvent) {
 
         ET_EXPECT_DEATH(
             etdump_gen[i]->log_evalue(evalue_tensor),
-            "Must set data sink before writing tensor-like data");
+            "failed to write tensor to debug buffer");
 
         if (j == 1) {
           // Set buffer data sink