Add flag to disable constant buffer

CMakeLists.txt

@@ -189,6 +189,9 @@ if(NOT EXECUTORCH_ENABLE_PROGRAM_VERIFICATION)
   add_definitions(-DET_ENABLE_PROGRAM_VERIFICATION=0)
 endif()
 
+# Disable the deprecated constant_buffer path.
+add_definitions(-DET_ENABLE_DEPRECATED_CONSTANT_BUFFER=0)
+
 if(EXECUTORCH_ENABLE_EVENT_TRACER)
   add_definitions(-DET_EVENT_TRACER_ENABLED)
 endif()

runtime/executor/program.cpp

@@ -29,6 +29,14 @@
 #define ET_ENABLE_PROGRAM_VERIFICATION 1
 #endif
 
+/*
+ * The constant_buffer path is deprecated from ExecuTorch 0.7. Disable it by
+ * passing -DET_ENABLE_DEPRECATED_CONSTANT_BUFFER=0.
+ */
+#ifndef ET_ENABLE_DEPRECATED_CONSTANT_BUFFER
+#define ET_ENABLE_DEPRECATED_CONSTANT_BUFFER 1
+#endif
+
 namespace executorch {
 namespace ET_RUNTIME_NAMESPACE {
 namespace {
@@ -306,6 +314,7 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
     // https://docs.pytorch.org/executorch/stable/api-life-cycle.html#deprecation-policy.
     // For support, contact the PyTorch Edge team or make an issue in:
     // https://github.com/pytorch/executorch/issues.
+#if ET_ENABLE_DEPRECATED_CONSTANT_BUFFER
     ET_LOG(
         Error,
         "!!DEPRECATED!! This branch is deprecated from ExecuTorch 0.7; re-export this PTE file to ensure support on newer runtimes.");
@@ -316,6 +325,13 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
         flatbuffer_program,
         /*constant_segment_data=*/FreeableBuffer{},
         std::move(pte_data_map));
+#else
+    ET_LOG(
+        Error,
+        "PTE file relies on the constant_buffer path, which is disabled in this"
+        " build (ET_ENABLE_DEPRECATED_CONSTANT_BUFFER=0). Please re-export the PTE file.");
+    return Error::InvalidProgram;
+#endif
   }
 }
 
@@ -449,6 +465,7 @@ Result<const void*> Program::get_constant_buffer_data(
         static_cast<const unsigned char*>(constant_segment_data_.data()) +
         offset);
   } else {
+#if ET_ENABLE_DEPRECATED_CONSTANT_BUFFER
     // Otherwise, the constant data is stored inside Program.constant_buffer.
     const auto* constant_buffer_ptr = internal_program->constant_buffer();
     size_t num_elems =
@@ -475,6 +492,14 @@ Result<const void*> Program::get_constant_buffer_data(
         static_cast<size_t>(storage_size));
 
     return storage->data();
+#else
+    (void)buffer_index;
+    (void)nbytes;
+    ET_LOG(
+        Error,
+        "constant_buffer path is disabled (ET_ENABLE_DEPRECATED_CONSTANT_BUFFER=0). Please re-export the PTE file.");
+    return Error::InvalidProgram;
+#endif
   }
 }
 

runtime/executor/test/program_test.cpp

