xref: /aosp_15_r20/external/executorch/extension/aten_util/aten_bridge.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 <executorch/extension/tensor/tensor.h>
12 #include <executorch/runtime/core/exec_aten/exec_aten.h>
13 
14 #include <ATen/Functions.h> // @manual=//caffe2/aten:ATen-cpu
15 #include <ATen/Tensor.h> // @manual=//caffe2/aten:ATen-core
16 #include <ATen/core/functional.h> // @manual=//caffe2/aten:ATen-core
17 #include <c10/core/ScalarTypeToTypeMeta.h> // @manual=//caffe2/c10:c10
18 
19 #include <memory>
20 #include <vector>
21 
22 namespace executorch {
23 namespace extension {
24 
25 torch::executor::ScalarType torch_to_executorch_scalar_type(
26     caffe2::TypeMeta type);
27 
28 c10::ScalarType executorch_to_torch_scalar_type(
29     torch::executor::ScalarType type);
30 
31 /*
32  * @param[in] aten_tensor Input at::Tensor
33  * @param[in,out] mutable_et ETensor whose underlying memory now will alias to
34  * aten_tensor
35  */
36 void alias_etensor_to_attensor(at::Tensor& at, torch::executor::Tensor& et);
37 
38 /*
39  * @param[in] et ETensor whose underlying memory now will alias to returned
40  * output tensor
41  * @param[ret] aten_tensor output at::Tensor
42  * Notes:
43  * It is owned by the caller of alias_attensor_to_etensor.
44  * Lifetime of tensor meta must be >= to that of the returned tensor since
45  * this function uses at::from_blob API that constructs non-owning tensor
46  * along with non-owning metadata, e.g. sizes.
47  * If such lifetime guarantees cannot be provided, returned tensor should be
48  * cloned.
49  */
50 at::Tensor alias_attensor_to_etensor(const torch::executor::Tensor& et);
51 
52 TensorPtr alias_tensor_ptr_to_attensor(at::Tensor& t);
53 
54 } // namespace extension
55 } // namespace executorch
56 
57 namespace torch {
58 namespace executor {
59 namespace util {
60 // TODO(T197294990): Remove these deprecated aliases once all users have moved
61 // to the new `::executorch` namespaces.
62 using ::executorch::extension::alias_attensor_to_etensor;
63 using ::executorch::extension::alias_etensor_to_attensor;
torchToExecuTorchScalarType(caffe2::TypeMeta type)64 inline torch::executor::ScalarType torchToExecuTorchScalarType(
65     caffe2::TypeMeta type) {
66   return ::executorch::extension::torch_to_executorch_scalar_type(type);
67 }
execuTorchtoTorchScalarType(torch::executor::ScalarType type)68 inline c10::ScalarType execuTorchtoTorchScalarType(
69     torch::executor::ScalarType type) {
70   return ::executorch::extension::executorch_to_torch_scalar_type(type);
71 }
72 } // namespace util
73 } // namespace executor
74 } // namespace torch
75 
76 // Some users refer to these as `torch::util::`.
77 namespace torch {
78 namespace util {
79 // TODO(T197294990): Remove these deprecated aliases once all users have moved
80 // to the new `::executorch` namespaces.
81 using ::torch::executor::util::alias_attensor_to_etensor;
82 using ::torch::executor::util::alias_etensor_to_attensor;
83 using ::torch::executor::util::execuTorchtoTorchScalarType;
84 using ::torch::executor::util::torchToExecuTorchScalarType;
85 } // namespace util
86 } // namespace torch
87