diff --git a/exir/passes/replace_view_copy_with_view_pass.py b/exir/passes/replace_view_copy_with_view_pass.py index 378b933211..b19cfbed95 100644 --- a/exir/passes/replace_view_copy_with_view_pass.py +++ b/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 diff --git a/exir/schema.py b/exir/schema.py index 9ef294abb6..17810ff693 100644 --- a/exir/schema.py +++ b/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 diff --git a/exir/tensor.py b/exir/tensor.py index 0c5218bb59..a26f15a238 100644 --- a/exir/tensor.py +++ b/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 diff --git a/schema/program.fbs b/schema/program.fbs index 064df063cf..7ab2175f8a 100644 --- a/schema/program.fbs +++ b/schema/program.fbs @@ -53,6 +53,13 @@ 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. @@ -60,11 +67,18 @@ 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 {