Update on "Introduce data schema to store raw tensors"

[ghstack-poisoned]

Author: lucylq

exir/schema_data.py

@@ -7,7 +7,7 @@
 # pyre-strict
 
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import List
 
 from executorch.exir.scalar_type import ScalarType
 
@@ -20,6 +20,8 @@ class TensorMetadata:
     scalar_type: ScalarType
     dim_sizes: List[int]
     dim_order: List[bytes]
+    storage_offset: int
+    layout: int
 
     offset: int
     size: int

schema/data.fbs

@@ -9,9 +9,31 @@ table TensorMetadata {
   // The unique id used to connect the data and program.
   fully_qualified_name:string;
   scalar_type:ScalarType;
+
+  // Size of each dimension.
   dim_sizes:[int];
+
+  // Specifies in what order the dimensions are laid out in memory (from outer
+  // to inner).
+  //
+  // For example, given a rank 3 Tensor of size (3, 5, 2). If we name
+  // dimensions: [row, column, batch], then a dim_order of:
+  // - (2, 0, 1) represents a [batch, row, column] ordering where "column" is
+  //   the innermost dimension, then comes "row", and the outermost dimension is
+  //   "batch".
+  // - (0, 2, 1) represents a [row, batch, column] ordering where "column" is
+  //   the innermost dimension, then comes "batch", and the outermost dimension
+  //   is "row".
   dim_order:[ubyte];
 
+  // Offset in scalar_type elements (e.g., multiples of 4 bytes for an int
+  // scalar type) from the beginning of the tensor buffer to the beginning of
+  // the actual data. Currently, the runtime only supports a value of zero.
+  storage_offset:int;
+
+  // May not be needed.
+  layout:byte;
+
   // Tensor offsets are relative to each TensorSegment.
   // To retrieve a given tensor:
   // 1. segment_base_offset: from the file header.
@@ -56,7 +78,6 @@ table Data {
 
   // Data segments.
   segments:[DataSegment];
-
 }
 
 root_type Data;