From 96f4b6a7aa4594e65656d365a8508dc90c2e1770 Mon Sep 17 00:00:00 2001
From: Tarun Karuturi <tkaruturi@meta.com>
Date: Fri, 24 Jan 2025 12:48:46 -0800
Subject: [PATCH] Use the proper offset for front_cursor in flatcc emitter
 (#7933)

Summary:

Updating the front_cursor pointer in the ETDumpGen class to point to the correct location in the buffer, and adding a check to ensure that the buffer size provided is greater than the minimum required size.
This fixes a long running issue that the Turing team was running into when generating ETDump's for large models.

The etdump_test.cpp file has been updated to use a smaller buffer size for testing purposes and we do more stress testing.

Reviewed By: Gasoonjia

Differential Revision: D68502867
---
 devtools/etdump/etdump_flatcc.cpp     | 9 ++++++++-
 devtools/etdump/etdump_flatcc.h       | 7 +++++--
 devtools/etdump/tests/etdump_test.cpp | 9 ++++++---
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/devtools/etdump/etdump_flatcc.cpp b/devtools/etdump/etdump_flatcc.cpp
index e175cc5d0fd..c8e55b18d7c 100644
--- a/devtools/etdump/etdump_flatcc.cpp
+++ b/devtools/etdump/etdump_flatcc.cpp
@@ -116,8 +116,15 @@ ETDumpGen::ETDumpGen(Span<uint8_t> buffer) {
     builder_ = (struct flatcc_builder*)alignPointer(buffer.data(), 64);
     uintptr_t buffer_with_builder =
         (uintptr_t)alignPointer(builder_ + sizeof(struct flatcc_builder), 64);
-    size_t buffer_size = buffer.size() -
+    size_t builder_size =
         (size_t)(buffer_with_builder - (uintptr_t)buffer.data());
+    size_t min_buf_size = max_alloc_buf_size + builder_size;
+    ET_CHECK_MSG(
+        buffer.size() > min_buf_size,
+        "Static buffer size provided to ETDumpGen is %zu, which is less than or equal to the minimum size of %zu",
+        buffer.size(),
+        min_buf_size);
+    size_t buffer_size = buffer.size() - builder_size;
     alloc_.set_buffer(
         (uint8_t*)buffer_with_builder,
         buffer_size,
diff --git a/devtools/etdump/etdump_flatcc.h b/devtools/etdump/etdump_flatcc.h
index dcfe687577a..4a818d18e55 100644
--- a/devtools/etdump/etdump_flatcc.h
+++ b/devtools/etdump/etdump_flatcc.h
@@ -31,8 +31,11 @@ struct ETDumpStaticAllocator {
     data_size = alloc_buf_size;
     allocated = 0;
     out_size = total_buf_size - alloc_buf_size;
-    front_cursor = &buffer[alloc_buf_size];
-    front_left = out_size / 2;
+    // The front of the buffer is the end of the allocation buffer.
+    // We start writing from the end of the allocation buffer, and
+    // move backwards.
+    front_cursor = &buffer[alloc_buf_size + out_size];
+    front_left = out_size;
   }
 
   // Pointer to backing buffer to allocate from.
diff --git a/devtools/etdump/tests/etdump_test.cpp b/devtools/etdump/tests/etdump_test.cpp
index b750e21eb07..f45652ab8f9 100644
--- a/devtools/etdump/tests/etdump_test.cpp
+++ b/devtools/etdump/tests/etdump_test.cpp
@@ -40,7 +40,7 @@ class ProfilerETDumpTest : public ::testing::Test {
   void SetUp() override {
     torch::executor::runtime_init();
     etdump_gen[0] = new ETDumpGen();
-    const size_t buf_size = 1024 * 1024;
+    const size_t buf_size = 512 * 1024;
     buf = (uint8_t*)malloc(buf_size * sizeof(uint8_t));
     etdump_gen[1] = new ETDumpGen(Span<uint8_t>(buf, buf_size));
   }
@@ -58,8 +58,11 @@ class ProfilerETDumpTest : public ::testing::Test {
 TEST_F(ProfilerETDumpTest, SingleProfileEvent) {
   for (size_t i = 0; i < 2; i++) {
     etdump_gen[i]->create_event_block("test_block");
-    EventTracerEntry entry = etdump_gen[i]->start_profiling("test_event", 0, 1);
-    etdump_gen[i]->end_profiling(entry);
+    for (size_t j = 0; j < 2048; j++) {
+      EventTracerEntry entry =
+          etdump_gen[i]->start_profiling("test_event", 0, 1);
+      etdump_gen[i]->end_profiling(entry);
+    }
 
     ETDumpResult result = etdump_gen[i]->get_etdump_data();
     ASSERT_TRUE(result.buf != nullptr);