xref: /aosp_15_r20/external/executorch/runtime/core/data_loader.h (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1 /*
2  * Copyright (c) Meta Platforms, Inc. and affiliates.
3  * All rights reserved.
4  *
5  * This source code is licensed under the BSD-style license found in the
6  * LICENSE file in the root directory of this source tree.
7  */
8 
9 #pragma once
10 
11 #include <cstddef>
12 
13 #include <executorch/runtime/core/freeable_buffer.h>
14 #include <executorch/runtime/core/result.h>
15 #include <executorch/runtime/platform/compiler.h>
16 
17 namespace executorch {
18 namespace runtime {
19 
20 /**
21  * Loads from a data source.
22  *
23  * See //executorch/extension/data_loader for common implementations.
24  */
25 class DataLoader {
26  public:
27   /**
28    * Describes the content of the segment.
29    */
30   struct SegmentInfo {
31     /**
32      * Represents the purpose of the segment.
33      */
34     enum class Type {
35       /**
36        * Data for the actual program.
37        */
38       Program,
39       /**
40        * Holds constant tensor data.
41        */
42       Constant,
43       /**
44        * Data used for initializing a backend.
45        */
46       Backend,
47       /**
48        * Data used for initializing mutable tensors.
49        */
50       Mutable,
51     };
52 
53     /// Type of the segment.
54     Type segment_type;
55 
56     /// Index of the segment within the segment list. Undefined for program
57     /// segments.
58     size_t segment_index;
59 
60     /// An optional, null-terminated string describing the segment. For
61     /// `Backend` segments, this is the backend ID. Null for other segment
62     /// types.
63     const char* descriptor;
64 
65     SegmentInfo() = default;
66 
67     explicit SegmentInfo(
68         Type segment_type,
69         size_t segment_index = 0,
70         const char* descriptor = nullptr)
segment_typeSegmentInfo71         : segment_type(segment_type),
72           segment_index(segment_index),
73           descriptor(descriptor) {}
74   };
75 
76   virtual ~DataLoader() = default;
77 
78   /**
79    * Loads data from the underlying data source.
80    *
81    * NOTE: This must be thread-safe. If this call modifies common state, the
82    * implementation must do its own locking.
83    *
84    * @param offset The byte offset in the data source to start loading from.
85    * @param size The number of bytes to load.
86    * @param segment_info Information about the segment being loaded.
87    *
88    * @returns a `FreeableBuffer` that owns the loaded data.
89    */
90   ET_NODISCARD virtual Result<FreeableBuffer>
91   load(size_t offset, size_t size, const SegmentInfo& segment_info) const = 0;
92 
93   /**
94    * Loads data from the underlying data source into the provided buffer.
95    *
96    * NOTE: This must be thread-safe. If this call modifies common state, the
97    * implementation must do its own locking.
98    *
99    * @param offset The byte offset in the data source to start loading from.
100    * @param size The number of bytes to load.
101    * @param segment_info Information about the segment being loaded.
102    * @param buffer The buffer to load data into. Must point to at least `size`
103    * bytes of memory.
104    *
105    * @returns an Error indicating if the load was successful.
106    */
load_into(size_t offset,size_t size,const SegmentInfo & segment_info,void * buffer)107   ET_NODISCARD virtual Error load_into(
108       size_t offset,
109       size_t size,
110       const SegmentInfo& segment_info,
111       void* buffer) const {
112     // Using a stub implementation here instead of pure virtual to expand the
113     // data_loader interface in a backwards compatible way.
114     (void)buffer;
115     (void)offset;
116     (void)size;
117     (void)segment_info;
118     ET_LOG(Error, "load_into() not implemented for this data loader.");
119     return Error::NotImplemented;
120   }
121 
122   /**
123    * Returns the length of the underlying data source, typically the file size.
124    */
125   ET_NODISCARD virtual Result<size_t> size() const = 0;
126 };
127 
128 } // namespace runtime
129 } // namespace executorch
130 
131 namespace torch {
132 namespace executor {
133 // TODO(T197294990): Remove these deprecated aliases once all users have moved
134 // to the new `::executorch` namespaces.
135 using ::executorch::runtime::DataLoader;
136 } // namespace executor
137 } // namespace torch
138