[executorch][schema] Add 'EXTERNAL' to DataLocation in schema

Differential Revision: D66523171

Pull Request resolved: #7191

exir/passes/replace_view_copy_with_view_pass.py

@@ -109,6 +109,7 @@ def __init__(self, base: TensorSpec, shape: List[int]) -> None:
             "mem_obj_id",
             "mem_offset",
             "dtype",  # property
+            "extra_tensor_info",  # property
         ]
 
         # Make sure _self_fields and _base_fields are disjoint

exir/schema.py

@@ -43,14 +43,20 @@ class TensorShapeDynamism(IntEnum):
     DYNAMIC_UNBOUND = 2
 
 
+class TensorDataLocation(IntEnum):
+    SEGMENT = 0
+    EXTERNAL = 1
+
+
 @dataclass
 class ExtraTensorInfo:
     """
     Check program.fbs for explanations of this enum.
     """
 
-    mutable_data_segments_idx: Optional[int] = None
+    mutable_data_segments_idx: int = 0
     fully_qualified_name: Optional[str] = None
+    location: TensorDataLocation = TensorDataLocation.SEGMENT
 
 
 @dataclass

exir/tensor.py

@@ -18,7 +18,7 @@
 import executorch.exir.schema as schema
 import torch
 from executorch.exir.error import internal_assert
-from executorch.exir.schema import ScalarType, TensorShapeDynamism
+from executorch.exir.schema import ExtraTensorInfo, ScalarType, TensorShapeDynamism
 from executorch.exir.sym_util import eval_shape
 
 
@@ -132,6 +132,7 @@ def __init__(
         is_sparse: bool = False,
         const: bool = False,
         requires_grad: bool = False,
+        extra_tensor_info: Optional[ExtraTensorInfo] = None,
     ) -> None:
         self.scalar_type = dtype
         self.const = const
@@ -146,6 +147,7 @@ def __init__(
         self.is_sparse = is_sparse
         self.init_mem_planning_fields()
         self.shape_dynamism: TensorShapeDynamism = determine_tensor_dynanism(self.shape)
+        self.extra_tensor_info = extra_tensor_info
 
     @property
     def allocated_memory(self) -> int:
@@ -346,6 +348,7 @@ def to_list(
         allocation_info=allocation_info,
         layout=layout_enum(spec.layout),
         shape_dynamism=spec.shape_dynamism,
+        extra_tensor_info=spec.extra_tensor_info,
     )
     return flatbuffer_tensor
 

schema/program.fbs

@@ -53,18 +53,32 @@ enum TensorShapeDynamism : byte {
   DYNAMIC_UNBOUND = 2,
 }
 
+// Indicates where a tensor is stored.
+enum TensorDataLocation : byte {
+  // Stored in a segment of the PTE file.
+  SEGMENT = 0,
+  // Stored outside of the PTE file.
+  EXTERNAL = 1,
+}
 
 // Table to put additional information about tensors in that is not applicable
 // to the vast majority of tensors in the vast majority of programs.
 table ExtraTensorInfo {
   // [Optional] Specifies the SubsegmentOffsets in
   //  program.mutable_data_segments that specifies where the data is located in.
   //  If not present and the data is located in a segment, then the data is in
-  //  the first index.
+  //  index zero.
   mutable_data_segments_idx: uint64;
 
   // [Optional] The unique name of the tensor. e.g. 'mod.linear.weight'
   fully_qualified_name: string;
+
+  // [Optional] Specifies where the tensor's data is stored.
+  // - SEGMENT (default): Data is stored in a segment.
+  // - EXTERNAL: Data is stored outside of the PTE file. fully_qualified_name
+  //   must be non-empty, and is used as a key to find the tensor's external
+  //   data. Tensor.data_buffer_idx is ignored.
+  location: TensorDataLocation;
 }
 
 table Tensor {