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

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.

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

[ghstack-poisoned]

Author: lucylq

exir/emit/_emitter.py

@@ -1125,6 +1125,7 @@ def _get_empty_tensor_evalue() -> EValue:
                     data_buffer_idx=0,
                     allocation_info=None,
                     shape_dynamism=TensorShapeDynamism.STATIC,
+                    location=DataLocation.SEGMENT,
                 )
             )
 

exir/passes/replace_view_copy_with_view_pass.py

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

exir/schema.py

@@ -53,6 +53,12 @@ class ExtraTensorInfo:
     fully_qualified_name: Optional[str] = None
 
 
+class DataLocation(IntEnum):
+    INLINE = 0
+    SEGMENT = 1
+    EXTERNAL = 2
+
+
 @dataclass
 class Tensor:
     scalar_type: ScalarType
@@ -66,6 +72,7 @@ class Tensor:
 
     # check program.fbs for explanations.
     shape_dynamism: TensorShapeDynamism
+    location: DataLocation
     extra_tensor_info: Optional[ExtraTensorInfo] = None
 
 
@@ -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,12 @@
 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 (
+    DataLocation,
+    ExtraTensorInfo,
+    ScalarType,
+    TensorShapeDynamism,
+)
 from executorch.exir.sym_util import eval_shape
 
 
@@ -132,6 +137,8 @@ def __init__(
         is_sparse: bool = False,
         const: bool = False,
         requires_grad: bool = False,
+        extra_tensor_info: Optional[ExtraTensorInfo] = None,
+        location: DataLocation = DataLocation.SEGMENT,
     ) -> None:
         self.scalar_type = dtype
         self.const = const
@@ -146,6 +153,8 @@ 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
+        self.location = location
 
     @property
     def allocated_memory(self) -> int:
@@ -346,6 +355,8 @@ def to_list(
         allocation_info=allocation_info,
         layout=layout_enum(spec.layout),
         shape_dynamism=spec.shape_dynamism,
+        location=spec.location,
+        extra_tensor_info=spec.extra_tensor_info,
     )
     return flatbuffer_tensor
 

exir/tests/common.py

@@ -14,6 +14,7 @@
     AllocationDetails,
     Chain,
     ContainerMetadata,
+    DataLocation,
     EValue,
     ExecutionPlan,
     Instruction,
@@ -56,6 +57,7 @@ def get_test_program() -> Program:
                                 memory_offset_low=16,
                             ),
                             shape_dynamism=TensorShapeDynamism.STATIC,
+                            location=DataLocation.SEGMENT,
                         )
                     ),
                 ],

schema/program.fbs

@@ -67,6 +67,16 @@ table ExtraTensorInfo {
   fully_qualified_name: string;
 }
 
+// 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 Tensor {
   scalar_type: ScalarType;
 
@@ -142,6 +152,12 @@ table Tensor {
   // [Optional] Additional information about the Tensor that is not applicable
   // to most tensors.
   extra_tensor_info: ExtraTensorInfo;
+
+  // 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 Int {
@@ -275,14 +291,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: