xref: /aosp_15_r20/external/pytorch/torch/csrc/inductor/aoti_torch/tensor_converter.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <ATen/Tensor.h>
4 #include <torch/csrc/inductor/aoti_torch/c/shim.h>
5 
6 namespace torch::aot_inductor {
7 
8 // Functions declared here are not meant to be called from the AOTInductor
9 // generated model.so
10 
11 // unsafe_alloc_new_handles_from_tensors is used for allocating new aten
12 // tensor objects and return them as a vector of AtenTensorHandle (raw
13 // pointers), and those pointers will be stolen by model.so.
14 TORCH_API std::vector<AtenTensorHandle> unsafe_alloc_new_handles_from_tensors(
15     std::vector<at::Tensor>& tensors);
16 
17 // alloc_tensors_by_stealing_from_handles is used for creating a vector of aten
18 // tensors by stealing from an array of handles. Only the handles are stolen,
19 // and the array itself is borrowed.
20 //
21 // WARNING: Can NOT be called in model.so unless in the non-ABI-compatible mode
22 TORCH_API std::vector<at::Tensor> alloc_tensors_by_stealing_from_handles(
23     AtenTensorHandle* handles,
24     size_t length);
25 
26 } // namespace torch::aot_inductor
27