Add constant_segment to schema (pytorch#1305)

Summary:
Pull Request resolved: pytorch#1305

Add `constant_segment`, `SubsegmentOffsets` to flatbuffer schema and tests. Used in subsequent diffs.

Generate fixtures with `buck2 run fbcode//executorch/exir/tests:generate_fixtures`

Issue/design: pytorch#885

Reviewed By: dbort

Differential Revision: D51595437

fbshipit-source-id: 721916bc95af62416b9c32b3032636f2b9de6359

exir/_serialize/test/test_program.py

@@ -31,6 +31,7 @@
     DataSegment,
     ExecutionPlan,
     Program,
+    SubsegmentOffsets,
 )
 from executorch.exir.tests.common import get_test_program
 
@@ -188,6 +189,7 @@ def make_execution_plan(
                 BackendDelegateInlineData(data=b"AA delegate [0,0] data"),
             ],
             segments=[],
+            constant_segment=SubsegmentOffsets(segment_index=0, offsets=[]),
         )
 
         # Demonstrate which data each delegate points to.

exir/emit/_emit_program.py

@@ -28,6 +28,7 @@
     Int,
     Program,
     String,
+    SubsegmentOffsets,
 )
 from executorch.exir.tensor import layout_enum, scalar_type_enum
 from executorch.exir.version import EXECUTORCH_SCHEMA_VERSION
@@ -204,5 +205,7 @@ def emit_program(
             backend_delegate_data=program_state.backend_delegate_data,
             # Segments may be added at serialization time.
             segments=[],
+            # Subsegment offsets may be added at serialization time.
+            constant_segment=SubsegmentOffsets(segment_index=0, offsets=[]),
         ),
     )

exir/schema.py

@@ -244,10 +244,17 @@ class DataSegment:
     size: int
 
 
+@dataclass
+class SubsegmentOffsets:
+    segment_index: int
+    offsets: List[int]
+
+
 @dataclass
 class Program:
     version: int
     execution_plan: List[ExecutionPlan]
     constant_buffer: List[Buffer]
     backend_delegate_data: List[BackendDelegateInlineData]
     segments: List[DataSegment]
+    constant_segment: SubsegmentOffsets

exir/tests/common.py

@@ -27,6 +27,7 @@
     Program,
     ScalarType,
     String,
+    SubsegmentOffsets,
     Tensor,
     TensorShapeDynamism,
 )
@@ -80,6 +81,7 @@ def get_test_program() -> Program:
         constant_buffer=[Buffer(storage=b"")],
         backend_delegate_data=[],
         segments=[],
+        constant_segment=SubsegmentOffsets(segment_index=0, offsets=[]),
     )
 
 

exir/tests/fixtures/basic_sin_max.txt

@@ -109,5 +109,11 @@
   ],
   "segments": [
 
-  ]
+  ],
+  "constant_segment": {
+    "segment_index": 0,
+    "offsets": [
+
+    ]
+  }
 }

exir/tests/fixtures/composite_delegate.txt

@@ -224,5 +224,11 @@
   ],
   "segments": [
 
-  ]
+  ],
+  "constant_segment": {
+    "segment_index": 0,
+    "offsets": [
+
+    ]
+  }
 }

schema/program.fbs

@@ -357,6 +357,16 @@ table DataSegment {
   size: uint64;
 }
 
+// Describes data offsets into a particular segment
+table SubsegmentOffsets {
+  // Index of the segment in Program.segments
+  segment_index: uint;
+
+  // Each element is an offset in bytes into the data of the segment pointed to
+  // by segment_index. Offsets must be aligned to @executorch-tensor-alignment.
+  offsets: [uint64];
+}
+
 table Program {
   // Schema version.
   version:uint;
@@ -367,7 +377,8 @@ table Program {
 
   // Tables of constant data, used for constant Values (e.g.data field of weight tensors).
   // Each constant is assigned an index into the table which are each individually aligned.
-  // 0 index is reserved to be pointed to by non-constant Tensors
+  // 0 index is reserved to be pointed to by non-constant Tensors.
+  // If this field is non-empty, constant_segment.offsets must be empty.
   constant_buffer:[Buffer];
 
   // List of delegate data. Pointed to by BackendDelegateDataReference.
@@ -376,6 +387,12 @@ table Program {
   // List of data segments that follow the Program data in this file, sorted by
   // offset. Elements in this schema can refer to these segments by index.
   segments:[DataSegment];
+
+  // Describes the offsets of each constant tensor, relative to the segment
+  // offset. If constant_segment.offsets field is non-empty, constant_buffer
+  // must be empty. constant_segment.offsets[0] is reserved to be pointed to by
+  // non-constant Tensors.
+  constant_segment: SubsegmentOffsets;
 }
 
 root_type Program;