Introduce data schema to store raw tensors

ghstack-source-id: da15b0f
Pull Request resolved: #6540

Author: lucylq

exir/schema_data.py

@@ -0,0 +1,45 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+from executorch.exir.scalar_type import ScalarType
+
+# Note: check executorch/schema/data.fbs for explanations of these fields.
+
+
+@dataclass
+class TensorMetadata:
+    fully_qualified_name: str
+    scalar_type: ScalarType
+    dim_sizes: List[int]
+    dim_order: List[bytes]
+
+    offset: int
+    size: int
+
+
+@dataclass
+class TensorSegment:
+    segment_index: int
+    tensor_metadata: List[TensorMetadata]
+
+
+@dataclass
+class DataSegment:
+    offset: int
+    size: int
+
+
+@dataclass
+class Data:
+    version: int
+    tensor_alignment: int
+    tensor_segments: List[TensorSegment]
+    data_segments: List[DataSegment]

schema/data.fbs

@@ -0,0 +1,62 @@
+include "scalar_type.fbs";
+namespace executorch_flatbuffer;
+
+// Update after BC breaking changes.
+file_identifier "DT01";
+file_extension "data";
+
+table TensorMetadata {
+  // The unique id used to connect the data and program.
+  fully_qualified_name:string;
+  scalar_type:ScalarType;
+  dim_sizes:[int];
+  dim_order:[ubyte];
+
+  // Tensor offsets are relative to each TensorSegment.
+  // To retrieve a given tensor:
+  // 1. segment_base_offset: from the file header.
+  // 2. segment offset: segments[tensor_segments[i].segment_index].offset
+  //    This is likely to be 0 (all the tensors in one segment).
+  // 3. tensor offset: tensor_segments[i].tensor_metadata[j].offset
+  //    May need to binary search over tensor_metadata to find the matching
+  //    tensor using fqn.
+  offset: uint64;
+  size: uint64;
+}
+
+table TensorSegment {
+  // Index of the segment in Data.segments.
+  segment_index: uint;
+
+  // Tensor information, including the offset and size.
+  tensor_metadata:[TensorMetadata];
+}
+
+table DataSegment {
+  // Segment offsets are relative to the segment base offset provided in
+  // the extended file header. Segments will typically be aligned in a
+  // way to make it possible to use mmap() to load them.
+  offset: uint64;
+
+  // The size in bytes of valid data starting at the offset. The segment
+  // data may be followed by padding before the segment that follows it,
+  // to make it easier to use mmap().
+  size: uint64;
+}
+
+table Data {
+  // Schema version.
+  version:uint;
+
+  // Alignment for each tensor.
+  tensor_alignment: uint32;
+
+  // Tensor information.
+  tensor_segments:[TensorSegment];
+
+  // Data segments.
+  segments:[DataSegment];
+
+}
+
+root_type Data;