Update on "[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/schema.py

@@ -43,10 +43,9 @@ class TensorShapeDynamism(IntEnum):
     DYNAMIC_UNBOUND = 2
 
 
-class DataLocation(IntEnum):
-    INLINE = 0
-    SEGMENT = 1
-    EXTERNAL = 2
+class TensorDataLocation(IntEnum):
+    CONSTANT_SEGMENT = 0
+    EXTERNAL = 1
 
 
 @dataclass
@@ -57,7 +56,7 @@ class ExtraTensorInfo:
 
     mutable_data_segments_idx: int = 0
     fully_qualified_name: Optional[str] = None
-    location: Optional[DataLocation] = None
+    location: TensorDataLocation = TensorDataLocation.CONSTANT_SEGMENT
 
 
 @dataclass
@@ -230,6 +229,11 @@ class FrameList:
     items: List[Frame]
 
 
+class DataLocation(IntEnum):
+    INLINE = 0
+    SEGMENT = 1
+
+
 @dataclass
 class BackendDelegateDataReference:
     location: DataLocation

schema/program.fbs

@@ -53,14 +53,12 @@ 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,
+// 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
@@ -75,11 +73,11 @@ 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;
+  // [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 {
@@ -290,6 +288,14 @@ 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: