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 <algorithm>
12 #include <functional>
13 #include <memory>
14 #include <vector>
15
16 #include <executorch/runtime/core/error.h>
17 #include <executorch/runtime/core/exec_aten/exec_aten.h>
18 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
19
20 namespace executorch {
21 namespace extension {
22
23 /**
24 * A smart pointer type for managing the lifecycle of a Tensor.
25 */
26 using TensorPtr = std::shared_ptr<executorch::aten::Tensor>;
27
28 /**
29 * Creates a TensorPtr that manages a Tensor with the specified properties.
30 *
31 * @param sizes A vector specifying the size of each dimension.
32 * @param data A pointer to the data buffer.
33 * @param dim_order A vector specifying the order of dimensions.
34 * @param strides A vector specifying the strides of the tensor.
35 * @param type The scalar type of the tensor elements.
36 * @param dynamism Specifies the mutability of the tensor's shape.
37 * @param deleter A custom deleter function for managing the lifetime of the
38 * data buffer. If provided, this deleter will be called when the managed Tensor
39 * object is destroyed.
40 * @return A TensorPtr that manages the newly created Tensor.
41 */
42 TensorPtr make_tensor_ptr(
43 std::vector<executorch::aten::SizesType> sizes,
44 void* data,
45 std::vector<executorch::aten::DimOrderType> dim_order,
46 std::vector<executorch::aten::StridesType> strides,
47 const executorch::aten::ScalarType type =
48 executorch::aten::ScalarType::Float,
49 const executorch::aten::TensorShapeDynamism dynamism =
50 executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND,
51 std::function<void(void*)> deleter = nullptr);
52
53 /**
54 * Creates a TensorPtr that manages a Tensor with the specified properties.
55 *
56 * @param sizes A vector specifying the size of each dimension.
57 * @param data A pointer to the data buffer.
58 * @param type The scalar type of the tensor elements.
59 * @param dynamism Specifies the mutability of the tensor's shape.
60 * @param deleter A custom deleter function for managing the lifetime of the
61 * data buffer. If provided, this deleter will be called when the managed Tensor
62 * object is destroyed.
63 * @return A TensorPtr that manages the newly created Tensor.
64 */
65 inline TensorPtr make_tensor_ptr(
66 std::vector<executorch::aten::SizesType> sizes,
67 void* data,
68 const executorch::aten::ScalarType type =
69 executorch::aten::ScalarType::Float,
70 const executorch::aten::TensorShapeDynamism dynamism =
71 executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND,
72 std::function<void(void*)> deleter = nullptr) {
73 return make_tensor_ptr(
74 std::move(sizes), data, {}, {}, type, dynamism, std::move(deleter));
75 }
76
77 /**
78 * Creates a TensorPtr that manages a Tensor with the specified properties.
79 *
80 * @param sizes A vector specifying the size of each dimension.
81 * @param data A pointer to the data buffer.
82 * @param type The scalar type of the tensor elements.
83 * @param dynamism Specifies the mutability of the tensor's shape.
84 * @param deleter A custom deleter function for managing the lifetime of the
85 * data buffer. If provided, this deleter will be called when the managed Tensor
86 * object is destroyed.
87 * @return A TensorPtr that manages the newly created Tensor.
88 */
89 inline TensorPtr make_tensor_ptr(
90 std::vector<executorch::aten::SizesType> sizes,
91 void* data,
92 const executorch::aten::ScalarType type =
93 executorch::aten::ScalarType::Float,
94 const executorch::aten::TensorShapeDynamism dynamism =
95 executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND,
96 std::function<void(void*)> deleter = nullptr) {
97 return make_tensor_ptr(make_tensor_impl_ptr(
98 std::move(sizes), data, {}, {}, type, dynamism, std::move(deleter)));
99 }
100
101 /**
102 * Creates a TensorPtr that manages a Tensor with the specified properties.
103 *
104 * This template overload is specialized for cases where the tensor data is
105 * provided as a vector. The scalar type is automatically deduced from the
106 * vector's data type. If the specified `type` differs from the deduced type of
107 * the vector's elements, and casting is allowed, the data will be cast to the
108 * specified `type`. This allows for flexible creation of tensors with data
109 * vectors of one type and a different scalar type.
110 *
111 * @tparam T The C++ type of the tensor elements, deduced from the vector.
112 * @param sizes A vector specifying the size of each dimension.
113 * @param data A vector containing the tensor's data.
114 * @param dim_order A vector specifying the order of dimensions.
115 * @param strides A vector specifying the strides of each dimension.
116 * @param type The scalar type of the tensor elements. If it differs from the
117 * deduced type, the data will be cast to this type if allowed.
118 * @param dynamism Specifies the mutability of the tensor's shape.
119 * @return A TensorPtr that manages the newly created TensorImpl.
120 */
121 template <
122 typename T = float,
123 executorch::aten::ScalarType deduced_type =
124 runtime::CppTypeToScalarType<T>::value>
125 inline TensorPtr make_tensor_ptr(
126 std::vector<executorch::aten::SizesType> sizes,
127 std::vector<T> data,
128 std::vector<executorch::aten::DimOrderType> dim_order = {},
129 std::vector<executorch::aten::StridesType> strides = {},
130 executorch::aten::ScalarType type = deduced_type,
131 executorch::aten::TensorShapeDynamism dynamism =
132 executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND) {
133 if (type != deduced_type) {
134 ET_CHECK_MSG(
135 runtime::canCast(deduced_type, type),
136 "Cannot cast deduced type to specified type.");
137 std::vector<uint8_t> casted_data(data.size() * runtime::elementSize(type));
138 ET_SWITCH_REALHBBF16_TYPES(type, nullptr, "make_tensor_ptr", CTYPE, [&] {
139 std::transform(
140 data.begin(),
141 data.end(),
142 reinterpret_cast<CTYPE*>(casted_data.data()),
143 [](const T& val) { return static_cast<CTYPE>(val); });
144 });
145 const auto raw_data_ptr = casted_data.data();
146 auto data_ptr =
147 std::make_shared<std::vector<uint8_t>>(std::move(casted_data));
148 return make_tensor_ptr(
149 std::move(sizes),
150 raw_data_ptr,
151 std::move(dim_order),
152 std::move(strides),
153 type,
154 dynamism,
155 [data_ptr = std::move(data_ptr)](void*) {});
156 }
157 const auto raw_data_ptr = data.data();
158 auto data_ptr = std::make_shared<std::vector<T>>(std::move(data));
159 return make_tensor_ptr(
160 std::move(sizes),
161 raw_data_ptr,
162 std::move(dim_order),
163 std::move(strides),
164 type,
165 dynamism,
166 [data_ptr = std::move(data_ptr)](void*) {});
167 }
168
169 /**
170 * Creates a TensorPtr that manages a Tensor with the specified properties.
171 *
172 * This template overload is specialized for cases where the tensor data is
173 * provided as a vector. The scalar type is automatically deduced from the
174 * vector's data type. If the specified `type` differs from the deduced type of
175 * the vector's elements, and casting is allowed, the data will be cast to the
176 * specified `type`. This allows for flexible creation of tensors with data
177 * vectors of one type and a different scalar type.
178 *
179 * @tparam T The C++ type of the tensor elements, deduced from the vector.
180 * @param data A vector containing the tensor's data.
181 * @param type The scalar type of the tensor elements. If it differs from the
182 * deduced type, the data will be cast to this type if allowed.
183 * @param dynamism Specifies the mutability of the tensor's shape.
184 * @return A TensorPtr that manages the newly created TensorImpl.
185 */
186 template <
187 typename T = float,
188 executorch::aten::ScalarType deduced_type =
189 runtime::CppTypeToScalarType<T>::value>
190 inline TensorPtr make_tensor_ptr(
191 std::vector<T> data,
192 executorch::aten::ScalarType type = deduced_type,
193 executorch::aten::TensorShapeDynamism dynamism =
194 executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND) {
195 std::vector<executorch::aten::SizesType> sizes{
196 executorch::aten::SizesType(data.size())};
197 return make_tensor_ptr(
198 std::move(sizes), std::move(data), {0}, {1}, type, dynamism);
199 }
200
201 /**
202 * Creates a TensorPtr that manages a Tensor with the specified properties.
203 *
204 * This template overload is specialized for cases where the tensor data is
205 * provided as an initializer list. The scalar type is automatically deduced
206 * from the initializer list's data type. If the specified `type` differs from
207 * the deduced type of the initializer list's elements, and casting is allowed,
208 * the data will be cast to the specified `type`. This allows for flexible
209 * creation of tensors with data vectors of one type and a different scalar
210 * type.
211 *
212 * @tparam T The C++ type of the tensor elements, deduced from the initializer
213 * list.
214 * @param sizes A vector specifying the size of each dimension.
215 * @param list An initializer list containing the tensor's data.
216 * @param dim_order A vector specifying the order of dimensions.
217 * @param strides A vector specifying the strides of each dimension.
218 * @param type The scalar type of the tensor elements. If it differs from the
219 * deduced type, the data will be cast to this type if allowed.
220 * @param dynamism Specifies the mutability of the tensor's shape.
221 * @return A TensorPtr that manages the newly created TensorImpl.
222 */
223 template <
224 typename T = float,
225 executorch::aten::ScalarType deduced_type =
226 runtime::CppTypeToScalarType<T>::value>
227 inline TensorPtr make_tensor_ptr(
228 std::vector<executorch::aten::SizesType> sizes,
229 std::initializer_list<T> list,
230 std::vector<executorch::aten::DimOrderType> dim_order = {},
231 std::vector<executorch::aten::StridesType> strides = {},
232 executorch::aten::ScalarType type = deduced_type,
233 executorch::aten::TensorShapeDynamism dynamism =
234 executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND) {
235 return make_tensor_ptr(
236 std::move(sizes),
237 std::vector<T>(std::move(list)),
238 std::move(dim_order),
239 std::move(strides),
240 type,
241 dynamism);
242 }
243
244 /**
245 * Creates a TensorPtr that manages a Tensor with the specified properties.
246 *
247 * This template overload allows creating a Tensor from an initializer list
248 * of data. The scalar type is automatically deduced from the type of the
249 * initializer list's elements. If the specified `type` differs from
250 * the deduced type of the initializer list's elements, and casting is allowed,
251 * the data will be cast to the specified `type`. This allows for flexible
252 * creation of tensors with data vectors of one type and a different scalar
253 * type.
254 *
255 * @tparam T The C++ type of the tensor elements, deduced from the initializer
256 * list.
257 * @param list An initializer list containing the tensor's data.
258 * @param type The scalar type of the tensor elements. If it differs from the
259 * deduced type, the data will be cast to this type if allowed.
260 * @param dynamism Specifies the mutability of the tensor's shape.
261 * @return A TensorPtr that manages the newly created TensorImpl.
262 */
263 template <
264 typename T = float,
265 executorch::aten::ScalarType deduced_type =
266 runtime::CppTypeToScalarType<T>::value>
267 inline TensorPtr make_tensor_ptr(
268 std::initializer_list<T> list,
269 executorch::aten::ScalarType type = deduced_type,
270 executorch::aten::TensorShapeDynamism dynamism =
271 executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND) {
272 std::vector<executorch::aten::SizesType> sizes{
273 executorch::aten::SizesType(list.size())};
274 return make_tensor_ptr(
275 std::move(sizes), std::move(list), {0}, {1}, type, dynamism);
276 }
277
278 /**
279 * Creates a TensorPtr that manages a Tensor with a single scalar value.
280 *
281 * @tparam T The C++ type of the scalar value.
282 * @param value The scalar value to be used for the Tensor.
283 * @return A TensorPtr that manages the newly created TensorImpl.
284 */
285 template <typename T>
make_tensor_ptr(T value)286 inline TensorPtr make_tensor_ptr(T value) {
287 return make_tensor_ptr({}, std::vector<T>{value});
288 }
289
290 /**
291 * Creates a TensorPtr that manages a Tensor with the specified properties.
292 *
293 * This overload accepts a raw memory buffer stored in a std::vector<uint8_t>
294 * and a scalar type to interpret the data. The vector is managed, and the
295 * memory's lifetime is tied to the TensorImpl.
296 *
297 * @param sizes A vector specifying the size of each dimension.
298 * @param data A vector containing the raw memory for the tensor's data.
299 * @param dim_order A vector specifying the order of dimensions.
300 * @param strides A vector specifying the strides of each dimension.
301 * @param type The scalar type of the tensor elements.
302 * @param dynamism Specifies the mutability of the tensor's shape.
303 * @return A TensorPtr managing the newly created Tensor.
304 */
305 TensorPtr make_tensor_ptr(
306 std::vector<executorch::aten::SizesType> sizes,
307 std::vector<uint8_t> data,
308 std::vector<executorch::aten::DimOrderType> dim_order,
309 std::vector<executorch::aten::StridesType> strides,
310 executorch::aten::ScalarType type = executorch::aten::ScalarType::Float,
311 executorch::aten::TensorShapeDynamism dynamism =
312 executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND);
313
314 /**
315 * Creates a TensorPtr that manages a Tensor with the specified properties.
316 *
317 * This overload accepts a raw memory buffer stored in a std::vector<uint8_t>
318 * and a scalar type to interpret the data. The vector is managed, and the
319 * memory's lifetime is tied to the TensorImpl.
320 *
321 * @param sizes A vector specifying the size of each dimension.
322 * @param data A vector containing the raw memory for the tensor's data.
323 * @param type The scalar type of the tensor elements.
324 * @param dynamism Specifies the mutability of the tensor's shape.
325 * @return A TensorPtr managing the newly created Tensor.
326 */
327 inline TensorPtr make_tensor_ptr(
328 std::vector<executorch::aten::SizesType> sizes,
329 std::vector<uint8_t> data,
330 executorch::aten::ScalarType type = executorch::aten::ScalarType::Float,
331 executorch::aten::TensorShapeDynamism dynamism =
332 executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND) {
333 return make_tensor_ptr(
334 std::move(sizes), std::move(data), {}, {}, type, dynamism);
335 }
336
337 /**
338 * Creates a TensorPtr to manage a new Tensor with the same properties
339 * as the given Tensor, sharing the same data without owning it.
340 *
341 * @param tensor The Tensor whose properties are used to create a new TensorPtr.
342 * @return A new TensorPtr managing a Tensor with the same properties as the
343 * original.
344 */
make_tensor_ptr(const executorch::aten::Tensor & tensor)345 inline TensorPtr make_tensor_ptr(const executorch::aten::Tensor& tensor) {
346 return make_tensor_ptr(
347 std::vector<executorch::aten::SizesType>(
348 tensor.sizes().begin(), tensor.sizes().end()),
349 tensor.mutable_data_ptr(),
350 #ifndef USE_ATEN_LIB
351 std::vector<executorch::aten::DimOrderType>(
352 tensor.dim_order().begin(), tensor.dim_order().end()),
353 std::vector<executorch::aten::StridesType>(
354 tensor.strides().begin(), tensor.strides().end()),
355 tensor.scalar_type(),
356 tensor.shape_dynamism()
357 #else // USE_ATEN_LIB
358 {},
359 std::vector<executorch::aten::StridesType>(
360 tensor.strides().begin(), tensor.strides().end()),
361 tensor.scalar_type()
362 #endif // USE_ATEN_LIB
363 );
364 }
365
366 /**
367 * Creates a TensorPtr that manages a Tensor with the specified properties.
368 *
369 * This overload accepts a raw memory buffer stored in a std::vector<uint8_t>
370 * and a scalar type to interpret the data. The vector is managed, and the
371 * memory's lifetime is tied to the TensorImpl.
372 *
373 * @param sizes A vector specifying the size of each dimension.
374 * @param data A vector containing the raw memory for the tensor's data.
375 * @param type The scalar type of the tensor elements.
376 * @param dynamism Specifies the mutability of the tensor's shape.
377 * @return A TensorPtr managing the newly created Tensor.
378 */
379 inline TensorPtr make_tensor_ptr(
380 std::vector<executorch::aten::SizesType> sizes,
381 std::vector<uint8_t> data,
382 executorch::aten::ScalarType type = executorch::aten::ScalarType::Float,
383 executorch::aten::TensorShapeDynamism dynamism =
384 executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND) {
385 return make_tensor_ptr(
386 make_tensor_impl_ptr(std::move(sizes), std::move(data), type, dynamism));
387 }
388
389 /**
390 * Creates a TensorPtr that manages a new Tensor with the same properties
391 * as the given Tensor, but with a copy of the data owned by the returned
392 * TensorPtr, or nullptr if the original data is null.
393 *
394 * @param tensor The Tensor to clone.
395 * @return A new TensorPtr that manages a Tensor with the same properties as the
396 * original but with copied data.
397 */
398 TensorPtr clone_tensor_ptr(const executorch::aten::Tensor& tensor);
399
400 /**
401 * Creates a new TensorPtr by cloning the given TensorPtr, copying the
402 * underlying data.
403 *
404 * @param tensor The TensorPtr to clone.
405 * @return A new TensorPtr that manages a Tensor with the same properties as the
406 * original but with copied data.
407 */
clone_tensor_ptr(const TensorPtr & tensor)408 inline TensorPtr clone_tensor_ptr(const TensorPtr& tensor) {
409 return clone_tensor_ptr(*tensor);
410 }
411
412 /**
413 * Resizes the Tensor managed by the provided TensorPtr to the new sizes.
414 *
415 * @param tensor A TensorPtr managing the Tensor to resize.
416 * @param sizes A vector representing the new sizes for each dimension.
417 * @return Error::Ok on success, or an appropriate error code on failure.
418 */
419 ET_NODISCARD
420 runtime::Error resize_tensor_ptr(
421 TensorPtr& tensor,
422 const std::vector<executorch::aten::SizesType>& sizes);
423
424 } // namespace extension
425 } // namespace executorch
426