@@ -199,6 +199,12 @@ TEST_F(ProgramTest, BadMagicFailsToLoad) {
 #define ET_ENABLE_PROGRAM_VERIFICATION 1
 #endif
 
+// ET_ENABLE_DEPRECATED_CONSTANT_BUFFER is 1 by default, and
+// set to 0 in the root CMakeLists.txt for OSS builds.
+#ifndef ET_ENABLE_DEPRECATED_CONSTANT_BUFFER
+#define ET_ENABLE_DEPRECATED_CONSTANT_BUFFER 1
+#endif
+
 #if ET_ENABLE_PROGRAM_VERIFICATION
 
 TEST_F(ProgramTest, VerificationCatchesTruncation) {
@@ -522,6 +528,7 @@ TEST_F(ProgramTest, LoadConstantSegmentWhenConstantBufferExists) {
       Program::HeaderStatus::CompatibleVersion);
 
   Result<Program> program = Program::load(&linear_loader.get());
+#if ET_ENABLE_DEPRECATED_CONSTANT_BUFFER
   ASSERT_EQ(program.error(), Error::Ok);
 
   const executorch_flatbuffer::Program* flatbuffer_program =
@@ -532,6 +539,11 @@ TEST_F(ProgramTest, LoadConstantSegmentWhenConstantBufferExists) {
 
   // The constant buffer should exist.
   EXPECT_GE(flatbuffer_program->constant_buffer()->size(), 1);
+#else
+  // The deprecated in-flatbuffer constant_buffer path is disabled; loading a
+  // PTE that relies on it must be rejected.
+  EXPECT_EQ(program.error(), Error::InvalidProgram);
+#endif
 }
 
 TEST_F(ProgramTest, LoadFromMutableSegment) {
@@ -706,14 +718,21 @@ TEST_F(ProgramTest, GetOutputFlatteningEncodingWithMissingContainerMetaType) {
   auto empty_segments = builder.CreateVector(
       std::vector<flatbuffers::Offset<executorch_flatbuffer::DataSegment>>{});
 
+  // Placeholder constant_segment.
+  auto constant_segment = executorch_flatbuffer::CreateSubsegmentOffsets(
+      builder,
+      /*segment_index=*/0,
+      builder.CreateVector(std::vector<uint64_t>{0}));
+
   // Build the Program
   auto program = executorch_flatbuffer::CreateProgram(
       builder,
       0, // version
       execution_plans,
       empty_constant_buffer,
       empty_backend_data,
-      empty_segments);
+      empty_segments,
+      constant_segment);
 
   builder.Finish(program, executorch_flatbuffer::ProgramIdentifier());
 
@@ -788,13 +807,19 @@ TEST_F(ProgramTest, GetOutputFlatteningEncodingWithMissingEncodedOutStr) {
   auto empty_segments = builder.CreateVector(
       std::vector<flatbuffers::Offset<executorch_flatbuffer::DataSegment>>{});
 
+  auto constant_segment = executorch_flatbuffer::CreateSubsegmentOffsets(
+      builder,
+      /*segment_index=*/0,
+      builder.CreateVector(std::vector<uint64_t>{0}));
+
   auto program = executorch_flatbuffer::CreateProgram(
       builder,
       0,
       execution_plans,
       empty_constant_buffer,
       empty_backend_data,
-      empty_segments);
+      empty_segments,
+      constant_segment);
 
   builder.Finish(program, executorch_flatbuffer::ProgramIdentifier());
 
@@ -837,6 +862,11 @@ TEST_F(ProgramTest, NullPlanNameDoesNotCrash) {
               flatbuffers::Offset<executorch_flatbuffer::BackendDelegate>>{}),
       builder.CreateVector(std::vector<int64_t>{0}));
 
+  auto constant_segment = executorch_flatbuffer::CreateSubsegmentOffsets(
+      builder,
+      /*segment_index=*/0,
+      builder.CreateVector(std::vector<uint64_t>{0}));
+
   auto program = executorch_flatbuffer::CreateProgram(
       builder,
       0,
@@ -851,7 +881,8 @@ TEST_F(ProgramTest, NullPlanNameDoesNotCrash) {
               executorch_flatbuffer::BackendDelegateInlineData>>{}),
       builder.CreateVector(
           std::vector<
-              flatbuffers::Offset<executorch_flatbuffer::DataSegment>>{}));
+              flatbuffers::Offset<executorch_flatbuffer::DataSegment>>{}),
+      constant_segment);
 
   builder.Finish(program, executorch_flatbuffer::ProgramIdentifier());
 
@@ -866,6 +897,7 @@ TEST_F(ProgramTest, NullPlanNameDoesNotCrash) {
   EXPECT_EQ(p->method_meta("forward").error(), Error::InvalidArgument);
 }
 
+#if ET_ENABLE_DEPRECATED_CONSTANT_BUFFER
 TEST_F(ProgramTest, GetConstantBufferDataRejectsOversizedRequest) {
   const char* path =
       std::getenv("DEPRECATED_ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH");
@@ -881,3 +913,4 @@ TEST_F(ProgramTest, GetConstantBufferDataRejectsOversizedRequest) {
 
   EXPECT_EQ(data.error(), Error::InvalidArgument);
 }
+#endif // ET_ENABLE_DEPRECATED_CONSTANT_BUFFER