[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: 257019860
@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,20 @@ class TensorShapeDynamism(IntEnum):
     DYNAMIC_UNBOUND = 2
 
 
+class TensorDataLocation(IntEnum):
+    CONSTANT_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.CONSTANT_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,6 +53,13 @@ enum TensorShapeDynamism : byte {
   DYNAMIC_UNBOUND = 2,
 }
 
+// Indicates where a tensor is stored.
+enum TensorDataLocation : byte {
+  // Stored in the constant segment of the PTE file.
+  CONSTANT_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.
@@ -65,6 +72,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 the constant segment
+  // (default case) or an 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: TensorDataLocation;
 }
 
 table Tensor {