From f24016d559d8aa3cebdb09d5442ca039314c2356 Mon Sep 17 00:00:00 2001 From: lucylq Date: Mon, 28 Oct 2024 15:39:28 -0700 Subject: [PATCH 1/4] Introduce data schema to store raw tensors [ghstack-poisoned] --- exir/schema_data.py | 45 ++++++++++++++++++++++++++++++++ schema/data.fbs | 62 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 107 insertions(+) create mode 100644 exir/schema_data.py create mode 100644 schema/data.fbs diff --git a/exir/schema_data.py b/exir/schema_data.py new file mode 100644 index 00000000000..affbb0e8093 --- /dev/null +++ b/exir/schema_data.py @@ -0,0 +1,45 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +from dataclasses import dataclass +from typing import List, Optional + +from executorch.exir.scalar_type import ScalarType + +# Note: check executorch/schema/data.fbs for explanations of these fields. + + +@dataclass +class TensorMetadata: + fully_qualified_name: str + scalar_type: ScalarType + dim_sizes: List[int] + dim_order: List[bytes] + + offset: int + size: int + + +@dataclass +class TensorSegment: + segment_index: int + tensor_metadata: List[TensorMetadata] + + +@dataclass +class DataSegment: + offset: int + size: int + + +@dataclass +class Data: + version: int + tensor_alignment: int + tensor_segments: List[TensorSegment] + data_segments: List[DataSegment] diff --git a/schema/data.fbs b/schema/data.fbs new file mode 100644 index 00000000000..8c059fb065b --- /dev/null +++ b/schema/data.fbs @@ -0,0 +1,62 @@ +include "scalar_type.fbs"; +namespace executorch_flatbuffer; + +// Update after BC breaking changes. +file_identifier "DT01"; +file_extension "data"; + +table TensorMetadata { + // The unique id used to connect the data and program. + fully_qualified_name:string; + scalar_type:ScalarType; + dim_sizes:[int]; + dim_order:[ubyte]; + + // Tensor offsets are relative to each TensorSegment. + // To retrieve a given tensor: + // 1. segment_base_offset: from the file header. + // 2. segment offset: segments[tensor_segments[i].segment_index].offset + // This is likely to be 0 (all the tensors in one segment). + // 3. tensor offset: tensor_segments[i].tensor_metadata[j].offset + // May need to binary search over tensor_metadata to find the matching + // tensor using fqn. + offset: uint64; + size: uint64; +} + +table TensorSegment { + // Index of the segment in Data.segments. + segment_index: uint; + + // Tensor information, including the offset and size. + tensor_metadata:[TensorMetadata]; +} + +table DataSegment { + // Segment offsets are relative to the segment base offset provided in + // the extended file header. Segments will typically be aligned in a + // way to make it possible to use mmap() to load them. + offset: uint64; + + // The size in bytes of valid data starting at the offset. The segment + // data may be followed by padding before the segment that follows it, + // to make it easier to use mmap(). + size: uint64; +} + +table Data { + // Schema version. + version:uint; + + // Alignment for each tensor. + tensor_alignment: uint32; + + // Tensor information. + tensor_segments:[TensorSegment]; + + // Data segments. + segments:[DataSegment]; + +} + +root_type Data; From 229527ebd17347d38de387c314ac26e3b0ad5525 Mon Sep 17 00:00:00 2001 From: lucylq Date: Mon, 28 Oct 2024 16:06:54 -0700 Subject: [PATCH 2/4] Update on "Introduce data schema to store raw tensors" [ghstack-poisoned] --- exir/schema_data.py | 4 +++- schema/data.fbs | 23 ++++++++++++++++++++++- 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/exir/schema_data.py b/exir/schema_data.py index affbb0e8093..03076540c51 100644 --- a/exir/schema_data.py +++ b/exir/schema_data.py @@ -7,7 +7,7 @@ # pyre-strict from dataclasses import dataclass -from typing import List, Optional +from typing import List from executorch.exir.scalar_type import ScalarType @@ -20,6 +20,8 @@ class TensorMetadata: scalar_type: ScalarType dim_sizes: List[int] dim_order: List[bytes] + storage_offset: int + layout: int offset: int size: int diff --git a/schema/data.fbs b/schema/data.fbs index 8c059fb065b..f5431974dce 100644 --- a/schema/data.fbs +++ b/schema/data.fbs @@ -9,9 +9,31 @@ table TensorMetadata { // The unique id used to connect the data and program. fully_qualified_name:string; scalar_type:ScalarType; + + // Size of each dimension. dim_sizes:[int]; + + // Specifies in what order the dimensions are laid out in memory (from outer + // to inner). + // + // For example, given a rank 3 Tensor of size (3, 5, 2). If we name + // dimensions: [row, column, batch], then a dim_order of: + // - (2, 0, 1) represents a [batch, row, column] ordering where "column" is + // the innermost dimension, then comes "row", and the outermost dimension is + // "batch". + // - (0, 2, 1) represents a [row, batch, column] ordering where "column" is + // the innermost dimension, then comes "batch", and the outermost dimension + // is "row". dim_order:[ubyte]; + // Offset in scalar_type elements (e.g., multiples of 4 bytes for an int + // scalar type) from the beginning of the tensor buffer to the beginning of + // the actual data. Currently, the runtime only supports a value of zero. + storage_offset:int; + + // May not be needed. + layout:byte; + // Tensor offsets are relative to each TensorSegment. // To retrieve a given tensor: // 1. segment_base_offset: from the file header. @@ -56,7 +78,6 @@ table Data { // Data segments. segments:[DataSegment]; - } root_type Data; From cec36f626bf02d9e88b28c0573ec186a7cb96458 Mon Sep 17 00:00:00 2001 From: lucylq Date: Tue, 29 Oct 2024 15:22:51 -0700 Subject: [PATCH 3/4] Update on "Introduce data schema to store raw tensors" Differential Revision: [D65156641](https://our.internmc.facebook.com/intern/diff/D65156641) [ghstack-poisoned] --- exir/schema_data.py | 1 - 1 file changed, 1 deletion(-) diff --git a/exir/schema_data.py b/exir/schema_data.py index 03076540c51..2c3b085b39c 100644 --- a/exir/schema_data.py +++ b/exir/schema_data.py @@ -4,7 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -# pyre-strict from dataclasses import dataclass from typing import List From 859cf529d61de5ae0fc0313d6a66059aa5c125f9 Mon Sep 17 00:00:00 2001 From: lucylq Date: Wed, 30 Oct 2024 15:55:35 -0700 Subject: [PATCH 4/4] Update on "Introduce data schema to store raw tensors" Differential Revision: [D65156641](https://our.internmc.facebook.com/intern/diff/D65156641) [ghstack-poisoned] --- schema/data.fbs | 44 +++++++++++++++----------------------------- 1 file changed, 15 insertions(+), 29 deletions(-) diff --git a/schema/data.fbs b/schema/data.fbs index f5431974dce..c2121d71c1a 100644 --- a/schema/data.fbs +++ b/schema/data.fbs @@ -3,15 +3,15 @@ namespace executorch_flatbuffer; // Update after BC breaking changes. file_identifier "DT01"; -file_extension "data"; +file_extension "ptd"; table TensorMetadata { // The unique id used to connect the data and program. - fully_qualified_name:string; - scalar_type:ScalarType; + fully_qualified_name: string; + scalar_type: ScalarType; // Size of each dimension. - dim_sizes:[int]; + dim_sizes: [int32]; // Specifies in what order the dimensions are laid out in memory (from outer // to inner). @@ -24,34 +24,19 @@ table TensorMetadata { // - (0, 2, 1) represents a [row, batch, column] ordering where "column" is // the innermost dimension, then comes "batch", and the outermost dimension // is "row". - dim_order:[ubyte]; + dim_order: [uint8]; - // Offset in scalar_type elements (e.g., multiples of 4 bytes for an int - // scalar type) from the beginning of the tensor buffer to the beginning of - // the actual data. Currently, the runtime only supports a value of zero. - storage_offset:int; - - // May not be needed. - layout:byte; + // Segment index that the tensor data is stored in. + segment_index: uint32; // Tensor offsets are relative to each TensorSegment. // To retrieve a given tensor: // 1. segment_base_offset: from the file header. - // 2. segment offset: segments[tensor_segments[i].segment_index].offset + // 2. segment offset: segments[segment_index].offset // This is likely to be 0 (all the tensors in one segment). // 3. tensor offset: tensor_segments[i].tensor_metadata[j].offset - // May need to binary search over tensor_metadata to find the matching - // tensor using fqn. + // Find the relevant index j by matching on tensor fqn. offset: uint64; - size: uint64; -} - -table TensorSegment { - // Index of the segment in Data.segments. - segment_index: uint; - - // Tensor information, including the offset and size. - tensor_metadata:[TensorMetadata]; } table DataSegment { @@ -66,18 +51,19 @@ table DataSegment { size: uint64; } -table Data { +table DataFile { // Schema version. - version:uint; + version: uint32; - // Alignment for each tensor. + // Alignment for each tensor in bytes. Offsets of the tensor provided + // in TensorMetadata.offset are aligned to tensor_alignment. tensor_alignment: uint32; // Tensor information. - tensor_segments:[TensorSegment]; + tensor_segments: [TensorMetadata]; // Data segments. - segments:[DataSegment]; + segments: [DataSegment]; } root_type Data;