Refuse a PTE or PTD file whose schema version is newer than the runtime can read (pytorch#22114)

### Summary

Every exported PTE file and every PTD data file carries a schema
version. The exporters stamp it, and nothing in the runtime ever read it
back. The only gate on a file was its four byte identifier, `ET12` for a
program and `FT01` for data. That says "this is an ExecuTorch file" and
says nothing about which shape of file it is. A file written by an
exporter newer than the runtime was therefore accepted, then misread
field by field, and failed later at load or during execution with an
error that points somewhere else.

This matters more now that the two halves can come from different
builds. The wheel ships prebuilt runtime libraries, so a user can export
with one installation of ExecuTorch and run with a runtime that was
never built next to it.

This PR reads the version in both readers, right after the root table is
obtained, and refuses anything above the highest version the runtime
supports. `Program::load` returns `InvalidProgram`,
`FlatTensorDataMap::load` returns `InvalidExternalData`, which is what
the checks around each of them return. Both name the two numbers:

```
Program schema version 1 is newer than the highest this runtime supports (0).
Export the model with an older ExecuTorch, or update the runtime.
```

The comparison is "less than or equal", not "equal", because the project
promises that an older file keeps working on a newer runtime, and the
schemas only ever grow by appending optional fields. Only a file from
the future is refused.

Each new constant sits next to the class it guards, and points in a
comment at the writer constants for the same number:
`EXECUTORCH_SCHEMA_VERSION` for the program, `kSchemaVersion` and
`_FLAT_TENSOR_VERSION` for the data file. A bump of one without the
other is then easy to spot. Nothing writes a version other than zero
today, so no file in the wild changes behavior. This is the reader side
that a future version bump needs in order to mean anything.

One limit worth stating plainly: a runtime that is already deployed has
no check at all, so this protects runtimes built from this change
onward. That is inherent to adding a reader.

### Test plan

Four new tests, two per format. Each pair builds a minimal file in
memory through the real FlatBuffer builder, one stamped at the supported
version and one above it. The first is expected to load and the second
to be refused. The positive test is the control: without it, a rejection
could just as well come from a malformed fixture as from the version.
The program test asks for the cheapest verification level, which shows
the check does not depend on full verification being enabled. The data
test follows the byte layout that `save_ptd()` writes, with no segments.

Built and ran both suites on Linux x86_64, before and after the change:

```
cmake -S . -B cmake-out -DEXECUTORCH_BUILD_TESTS=ON \
      -DEXECUTORCH_ENABLE_PROGRAM_VERIFICATION=ON \
      -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
      -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON
cmake --build cmake-out -j
./cmake-out/runtime/executor/test/program_test
./cmake-out/extension/flat_tensor/test/extension_flat_tensor_test
```

The existing tests pass unchanged in both runs, plus the four new ones.
As a control, deleting either check while keeping the tests makes that
format's negative test fail, with the newer file loading and reporting
no error, so the check is what catches it.

Binary size, measured on the two object files. `program.cpp.o` grows by
24 bytes with logging disabled and 239 bytes with logging enabled.
`flat_tensor_data_map.cpp.o` grows by 32 bytes with logging disabled and
243 bytes with logging enabled.

Author: shoumikhin

exir/version.py

@@ -6,4 +6,12 @@
 
 # pyre-strict
 
+# Stamped into Program.version of every exported PTE file. Keep in sync with
+# Program::kMaxSupportedSchemaVersion in //executorch/runtime/executor/program.h,
+# which is the highest version the C++ runtime agrees to load.
+#
+# Bump this only for a change that an older runtime would misread if it went
+# unnoticed: a semantic change to an existing field, or a new field the runtime
+# must understand to execute correctly. Purely additive optional fields stay
+# backward/forward compatible (see schema/README.md) and need no bump.
 EXECUTORCH_SCHEMA_VERSION = 0

extension/flat_tensor/flat_tensor_data_map.cpp

@@ -293,6 +293,21 @@ ET_NODISCARD Result<const char*> FlatTensorDataMap::get_key(
   const flat_tensor_flatbuffer::FlatTensor* flat_tensor =
       flat_tensor_flatbuffer::GetFlatTensor(flat_tensor_data->data());
 
+  // The file identifier above ("FT01") is bumped by convention only on a
+  // backward-incompatible schema change, so it selects a schema family. The
+  // version is the finer gate within that family: a file written by a newer
+  // exporter is refused here instead of being misread field by field. Older
+  // files stay loadable because the schema only grows by appending optional
+  // fields.
+  ET_CHECK_OR_RETURN_ERROR(
+      flat_tensor->version() <= kMaxSupportedSchemaVersion,
+      InvalidExternalData,
+      "FlatTensor schema version %u is newer than the highest this runtime "
+      "supports (%u). Export the data with an older ExecuTorch, or update the "
+      "runtime.",
+      flat_tensor->version(),
+      kMaxSupportedSchemaVersion);
+
   // Validate flat_tensor.
   ET_CHECK_OR_RETURN_ERROR(
       flat_tensor->named_data() != nullptr,

extension/flat_tensor/flat_tensor_data_map.h

@@ -35,6 +35,17 @@ namespace extension {
 class FlatTensorDataMap final
     : public executorch::ET_RUNTIME_NAMESPACE::NamedDataMap {
  public:
+  /**
+   * The highest data schema version that this runtime can read.
+   *
+   * Keep in sync with kSchemaVersion in
+   * //executorch/extension/flat_tensor/serialize/serialize.h and with
+   * _FLAT_TENSOR_VERSION in
+   * //executorch/extension/flat_tensor/serialize/serialize.py, which are the
+   * versions that the two writers stamp into every PTD file.
+   */
+  static constexpr uint32_t kMaxSupportedSchemaVersion = 0;
+
   /**
    * Creates a new DataMap that wraps FlatTensor data.
    *

extension/flat_tensor/serialize/serialize.h

@@ -19,6 +19,9 @@ namespace flat_tensor {
 
 /**
  * Schema version of the .ptd format. Should be kept in sync with serialize.py
+ * and with FlatTensorDataMap::kMaxSupportedSchemaVersion in
+ * //executorch/extension/flat_tensor/flat_tensor_data_map.h, which is the
+ * highest version the runtime agrees to load.
  */
 constexpr uint32_t kSchemaVersion = 0;
 

extension/flat_tensor/serialize/serialize.py

@@ -42,7 +42,10 @@
 # Alignment of the flatbuffer (after the header).
 _FLATBUFFER_ALIGNMENT: int = 16
 
-# Current version. Keep in sync with c++ version number in serialize.
+# Current version. Keep in sync with c++ version number in serialize, and with
+# FlatTensorDataMap::kMaxSupportedSchemaVersion in
+# //executorch/extension/flat_tensor/flat_tensor_data_map.h, which is the
+# highest version the runtime agrees to load.
 _FLAT_TENSOR_VERSION: int = 0
 
 

extension/flat_tensor/test/flat_tensor_data_map_test.cpp

@@ -18,14 +18,8 @@
 #include <gtest/gtest.h>
 
 using namespace ::testing;
-using executorch::extension::BufferDataLoader;
-using executorch::extension::FileDataLoader;
-using executorch::extension::FlatTensorDataMap;
-using executorch::runtime::DataLoader;
-using executorch::runtime::Error;
-using executorch::runtime::FreeableBuffer;
-using executorch::runtime::Result;
-using executorch::runtime::TensorLayout;
+using namespace executorch::extension;
+using namespace executorch::runtime;
 
 class FlatTensorDataMapTest : public ::testing::Test {
  protected:
@@ -181,3 +175,94 @@ TEST_F(FlatTensorDataMapTest, LoadAndCheckSize) {
       FlatTensorDataMap::load(&truncated_loader);
   ASSERT_EQ(truncated_program.error(), Error::InvalidExternalData);
 }
+
+namespace {
+
+constexpr size_t kAlignment = 16;
+
+size_t aligned_up(size_t size) {
+  return (size + kAlignment - 1) & ~(kAlignment - 1);
+}
+
+// Builds the smallest PTD file that FlatTensorDataMap::load() accepts, stamped
+// with the given schema version and holding no data. Follows the layout written
+// by save_ptd(): the header is embedded in the flatbuffer region, so the offset
+// to the root table shifts by the size of the header.
+std::vector<uint8_t> CreateDataWithVersion(uint32_t version) {
+  flatbuffers::FlatBufferBuilder builder;
+  auto flat_tensor = flat_tensor_flatbuffer::CreateFlatTensor(
+      builder,
+      version,
+      builder.CreateVector(
+          std::vector<
+              flatbuffers::Offset<flat_tensor_flatbuffer::DataSegment>>{}),
+      builder.CreateVector(
+          std::vector<
+              flatbuffers::Offset<flat_tensor_flatbuffer::NamedData>>{}));
+  builder.Finish(flat_tensor, flat_tensor_flatbuffer::FlatTensorIdentifier());
+
+  const uint8_t* flatbuffer = builder.GetBufferPointer();
+  const size_t flatbuffer_size = builder.GetSize();
+  const size_t header_size =
+      aligned_up(FlatTensorHeader::kHeaderExpectedLength);
+
+  std::vector<uint8_t> data;
+  auto append = [&data](const void* bytes, size_t size) {
+    const uint8_t* begin = static_cast<const uint8_t*>(bytes);
+    data.insert(data.end(), begin, begin + size);
+  };
+
+  uint32_t root_table_offset = *reinterpret_cast<const uint32_t*>(flatbuffer) +
+      static_cast<uint32_t>(header_size);
+  append(&root_table_offset, sizeof(root_table_offset));
+  append(flatbuffer + sizeof(root_table_offset), 4); // File identifier.
+
+  append(FlatTensorHeader::kMagic, sizeof(FlatTensorHeader::kMagic));
+  uint32_t header_length = FlatTensorHeader::kHeaderExpectedLength;
+  append(&header_length, sizeof(header_length));
+  uint64_t header_fields[] = {
+      header_size, // Offset to the flatbuffer.
+      flatbuffer_size,
+      header_size + aligned_up(flatbuffer_size), // Offset to the segments.
+      0, // Segment data size.
+  };
+  append(header_fields, sizeof(header_fields));
+  data.resize(sizeof(root_table_offset) + 4 + header_size, 0);
+
+  // The first eight bytes of the flatbuffer were written above, before the
+  // header.
+  append(flatbuffer + 8, flatbuffer_size - 8);
+  data.resize(header_size + aligned_up(flatbuffer_size), 0);
+
+  return data;
+}
+
+} // namespace
+
+TEST_F(FlatTensorDataMapTest, SupportedSchemaVersionLoads) {
+  std::vector<uint8_t> data =
+      CreateDataWithVersion(FlatTensorDataMap::kMaxSupportedSchemaVersion);
+
+  alignas(16) uint8_t aligned_buffer[512];
+  ASSERT_LE(data.size(), sizeof(aligned_buffer));
+  memcpy(aligned_buffer, data.data(), data.size());
+
+  BufferDataLoader loader(aligned_buffer, data.size());
+  Result<FlatTensorDataMap> data_map = FlatTensorDataMap::load(&loader);
+
+  EXPECT_EQ(data_map.error(), Error::Ok);
+}
+
+TEST_F(FlatTensorDataMapTest, NewerSchemaVersionFailsToLoad) {
+  std::vector<uint8_t> data =
+      CreateDataWithVersion(FlatTensorDataMap::kMaxSupportedSchemaVersion + 1);
+
+  alignas(16) uint8_t aligned_buffer[512];
+  ASSERT_LE(data.size(), sizeof(aligned_buffer));
+  memcpy(aligned_buffer, data.data(), data.size());
+
+  BufferDataLoader loader(aligned_buffer, data.size());
+  Result<FlatTensorDataMap> data_map = FlatTensorDataMap::load(&loader);
+
+  EXPECT_EQ(data_map.error(), Error::InvalidExternalData);
+}

runtime/executor/program.cpp

@@ -226,6 +226,21 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
   const executorch_flatbuffer::Program* flatbuffer_program =
       executorch_flatbuffer::GetProgram(program_data->data());
 
+  // The file identifier above ("ET12") is bumped by convention only on a
+  // backward-incompatible schema change, so it selects a schema family. The
+  // version is the finer gate within that family: a file written by a newer
+  // exporter is refused here instead of being misread field by field. Older
+  // files stay loadable because the schema only grows by appending optional
+  // fields.
+  ET_CHECK_OR_RETURN_ERROR(
+      flatbuffer_program->version() <= kMaxSupportedSchemaVersion,
+      InvalidProgram,
+      "Program schema version %u is newer than the highest this runtime "
+      "supports (%u). Export the model with an older ExecuTorch, or update "
+      "the runtime.",
+      flatbuffer_program->version(),
+      kMaxSupportedSchemaVersion);
+
   // Instantiate PteDataMap if named_data is present.
   const auto named_data = flatbuffer_program->named_data();
   std::optional<internal::PteDataMap> pte_data_map = std::nullopt;

runtime/executor/program.h

@@ -75,6 +75,15 @@ class Program final {
     InternalConsistency,
   };
 
+  /**
+   * The highest program schema version that this runtime can read.
+   *
+   * Keep in sync with EXECUTORCH_SCHEMA_VERSION in
+   * //executorch/exir/version.py, which is the version that the exporter stamps
+   * into every PTE file.
+   */
+  static constexpr uint32_t kMaxSupportedSchemaVersion = 0;
+
   /**
    * Loads a Program from the provided loader. The Program will hold a pointer
    * to the loader, which must outlive the returned Program instance.

runtime/executor/test/program_test.cpp

@@ -210,6 +210,96 @@ TEST_F(ProgramTest, BadMagicFailsToLoad) {
   }
 }
 
+namespace {
+
+// Builds the smallest program that Program::load() accepts, stamped with the
+// given schema version.
+std::vector<uint8_t> CreateProgramWithVersion(uint32_t version) {
+  flatbuffers::FlatBufferBuilder builder(1024);
+
+  auto plan_name = builder.CreateString("forward");
+  auto empty_values = builder.CreateVector(
+      std::vector<flatbuffers::Offset<executorch_flatbuffer::EValue>>{});
+  auto empty_inputs = builder.CreateVector(std::vector<int32_t>{});
+  auto empty_outputs = builder.CreateVector(std::vector<int32_t>{});
+  auto empty_chains = builder.CreateVector(
+      std::vector<flatbuffers::Offset<executorch_flatbuffer::Chain>>{});
+  auto empty_operators = builder.CreateVector(
+      std::vector<flatbuffers::Offset<executorch_flatbuffer::Operator>>{});
+  auto empty_delegates = builder.CreateVector(
+      std::vector<
+          flatbuffers::Offset<executorch_flatbuffer::BackendDelegate>>{});
+  auto buffer_sizes = builder.CreateVector(std::vector<int64_t>{0});
+
+  auto execution_plan = executorch_flatbuffer::CreateExecutionPlan(
+      builder,
+      plan_name,
+      /*container_meta_type=*/0,
+      empty_values,
+      empty_inputs,
+      empty_outputs,
+      empty_chains,
+      empty_operators,
+      empty_delegates,
+      buffer_sizes);
+  auto execution_plans = builder.CreateVector(
+      std::vector<flatbuffers::Offset<executorch_flatbuffer::ExecutionPlan>>{
+          execution_plan});
+
+  // A constant segment holding only the placeholder offset means "no
+  // constants", which keeps this program off the deprecated constant_buffer
+  // path that OSS builds compile out.
+  auto constant_segment = executorch_flatbuffer::CreateSubsegmentOffsets(
+      builder,
+      /*segment_index=*/0,
+      builder.CreateVector(std::vector<uint64_t>{0}));
+
+  auto program = executorch_flatbuffer::CreateProgram(
+      builder,
+      version,
+      execution_plans,
+      /*constant_buffer=*/0,
+      /*backend_delegate_data=*/0,
+      /*segments=*/0,
+      constant_segment);
+  builder.Finish(program, executorch_flatbuffer::ProgramIdentifier());
+
+  const uint8_t* data = builder.GetBufferPointer();
+  return std::vector<uint8_t>(data, data + builder.GetSize());
+}
+
+} // namespace
+
+TEST_F(ProgramTest, SupportedSchemaVersionLoads) {
+  std::vector<uint8_t> data =
+      CreateProgramWithVersion(Program::kMaxSupportedSchemaVersion);
+
+  alignas(16) uint8_t aligned_buffer[2048];
+  ASSERT_LE(data.size(), sizeof(aligned_buffer));
+  memcpy(aligned_buffer, data.data(), data.size());
+
+  BufferDataLoader data_loader(aligned_buffer, data.size());
+  Result<Program> program = Program::load(&data_loader, kDefaultVerification);
+
+  EXPECT_EQ(program.error(), Error::Ok);
+}
+
+TEST_F(ProgramTest, NewerSchemaVersionFailsToLoad) {
+  std::vector<uint8_t> data =
+      CreateProgramWithVersion(Program::kMaxSupportedSchemaVersion + 1);
+
+  alignas(16) uint8_t aligned_buffer[2048];
+  ASSERT_LE(data.size(), sizeof(aligned_buffer));
+  memcpy(aligned_buffer, data.data(), data.size());
+
+  // Use minimal verification to show that even the cheapest level catches it.
+  BufferDataLoader data_loader(aligned_buffer, data.size());
+  Result<Program> program =
+      Program::load(&data_loader, Program::Verification::Minimal);
+
+  EXPECT_EQ(program.error(), Error::InvalidProgram);
+}
+
 // These tests require ET_ENABLE_PROGRAM_VERIFICATION to be enabled.
 // In Release builds, verification is disabled by default to save binary size.
 #ifndef ET_ENABLE_PROGRAM_VERIFICATION