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

Pull Request resolved: #7191

To indicate if a tensor is external to the PTE file or not.

Currently, we can also use the existence of 'fqn' to determine if a tensor is external or not.
I think it's better to have a specific location field as fqn may be required for cases besides external tensor storage.
ghstack-source-id: 256993560
@exported-using-ghexport

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

Author: lucylq

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,21 @@ class TensorShapeDynamism(IntEnum):
     DYNAMIC_UNBOUND = 2
 
 
+class DataLocation(IntEnum):
+    INLINE = 0
+    SEGMENT = 1
+    EXTERNAL = 2
+
+
 @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: Optional[DataLocation] = None
 
 
 @dataclass
@@ -223,11 +230,6 @@ class FrameList:
     items: List[Frame]
 
 
-class DataLocation(IntEnum):
-    INLINE = 0
-    SEGMENT = 1
-
-
 @dataclass
 class BackendDelegateDataReference:
     location: DataLocation

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,6 +53,15 @@ enum TensorShapeDynamism : byte {
   DYNAMIC_UNBOUND = 2,
 }
 
+// Indicates where a piece of data is stored.
+enum DataLocation : byte {
+  // Stored directly in the flatbuffer.
+  INLINE = 0,
+  // Stored in a segment.
+  SEGMENT = 1,
+  // Stored outside of the flatbuffer.
+  EXTERNAL = 2,
+}
 
 // Table to put additional information about tensors in that is not applicable
 // to the vast majority of tensors in the vast majority of programs.
@@ -65,6 +74,12 @@ table ExtraTensorInfo {
 
   // [Optional] The unique name of the tensor. e.g. 'mod.linear.weight'
   fully_qualified_name: string;
+
+  // [Optional] Specifies where the tensor is stored; either in a segment or external
+  // location. If a constant tensor is stored externally, data_buffer_idx
+  // is not relevant; use extra_tensor_info.fully_qualified_name to match up
+  // the external tensor.
+  location: DataLocation;
 }
 
 table Tensor {
@@ -275,14 +290,6 @@ table FrameList {
   items: [Frame];
 }
 
-// Indicates where a piece of data is stored.
-enum DataLocation : byte {
-  // Stored directly in the flatbuffer.
-  INLINE = 0,
-  // Stored in a segment.
-  SEGMENT = 1,
-}
-
 // Indicates where the delegate data is stored
 table BackendDelegateDataReference {
   // Indicates which list to index into: