From d295db17dd14e2b97810e850b24f3e8b0584a458 Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Wed, 26 Feb 2025 16:12:46 -0800
Subject: [PATCH] [executorch][weight sharing] Introduce NamedData to PTE
 schema

See 'Schema Changes' in the [RFC]

Differential Revision: [D69430152](https://our.internmc.facebook.com/intern/diff/D69430152/)

[ghstack-poisoned]
---
 exir/schema.py     |  7 +++++++
 schema/program.fbs | 16 ++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/exir/schema.py b/exir/schema.py
index 8e1434a2fe..7dba623aeb 100644
--- a/exir/schema.py
+++ b/exir/schema.py
@@ -290,6 +290,12 @@ class SubsegmentOffsets:
     offsets: List[int]
 
 
+@dataclass
+class NamedData:
+    key: str
+    segment_index: int
+
+
 @dataclass
 class Program:
     version: int
@@ -299,3 +305,4 @@ class Program:
     segments: List[DataSegment]
     constant_segment: SubsegmentOffsets
     mutable_data_segments: Optional[List[SubsegmentOffsets]] = None
+    named_data: Optional[List[NamedData]] = None
diff --git a/schema/program.fbs b/schema/program.fbs
index 7ab2175f8a..7308cc6319 100644
--- a/schema/program.fbs
+++ b/schema/program.fbs
@@ -431,6 +431,17 @@ table SubsegmentOffsets {
   offsets: [uint64];
 }
 
+// Attributes a name to data referenced by Program.segments. Used when data is
+// referenced by multiple users, in cases where indices are not guaranteed to
+// be consistent across the users.
+table NamedData {
+  // The unique id of the data blob.
+  key: string;
+
+  // Index of the segment in Program.segments.
+  segment_index: uint32;
+}
+
 table Program {
   // Schema version.
   version: uint;
@@ -468,6 +479,11 @@ table Program {
   // constant memory, copying it over, and then being unable to release the
   // constant segment. No two elements should point to the same segment.
   mutable_data_segments: [SubsegmentOffsets];
+
+  // [Optional] List of blobs keyed by a unique name. Note that multiple
+  // 'NamedData' entries could point to the same segment index. Stored in
+  // segments attached to the PTE file.
+  named_data: [NamedData];
 }
 
 root_type Program;