executorch_sys/cxx_bridge/
tensor_ptr.rs

1// Clippy doesnt detect the 'Safety' comments in the cxx bridge.
2#![allow(clippy::missing_safety_doc)]
3// TensorPtr_new mirrors the Cpp make_tensor_ptr signature, which takes 8 arguments.
4#![allow(clippy::too_many_arguments)]
5// The ET_-prefixed C bridge type names are not UpperCamelCase.
6#![allow(non_camel_case_types)]
7
8pub mod cxx_util {
9    /// A wrapper around `std::any::Any` that can be used in a cxx bridge.
10    ///
11    /// This struct is useful to pass any Rust object to C++ code as `Box<RustAny>`, and the C++ code will call
12    /// the destructor of the object when the `RustAny` object is dropped.
13    pub struct RustAny {
14        #[allow(unused)]
15        inner: Box<dyn std::any::Any>,
16    }
17    impl RustAny {
18        /// Create a new `RustAny` object.
19        pub fn new(inner: Box<dyn std::any::Any>) -> Self {
20            Self { inner }
21        }
22    }
23}
24
25use cxx_util::RustAny;
26
27#[cxx::bridge]
28pub(crate) mod ffi {
29
30    extern "Rust" {
31        #[namespace = "executorch_rs::cxx_util"]
32        type RustAny;
33    }
34
35    unsafe extern "C++" {
36        include!("executorch-sys/cpp/executorch_rs/cxx_bridge.hpp");
37
38        /// Redefinition of the [`ET_ScalarType`](crate::ET_ScalarType).
39        type ET_ScalarType = crate::ET_ScalarType;
40        /// Redefinition of the [`ET_TensorShapeDynamism`](crate::ET_TensorShapeDynamism).
41        type ET_TensorShapeDynamism = crate::ET_TensorShapeDynamism;
42        /// Redefinition of the [`ET_Device`](crate::ET_Device).
43        type ET_Device = crate::ET_Device;
44        /// A minimal Tensor type whose API is a source compatible subset of at::Tensor.
45        #[namespace = "executorch::aten"]
46        type Tensor;
47
48        /// Create a new tensor pointer.
49        ///
50        /// The `device` parameter sets the Tensor's device location only — no data is allocated or
51        /// copied. The caller is responsible for ensuring `data` already lives on the requested
52        /// device. To copy CPU data to a device, use `TensorPtr_clone_to` instead.
53        ///
54        /// Arguments:
55        /// - `sizes`: The dimensions of the tensor.
56        /// - `data`: A pointer to the beginning of the data buffer.
57        /// - `dim_order`: The order of the dimensions.
58        /// - `strides`: The strides of the tensor, in units of elements (not bytes).
59        /// - `scalar_type`: The scalar type of the tensor.
60        /// - `dynamism`: The dynamism of the tensor.
61        /// - `allocation`: A `Box<RustAny>` object that will be dropped when the tensor is dropped. Can be used to
62        ///    manage the lifetime of the data buffer.
63        /// - `device`: The device on which `data` resides.
64        ///
65        /// Returns a shared pointer to the tensor.
66        ///
67        /// # Safety
68        ///
69        /// The `data` pointer must be valid for the lifetime of the tensor, and accessing it according to the data
70        /// type, sizes, dim order, and strides must be valid. The `data` pointer must reside on `device`.
71        #[namespace = "executorch_rs"]
72        unsafe fn TensorPtr_new(
73            sizes: UniquePtr<CxxVector<i32>>,
74            data: *mut u8,
75            dim_order: UniquePtr<CxxVector<u8>>,
76            strides: UniquePtr<CxxVector<i32>>,
77            scalar_type: ET_ScalarType,
78            dynamism: ET_TensorShapeDynamism,
79            allocation: Box<RustAny>,
80            device: ET_Device,
81        ) -> SharedPtr<Tensor>;
82
83        /// Creates a TensorPtr that manages a new Tensor with the same properties
84        /// as the given Tensor, but with a copy of the data owned by the returned
85        /// TensorPtr, or nullptr if the original data is null.
86        ///
87        /// Arguments:
88        ///
89        /// - `tensor`: The Tensor to clone.
90        /// - `scalar_type`: The data type for the cloned tensor. The data will be
91        ///   cast from the source tensor's type.
92        ///
93        /// Returns a new TensorPtr that manages a Tensor with the specified type
94        /// and copied/cast data.
95        #[namespace = "executorch_rs"]
96        fn TensorPtr_clone(tensor: &Tensor, scalar_type: ET_ScalarType) -> SharedPtr<Tensor>;
97
98        /// Clones a TensorPtr's data onto the given target device, allocating and copying as
99        /// needed.
100        ///
101        /// The transfer direction is inferred from the source and target device: host-to-device
102        /// when `target` is an accelerator, and device-to-host when `target` is CPU. Copies use the
103        /// DeviceAllocator registered for the accelerator side; a device-backed result owns its
104        /// memory and frees it via that allocator when destroyed.
105        ///
106        /// Source and target must differ in device domain: for a CPU-to-CPU copy use
107        /// `TensorPtr_clone`, and device-to-device transfers are not supported.
108        ///
109        /// Arguments:
110        ///
111        /// - `tensor`: The source tensor whose data will be copied.
112        /// - `device`: The destination device (CPU or an accelerator).
113        ///
114        /// Returns a TensorPtr backed by `device` memory containing the copied data.
115        #[namespace = "executorch_rs"]
116        fn TensorPtr_clone_to(tensor: SharedPtr<Tensor>, device: ET_Device) -> SharedPtr<Tensor>;
117    }
118
119    impl SharedPtr<Tensor> {}
120}
executorch_sys/cxx_bridge/tensor_ptr.rs

executorch_sys/cxx_bridge/
tensor_ptr.rs
