diff --git a/.github/workflows/ci_requirements.txt b/.github/workflows/ci_requirements.txt index a0ea9f0d..5359e9b9 100644 --- a/.github/workflows/ci_requirements.txt +++ b/.github/workflows/ci_requirements.txt @@ -1,4 +1,4 @@ -# Copy of executorch/requirements-dev.txt as of v1.3.1 +# Copy of executorch/requirements-dev.txt as of v1.4.0 cmake>=3.29, <4.0.0 # For building binary targets in the wheel. packaging>=24.2 # Lower bound required by setuptools @@ -10,4 +10,4 @@ zstd # Imported by resolve_buck.py. certifi # Imported by resolve_buck.py. lintrunner==0.12.7 lintrunner-adapters==0.13.0 -torch==2.12.0 +torch==2.13.0 diff --git a/README.md b/README.md index 51b65858..370c7117 100644 --- a/README.md +++ b/README.md @@ -63,7 +63,7 @@ See `example/hello_world` for a complete example. ## Build To use the library you must compile the C++ executorch library yourself, as there are many configurations that determines which modules, backends, and operations are supported. See the `executorch-sys` crate for more info. -Currently the supported Cpp executorch version is `1.3.1`. +Currently the supported Cpp executorch version is `1.4.0`. The table below shows which version of the Rust crate is compatible with which C++ ExecuTorch version: diff --git a/etc/download_headers.py b/etc/download_headers.py index 861ea1e4..967a967d 100644 --- a/etc/download_headers.py +++ b/etc/download_headers.py @@ -63,8 +63,6 @@ def main(): "extension/data_loader/mman.h", "extension/data_loader/mman_windows.h", "extension/module/bundled_module.h", # TODO - "runtime/core/device_allocator.h", # TODO - "runtime/core/device_memory_buffer.h", # TODO "extension/flat_tensor/serialize/serialize.h", "**/test/**", "**/testing_util/**", @@ -74,9 +72,7 @@ def main(): for include in includes: files.update(glob.glob(include, root_dir=tmp_headers_dir, recursive=True)) for exclude in excludes: - files.difference_update( - glob.glob(exclude, root_dir=tmp_headers_dir, recursive=True) - ) + files.difference_update(glob.glob(exclude, root_dir=tmp_headers_dir, recursive=True)) for f in files: dst = HEADERS_DIR / f if not dst.parent.exists(): diff --git a/etc/setup_dev_env.py b/etc/setup_dev_env.py index 6362cc3f..ef4dc26e 100644 --- a/etc/setup_dev_env.py +++ b/etc/setup_dev_env.py @@ -7,9 +7,7 @@ import warnings from pathlib import Path -DEV_EXECUTORCH_DIR = ( - Path(__file__).parent.parent.resolve() / "etc" / ".dev-env" / "executorch" -) +DEV_EXECUTORCH_DIR = Path(__file__).parent.parent.resolve() / "etc" / ".dev-env" / "executorch" def main(): @@ -49,16 +47,14 @@ def main(): "install", "-r", DEV_EXECUTORCH_DIR / "requirements-dev.txt", - "torch==2.12.0", + "torch==2.13.0", "--extra-index-url", "https://download.pytorch.org/whl/test/cpu", ] ) build_executorch_with_dev_cfg() - subprocess.check_call( - [sys.executable, "-m", "pip", "install", "huggingface_hub[cli]"] - ) + subprocess.check_call([sys.executable, "-m", "pip", "install", "huggingface_hub[cli]"]) def clone_executorch(): @@ -71,7 +67,7 @@ def clone_executorch(): "--depth", "1", "--branch", - "v1.3.1", + "v1.4.0", "https://github.com/pytorch/executorch.git", ".", ], @@ -94,12 +90,8 @@ def clone_executorch(): cwd=DEV_EXECUTORCH_DIR / "backends" / "apple" / "coreml" / "scripts", ) - subprocess.check_call( - ["git", "submodule", "update", "--init", "--recursive"], cwd=DEV_EXECUTORCH_DIR - ) - subprocess.check_call( - ["git", "submodule", "sync", "--recursive"], cwd=DEV_EXECUTORCH_DIR - ) + subprocess.check_call(["git", "submodule", "update", "--init", "--recursive"], cwd=DEV_EXECUTORCH_DIR) + subprocess.check_call(["git", "submodule", "sync", "--recursive"], cwd=DEV_EXECUTORCH_DIR) def build_executorch_with_dev_cfg(): diff --git a/examples/llama3/README.md b/examples/llama3/README.md index deec033e..4a7521fa 100644 --- a/examples/llama3/README.md +++ b/examples/llama3/README.md @@ -72,4 +72,4 @@ To run the example, follow these steps (note that some steps should run from the The llama model can be exported with many options, such as quantization, different data types (f32, bf16), different backends, kv caching, etc. This example use a specific set of options, as specified above. Different options require different export commands and modifications to the code and build script which you can play around with. -See the [llama README](https://github.com/pytorch/executorch/blob/v1.3.1/examples/models/llama/README.md) at the Cpp executorch repository for more details. +See the [llama README](https://github.com/pytorch/executorch/blob/v1.4.0/examples/models/llama/README.md) at the Cpp executorch repository for more details. diff --git a/executorch-sys/README.md b/executorch-sys/README.md index 77ce85d8..952869e7 100644 --- a/executorch-sys/README.md +++ b/executorch-sys/README.md @@ -4,14 +4,14 @@ For a general description of the project, see the the `executorch` crate. ## Build To build the library, you need to build the C++ library yourself first. -Currently the supported Cpp executorch version is `1.3.1`. +Currently the supported Cpp executorch version is `1.4.0`. The C++ library allow for great flexibility with many flags, customizing which modules, kernels, and extensions are built. Multiple static libraries are built, and the Rust library links to them. In the following example we build the C++ library with the necessary flags to run example `hello_world`: ```bash # Clone the C++ library cd ${EXECUTORCH_CPP_DIR} -git clone --depth 1 --branch v1.3.1 https://github.com/pytorch/executorch.git . +git clone --depth 1 --branch v1.4.0 https://github.com/pytorch/executorch.git . git submodule sync --recursive git submodule update --init --recursive diff --git a/executorch-sys/build.rs b/executorch-sys/build.rs index c0a95f8d..fbb24823 100644 --- a/executorch-sys/build.rs +++ b/executorch-sys/build.rs @@ -1,6 +1,6 @@ use std::path::{Path, PathBuf}; -// const EXECUTORCH_VERSION: &str = "1.3.1"; +// const EXECUTORCH_VERSION: &str = "1.4.0"; fn main() { // TODO: verify on runtime we use the correct version of executorch diff --git a/executorch-sys/cpp/executorch_rs/c_bridge.cpp b/executorch-sys/cpp/executorch_rs/c_bridge.cpp index 69f0b2e3..9e40c5e8 100644 --- a/executorch-sys/cpp/executorch_rs/c_bridge.cpp +++ b/executorch-sys/cpp/executorch_rs/c_bridge.cpp @@ -67,6 +67,10 @@ namespace static_assert(is_equal_layout()); static_assert(std::is_trivially_move_constructible_v); + static_assert(is_equal_layout()); + static_assert(std::is_trivially_move_constructible_v); + static_assert(is_equal_layout()); + static_assert(std::is_trivially_move_constructible_v); static_assert(is_equal_layout()); static_assert(std::is_trivially_move_constructible_v); @@ -129,6 +133,7 @@ namespace // // static_assert(std::is_trivially_move_constructible_v); + static_assert(is_equal_layout>()); static_assert(is_equal_layout()); static_assert(std::is_trivially_move_constructible_v); @@ -200,6 +205,16 @@ void *executorch_MemoryAllocator_allocate(struct ET_MemoryAllocator *self, size_ auto self_ = checked_reinterpret_cast(self); return self_->allocate(size, alignment); } +size_t executorch_MemoryAllocator_used_size(const struct ET_MemoryAllocator *self) +{ + auto self_ = checked_reinterpret_cast(self); + return self_->used_size(); +} +size_t executorch_MemoryAllocator_free_size(const struct ET_MemoryAllocator *self) +{ + auto self_ = checked_reinterpret_cast(self); + return self_->free_size(); +} struct ET_HierarchicalAllocator executorch_HierarchicalAllocator_new(struct ET_SpanSpanU8 buffers) { auto buffers_ = *checked_reinterpret_cast>>(&buffers); @@ -210,6 +225,28 @@ struct ET_HierarchicalAllocator executorch_HierarchicalAllocator_new(struct ET_S new (self_) executorch::runtime::HierarchicalAllocator(buffers_); return self; } +struct ET_HierarchicalAllocator executorch_HierarchicalAllocator_new_with_devices(struct ET_SpanSpanU8 buffers, struct ET_SpanDevice planned_buffer_devices) +{ + auto buffers_ = *checked_reinterpret_cast>>(&buffers); + ET_CHECK((void *)buffers_.begin() == (void *)buffers.data); + ET_CHECK(buffers_.size() == buffers.len); + auto devices_ = *checked_reinterpret_cast>(&planned_buffer_devices); + ET_CHECK((const void *)devices_.begin() == (const void *)planned_buffer_devices.data); + ET_CHECK(devices_.size() == planned_buffer_devices.len); + struct ET_HierarchicalAllocator self; + auto self_ = checked_reinterpret_cast(&self); + new (self_) executorch::runtime::HierarchicalAllocator(buffers_, devices_); + return self; +} +struct ET_SpanDevice executorch_HierarchicalAllocator_planned_buffer_devices(const struct ET_HierarchicalAllocator *self) +{ + auto self_ = checked_reinterpret_cast(self); + auto devices = self_->planned_buffer_devices(); + return ET_SpanDevice{ + .data = reinterpret_cast(devices.data()), + .len = devices.size(), + }; +} void executorch_HierarchicalAllocator_destructor(struct ET_HierarchicalAllocator *self) { auto self_ = checked_reinterpret_cast(self); @@ -229,6 +266,20 @@ struct ET_MemoryManager executorch_MemoryManager_new( new (self_) executorch::runtime::MemoryManager(method_allocator_, planned_memory_, temp_allocator_); return self; } +struct ET_SpanDevice executorch_MemoryManager_planned_buffer_devices(const struct ET_MemoryManager *self) +{ + auto self_ = checked_reinterpret_cast(self); + auto devices = self_->planned_buffer_devices(); + return ET_SpanDevice{ + .data = reinterpret_cast(devices.data()), + .len = devices.size(), + }; +} +bool executorch_MemoryManager_has_device_memory(const struct ET_MemoryManager *self) +{ + auto self_ = checked_reinterpret_cast(self); + return self_->has_device_memory(); +} // Loaders static executorch::runtime::DataLoader *cast_data_loader_mut(struct ET_DataLoaderRefMut loader) @@ -448,7 +499,7 @@ enum ET_ScalarType executorch_Tensor_scalar_type(struct ET_TensorRef self) struct ET_Device executorch_Tensor_device(struct ET_TensorRef self) { auto self_ = cast_tensor(self); - auto d = self_->unsafeGetTensorImpl()->device(); + auto d = self_->device(); return ET_Device{static_cast(d.type()), static_cast(d.index())}; } size_t executorch_Tensor_element_size(struct ET_TensorRef self) diff --git a/executorch-sys/cpp/executorch_rs/c_bridge.h b/executorch-sys/cpp/executorch_rs/c_bridge.h index 7297cf92..4b7ebd78 100644 --- a/executorch-sys/cpp/executorch_rs/c_bridge.h +++ b/executorch-sys/cpp/executorch_rs/c_bridge.h @@ -505,7 +505,7 @@ extern "C" }; struct ET_HierarchicalAllocator { - size_t _blob[34]; + size_t _blob[36]; }; struct ET_MemoryManager { @@ -631,6 +631,11 @@ extern "C" uint8_t *data; size_t len; }; + struct ET_SpanDevice + { + const struct ET_Device *data; + size_t len; + }; struct ET_SpanSpanU8 { struct ET_SpanU8 *data; @@ -712,12 +717,19 @@ extern "C" struct ET_MemoryAllocator executorch_MemoryAllocator_new(uint32_t size, uint8_t *base_address); void *executorch_MemoryAllocator_allocate(struct ET_MemoryAllocator *self, size_t size, size_t alignment); + size_t executorch_MemoryAllocator_used_size(const struct ET_MemoryAllocator *self); + size_t executorch_MemoryAllocator_free_size(const struct ET_MemoryAllocator *self); struct ET_HierarchicalAllocator executorch_HierarchicalAllocator_new(struct ET_SpanSpanU8 buffers); + struct ET_HierarchicalAllocator executorch_HierarchicalAllocator_new_with_devices(struct ET_SpanSpanU8 buffers, struct ET_SpanDevice planned_buffer_devices); + struct ET_SpanDevice executorch_HierarchicalAllocator_planned_buffer_devices( + const struct ET_HierarchicalAllocator *self); void executorch_HierarchicalAllocator_destructor(struct ET_HierarchicalAllocator *self); struct ET_MemoryManager executorch_MemoryManager_new( struct ET_MemoryAllocator *method_allocator, struct ET_HierarchicalAllocator *planned_memory, struct ET_MemoryAllocator *temp_allocator); + struct ET_SpanDevice executorch_MemoryManager_planned_buffer_devices(const struct ET_MemoryManager *self); + bool executorch_MemoryManager_has_device_memory(const struct ET_MemoryManager *self); // Loaders struct ET_BufferDataLoader executorch_BufferDataLoader_new(const void *data, size_t size); diff --git a/executorch-sys/cpp/executorch_rs/cxx_bridge.cpp b/executorch-sys/cpp/executorch_rs/cxx_bridge.cpp index 977d7c4c..fcb4242d 100644 --- a/executorch-sys/cpp/executorch_rs/cxx_bridge.cpp +++ b/executorch-sys/cpp/executorch_rs/cxx_bridge.cpp @@ -70,7 +70,8 @@ namespace executorch_rs std::unique_ptr> strides, ET_ScalarType scalar_type, ET_TensorShapeDynamism dynamism, - rust::Box allocation) + rust::Box allocation, + struct ET_Device device) { // std::function must be copyable, so we need to wrap the allocation in a shared_ptr std::shared_ptr> allocation_ptr = @@ -83,7 +84,10 @@ namespace executorch_rs std::move(*strides), static_cast(scalar_type), static_cast(dynamism), - [allocation_ptr = allocation_ptr](void *) mutable {}); + [allocation_ptr = allocation_ptr](void *) mutable {}, + executorch::aten::Device( + static_cast(device.type), + device.index)); } std::shared_ptr TensorPtr_clone( @@ -93,6 +97,17 @@ namespace executorch_rs return executorch::extension::clone_tensor_ptr( tensor, static_cast(scalar_type)); } + + std::shared_ptr TensorPtr_clone_to( + std::shared_ptr tensor, + struct ET_Device device) + { + return executorch::extension::clone_tensor_ptr_to( + tensor, + executorch::aten::Device( + static_cast(device.type), + device.index)); + } #endif #if defined(EXECUTORCH_RS_MODULE) diff --git a/executorch-sys/cpp/executorch_rs/cxx_bridge.hpp b/executorch-sys/cpp/executorch_rs/cxx_bridge.hpp index 51affb0a..5ec30732 100644 --- a/executorch-sys/cpp/executorch_rs/cxx_bridge.hpp +++ b/executorch-sys/cpp/executorch_rs/cxx_bridge.hpp @@ -36,10 +36,14 @@ namespace executorch_rs std::unique_ptr> strides, ET_ScalarType scalar_type, ET_TensorShapeDynamism dynamism, - rust::Box allocation); + rust::Box allocation, + struct ET_Device device); std::shared_ptr TensorPtr_clone( const executorch::aten::Tensor &tensor, ET_ScalarType scalar_type); + std::shared_ptr TensorPtr_clone_to( + std::shared_ptr tensor, + struct ET_Device device); #endif #if defined(EXECUTORCH_RS_MODULE) diff --git a/executorch-sys/src/c_bridge.rs b/executorch-sys/src/c_bridge.rs index 278707cf..a4004cc5 100644 --- a/executorch-sys/src/c_bridge.rs +++ b/executorch-sys/src/c_bridge.rs @@ -60,6 +60,7 @@ impl_ref_clone_copy!(ET_ArrayRefU8); impl_ref_clone_copy!(ET_ArrayRefUsizeType); // Span +impl_ref_clone_copy!(ET_SpanDevice); impl_ref_clone_copy!(ET_SpanI64); impl_ref_clone_copy!(ET_SpanOptionalTensor); impl_ref_clone_copy!(ET_SpanSpanU8); diff --git a/executorch-sys/src/cxx_bridge/core.rs b/executorch-sys/src/cxx_bridge/core.rs index 2adb6234..c261b0e9 100644 --- a/executorch-sys/src/cxx_bridge/core.rs +++ b/executorch-sys/src/cxx_bridge/core.rs @@ -93,3 +93,8 @@ unsafe impl ExternType for crate::ET_MemoryAllocator { type Id = type_id!("ET_MemoryAllocator"); type Kind = cxx::kind::Opaque; } + +unsafe impl ExternType for crate::ET_Device { + type Id = type_id!("ET_Device"); + type Kind = cxx::kind::Trivial; +} diff --git a/executorch-sys/src/cxx_bridge/mod.rs b/executorch-sys/src/cxx_bridge/mod.rs index 56661037..9263aeaf 100644 --- a/executorch-sys/src/cxx_bridge/mod.rs +++ b/executorch-sys/src/cxx_bridge/mod.rs @@ -17,4 +17,4 @@ pub use module::ffi::{ #[cfg(feature = "tensor-ptr")] pub(crate) mod tensor_ptr; #[cfg(feature = "tensor-ptr")] -pub use tensor_ptr::ffi::{Tensor, TensorPtr_clone, TensorPtr_new}; +pub use tensor_ptr::ffi::{Tensor, TensorPtr_clone, TensorPtr_clone_to, TensorPtr_new}; diff --git a/executorch-sys/src/cxx_bridge/tensor_ptr.rs b/executorch-sys/src/cxx_bridge/tensor_ptr.rs index ea33a491..f8f29afb 100644 --- a/executorch-sys/src/cxx_bridge/tensor_ptr.rs +++ b/executorch-sys/src/cxx_bridge/tensor_ptr.rs @@ -1,5 +1,7 @@ // Clippy doesnt detect the 'Safety' comments in the cxx bridge. #![allow(clippy::missing_safety_doc)] +// TensorPtr_new mirrors the Cpp make_tensor_ptr signature, which takes 8 arguments. +#![allow(clippy::too_many_arguments)] // The ET_-prefixed C bridge type names are not UpperCamelCase. #![allow(non_camel_case_types)] @@ -37,12 +39,18 @@ pub(crate) mod ffi { type ET_ScalarType = crate::ET_ScalarType; /// Redefinition of the [`ET_TensorShapeDynamism`](crate::ET_TensorShapeDynamism). type ET_TensorShapeDynamism = crate::ET_TensorShapeDynamism; + /// Redefinition of the [`ET_Device`](crate::ET_Device). + type ET_Device = crate::ET_Device; /// A minimal Tensor type whose API is a source compatible subset of at::Tensor. #[namespace = "executorch::aten"] type Tensor; /// Create a new tensor pointer. /// + /// The `device` parameter sets the Tensor's device location only — no data is allocated or + /// copied. The caller is responsible for ensuring `data` already lives on the requested + /// device. To copy CPU data to a device, use `TensorPtr_clone_to` instead. + /// /// Arguments: /// - `sizes`: The dimensions of the tensor. /// - `data`: A pointer to the beginning of the data buffer. @@ -52,13 +60,14 @@ pub(crate) mod ffi { /// - `dynamism`: The dynamism of the tensor. /// - `allocation`: A `Box` object that will be dropped when the tensor is dropped. Can be used to /// manage the lifetime of the data buffer. + /// - `device`: The device on which `data` resides. /// /// Returns a shared pointer to the tensor. /// /// # Safety /// /// The `data` pointer must be valid for the lifetime of the tensor, and accessing it according to the data - /// type, sizes, dim order, and strides must be valid. + /// type, sizes, dim order, and strides must be valid. The `data` pointer must reside on `device`. #[namespace = "executorch_rs"] unsafe fn TensorPtr_new( sizes: UniquePtr>, @@ -68,6 +77,7 @@ pub(crate) mod ffi { scalar_type: ET_ScalarType, dynamism: ET_TensorShapeDynamism, allocation: Box, + device: ET_Device, ) -> SharedPtr; /// Creates a TensorPtr that manages a new Tensor with the same properties @@ -84,6 +94,26 @@ pub(crate) mod ffi { /// and copied/cast data. #[namespace = "executorch_rs"] fn TensorPtr_clone(tensor: &Tensor, scalar_type: ET_ScalarType) -> SharedPtr; + + /// Clones a TensorPtr's data onto the given target device, allocating and copying as + /// needed. + /// + /// The transfer direction is inferred from the source and target device: host-to-device + /// when `target` is an accelerator, and device-to-host when `target` is CPU. Copies use the + /// DeviceAllocator registered for the accelerator side; a device-backed result owns its + /// memory and frees it via that allocator when destroyed. + /// + /// Source and target must differ in device domain: for a CPU-to-CPU copy use + /// `TensorPtr_clone`, and device-to-device transfers are not supported. + /// + /// Arguments: + /// + /// - `tensor`: The source tensor whose data will be copied. + /// - `device`: The destination device (CPU or an accelerator). + /// + /// Returns a TensorPtr backed by `device` memory containing the copied data. + #[namespace = "executorch_rs"] + fn TensorPtr_clone_to(tensor: SharedPtr, device: ET_Device) -> SharedPtr; } impl SharedPtr {} diff --git a/executorch-sys/src/lib.rs b/executorch-sys/src/lib.rs index 55bd14d5..e00fabf3 100644 --- a/executorch-sys/src/lib.rs +++ b/executorch-sys/src/lib.rs @@ -11,7 +11,7 @@ //! //! //! To build the library, you need to build the C++ library yourself first. -//! Currently the supported Cpp executorch version is `1.3.1`. +//! Currently the supported Cpp executorch version is `1.4.0`. //! The C++ library allow for great flexibility with many flags, customizing which modules, kernels, and extensions are //! built. //! Multiple static libraries are built, and the Rust library links to them. @@ -19,7 +19,7 @@ //! ```bash //! # Clone the C++ library //! cd ${EXECUTORCH_CPP_DIR} -//! git clone --depth 1 --branch v1.3.1 https://github.com/pytorch/executorch.git . +//! git clone --depth 1 --branch v1.4.0 https://github.com/pytorch/executorch.git . //! git submodule sync --recursive //! git submodule update --init --recursive //! @@ -133,7 +133,7 @@ extern crate core as std; extern crate link_cplusplus; /// The version of the ExecuTorch C++ library that this crate is compatible and linked with. -pub const EXECUTORCH_CPP_VERSION: &str = "1.3.1"; +pub const EXECUTORCH_CPP_VERSION: &str = "1.4.0"; mod c_bridge; pub use c_bridge::*; diff --git a/executorch-sys/third-party/executorch/extension/flat_tensor/flat_tensor_data_map.h b/executorch-sys/third-party/executorch/extension/flat_tensor/flat_tensor_data_map.h index 751e312f..7b66eeab 100644 --- a/executorch-sys/third-party/executorch/extension/flat_tensor/flat_tensor_data_map.h +++ b/executorch-sys/third-party/executorch/extension/flat_tensor/flat_tensor_data_map.h @@ -54,7 +54,7 @@ class FlatTensorDataMap final ET_NODISCARD executorch::runtime::Result< const executorch::ET_RUNTIME_NAMESPACE::TensorLayout> - get_tensor_layout(executorch::aten::string_view key) const override; + get_tensor_layout(std::string_view key) const override; /** * Retrieve read-only data for the specified key. @@ -65,7 +65,7 @@ class FlatTensorDataMap final */ ET_NODISCARD executorch::runtime::Result get_data( - executorch::aten::string_view key) const override; + std::string_view key) const override; /** * Loads the data of the specified tensor into the provided buffer. @@ -78,7 +78,7 @@ class FlatTensorDataMap final * @returns an Error indicating if the load was successful. */ ET_NODISCARD executorch::runtime::Error load_data_into( - executorch::aten::string_view key, + std::string_view key, void* buffer, size_t size) const override; diff --git a/executorch-sys/third-party/executorch/extension/module/module.h b/executorch-sys/third-party/executorch/extension/module/module.h index 47ead230..91c7feaa 100644 --- a/executorch-sys/third-party/executorch/extension/module/module.h +++ b/executorch-sys/third-party/executorch/extension/module/module.h @@ -18,6 +18,8 @@ #include #include +#include + #ifdef USE_ATEN_LIB #define ET_MODULE_NAMESPACE module::aten #else // !USE_ATEN_LIB @@ -716,6 +718,11 @@ class Module { struct PlannedMemory { std::vector> planned_buffers; std::vector> planned_spans; + std::vector device_buffers; + /// Per-buffer Device (type + index) metadata used by + /// HierarchicalAllocator. Owns the storage backing the device span the + /// allocator references, so it must outlive `planned_memory`. + std::vector planned_devices; std::unique_ptr planned_memory; }; std::unique_ptr make_planned_memory( @@ -723,6 +730,8 @@ class Module { std::unique_ptr make_planned_memory_with_shared_arenas( const std::vector& buffer_sizes, std::vector>& shared_arenas); + std::unique_ptr make_planned_memory_with_devices( + const ET_RUNTIME_NAMESPACE::MethodMeta& method_meta); runtime::Result> get_mem_planned_buffer_sizes( const std::string& method_name); runtime::Result> get_max_mem_planned_buffer_sizes(); diff --git a/executorch-sys/third-party/executorch/extension/tensor/tensor_ptr.h b/executorch-sys/third-party/executorch/extension/tensor/tensor_ptr.h index 0ed06cbe..ffe13cb5 100644 --- a/executorch-sys/third-party/executorch/extension/tensor/tensor_ptr.h +++ b/executorch-sys/third-party/executorch/extension/tensor/tensor_ptr.h @@ -32,8 +32,14 @@ using TensorPtr = std::shared_ptr; /** * Creates a TensorPtr that manages a Tensor with the specified properties. * + * The `device` parameter sets the Tensor's device location only — no data is + * allocated or copied. The caller is responsible for ensuring `data` already + * lives on the requested device; construct the `executorch::aten::Device` from + * the runtime environment and pass it in. To copy CPU data to a device, use + * `clone_tensor_ptr_to` instead. + * * @param sizes A vector specifying the size of each dimension. - * @param data A pointer to the data buffer. + * @param data A pointer to the data buffer (CPU or device, see device). * @param dim_order A vector specifying the order of dimensions. * @param strides A vector specifying the strides of the tensor. * @param type The scalar type of the tensor elements. @@ -41,6 +47,7 @@ using TensorPtr = std::shared_ptr; * @param deleter A custom deleter function for managing the lifetime of the * data buffer. If provided, this deleter will be called when the managed Tensor * object is destroyed. + * @param device The device on which `data` resides (default CPU). * @return A TensorPtr that manages the newly created Tensor. */ TensorPtr make_tensor_ptr( @@ -52,18 +59,23 @@ TensorPtr make_tensor_ptr( executorch::aten::ScalarType::Float, const executorch::aten::TensorShapeDynamism dynamism = executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND, - std::function deleter = nullptr); + std::function deleter = nullptr, + executorch::aten::Device device = + executorch::aten::Device(executorch::aten::DeviceType::CPU)); /** * Creates a TensorPtr that manages a Tensor with the specified properties. * + * Convenience overload for the primary factory; see the primary overload for + * device semantics. + * * @param sizes A vector specifying the size of each dimension. - * @param data A pointer to the data buffer. + * @param data A pointer to the data buffer (CPU or device, see device_type). * @param type The scalar type of the tensor elements. * @param dynamism Specifies the mutability of the tensor's shape. * @param deleter A custom deleter function for managing the lifetime of the - * data buffer. If provided, this deleter will be called when the managed Tensor - * object is destroyed. + * data buffer. + * @param device The device on which `data` resides (default CPU). * @return A TensorPtr that manages the newly created Tensor. */ inline TensorPtr make_tensor_ptr( @@ -73,9 +85,18 @@ inline TensorPtr make_tensor_ptr( executorch::aten::ScalarType::Float, const executorch::aten::TensorShapeDynamism dynamism = executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND, - std::function deleter = nullptr) { + std::function deleter = nullptr, + executorch::aten::Device device = + executorch::aten::Device(executorch::aten::DeviceType::CPU)) { return make_tensor_ptr( - std::move(sizes), data, {}, {}, type, dynamism, std::move(deleter)); + std::move(sizes), + data, + {}, + {}, + type, + dynamism, + std::move(deleter), + device); } /** @@ -88,6 +109,9 @@ inline TensorPtr make_tensor_ptr( * specified `type`. This allows for flexible creation of tensors with data * vectors of one type and a different scalar type. * + * The result is always a CPU tensor. To move it to a device, use + * `clone_tensor_ptr_to`. + * * @tparam T The C++ type of the tensor elements, deduced from the vector. * @param sizes A vector specifying the size of each dimension. * @param data A vector containing the tensor's data. @@ -177,10 +201,10 @@ inline TensorPtr make_tensor_ptr( * * This template overload is specialized for cases where the tensor data is * provided as a vector. The scalar type is automatically deduced from the - * vector's data type. If the specified `type` differs from the deduced type of - * the vector's elements, and casting is allowed, the data will be cast to the - * specified `type`. This allows for flexible creation of tensors with data - * vectors of one type and a different scalar type. + * vector's data type. + * + * The result is always a CPU tensor. To move it to a device, use + * `clone_tensor_ptr_to`. * * @tparam T The C++ type of the tensor elements, deduced from the vector. * @param data A vector containing the tensor's data. @@ -209,11 +233,10 @@ inline TensorPtr make_tensor_ptr( * * This template overload is specialized for cases where the tensor data is * provided as an initializer list. The scalar type is automatically deduced - * from the initializer list's data type. If the specified `type` differs from - * the deduced type of the initializer list's elements, and casting is allowed, - * the data will be cast to the specified `type`. This allows for flexible - * creation of tensors with data vectors of one type and a different scalar - * type. + * from the initializer list's data type. + * + * The result is always a CPU tensor. To move it to a device, use + * `clone_tensor_ptr_to`. * * @tparam T The C++ type of the tensor elements, deduced from the initializer * list. @@ -252,11 +275,10 @@ inline TensorPtr make_tensor_ptr( * * This template overload allows creating a Tensor from an initializer list * of data. The scalar type is automatically deduced from the type of the - * initializer list's elements. If the specified `type` differs from - * the deduced type of the initializer list's elements, and casting is allowed, - * the data will be cast to the specified `type`. This allows for flexible - * creation of tensors with data vectors of one type and a different scalar - * type. + * initializer list's elements. + * + * The result is always a CPU tensor. To move it to a device, use + * `clone_tensor_ptr_to`. * * @tparam T The C++ type of the tensor elements, deduced from the initializer * list. @@ -299,7 +321,8 @@ inline TensorPtr make_tensor_ptr(T value) { * * This overload accepts a raw memory buffer stored in a std::vector * and a scalar type to interpret the data. The vector is managed, and the - * memory's lifetime is tied to the TensorImpl. + * memory's lifetime is tied to the TensorImpl. The result is always a CPU + * tensor. * * @param sizes A vector specifying the size of each dimension. * @param data A vector containing the raw memory for the tensor's data. @@ -321,9 +344,8 @@ TensorPtr make_tensor_ptr( /** * Creates a TensorPtr that manages a Tensor with the specified properties. * - * This overload accepts a raw memory buffer stored in a std::vector - * and a scalar type to interpret the data. The vector is managed, and the - * memory's lifetime is tied to the TensorImpl. + * Convenience overload for the raw-buffer factory; see above. The result is + * always a CPU tensor. * * @param sizes A vector specifying the size of each dimension. * @param data A vector containing the raw memory for the tensor's data. @@ -352,6 +374,9 @@ inline TensorPtr make_tensor_ptr( * configuration. If `dim_order` is empty but `strides` is provided, `dim_order` * is left empty so the core may infer it from the provided strides. * + * This overload always aliases — it never copies. To copy a tensor's data to + * a device, use `clone_tensor_ptr_to`. + * * @param tensor The source tensor to alias. * @param sizes Optional sizes override. * @param dim_order Optional dimension order override. @@ -401,16 +426,22 @@ inline TensorPtr make_tensor_ptr( tensor.scalar_type(), #ifndef USE_ATEN_LIB tensor.shape_dynamism(), + std::move(deleter), + executorch::aten::Device(tensor.device_type(), tensor.device_index())); #else // USE_ATEN_LIB executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND, + std::move(deleter), + tensor.device()); #endif // USE_ATEN_LIB - std::move(deleter)); } /** * Convenience overload identical to make_tensor_ptr(*tensor_ptr, ...). * Keeps the original TensorPtr alive until the returned TensorPtr is destroyed. * + * This overload always aliases — it never copies. To copy a tensor's data to + * a device, use `clone_tensor_ptr_to`. + * * @param tensor_ptr The source tensor pointer to alias. * @param sizes Optional sizes override. * @param dim_order Optional dimension order override. @@ -498,6 +529,32 @@ runtime::Error resize_tensor_ptr( TensorPtr& tensor, const std::vector& sizes); +/** + * Clones a TensorPtr's data onto the given target device, allocating and + * copying as needed. + * + * The transfer direction is inferred from the source and target device: + * host-to-device when `target` is an accelerator, and device-to-host when + * `target` is CPU. Copies use the DeviceAllocator registered for the + * accelerator side; a device-backed result owns its memory and frees it via + * that allocator when destroyed. + * + * Source and target must differ in device domain: for a CPU-to-CPU copy use + * clone_tensor_ptr, and device-to-device transfers are not supported. + * + * Only available in the ExecuTorch portable build: it relies on the ExecuTorch + * DeviceAllocator, which has no equivalent in USE_ATEN_LIB builds. + * + * @param tensor The source tensor whose data will be copied. + * @param target The destination device (CPU or an accelerator). + * @return A TensorPtr backed by `target` memory containing the copied data. + */ +#ifndef USE_ATEN_LIB +TensorPtr clone_tensor_ptr_to( + const TensorPtr& tensor, + executorch::aten::Device target); +#endif // USE_ATEN_LIB + } // namespace extension } // namespace executorch diff --git a/executorch-sys/third-party/executorch/runtime/core/device_allocator.h b/executorch-sys/third-party/executorch/runtime/core/device_allocator.h new file mode 100644 index 00000000..259bdb9f --- /dev/null +++ b/executorch-sys/third-party/executorch/runtime/core/device_allocator.h @@ -0,0 +1,198 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +#include +#include +#include +#include + +namespace executorch { +namespace runtime { + +/** + * Abstract interface for device-specific memory allocation. + * + * Each device type (CUDA, etc.) provides a concrete implementation + * that handles memory allocation on that device. Implementations are + * expected to be singletons with static lifetime, registered via + * DeviceAllocatorRegistry. + */ +class DeviceAllocator { + public: + /** + * Default alignment of memory returned by allocate(). Reuses + * MemoryAllocator::kDefaultAlignment so host- and device-side allocations + * share the same baseline contract. Backends whose underlying device APIs + * already provide stronger guarantees (e.g. cudaMalloc returns 256-byte + * aligned pointers) will trivially satisfy this. + */ + static constexpr size_t kDefaultAlignment = + MemoryAllocator::kDefaultAlignment; + + virtual ~DeviceAllocator() = default; + /** + * Allocate device memory. + * + * @param nbytes Number of bytes to allocate. + * @param index The device index. + * @param alignment Minimum alignment of the returned pointer in bytes. + * Must be a power of 2. Defaults to kDefaultAlignment. + * @return A Result containing the device pointer on success, or an error. + */ + virtual Result allocate( + size_t nbytes, + etensor::DeviceIndex index, + size_t alignment = kDefaultAlignment) = 0; + + /** + * Deallocate device memory previously allocated via allocate(). + * + * @param ptr Pointer to the memory to deallocate. + * @param index The device index. + */ + virtual void deallocate(void* ptr, etensor::DeviceIndex index) = 0; + + /** + * Copy data from host memory to device memory. + * + * @param dst Destination pointer (device memory). + * @param src Source pointer (host memory). + * @param nbytes Number of bytes to copy. + * @param index The device index. + * @return Error::Ok on success, or an appropriate error code on failure. + */ + virtual Error copy_host_to_device( + void* dst, + const void* src, + size_t nbytes, + etensor::DeviceIndex index) = 0; + + /** + * Copy data from device memory to host memory. + * + * @param dst Destination pointer (host memory). + * @param src Source pointer (device memory). + * @param nbytes Number of bytes to copy. + * @param index The device index. + * @return Error::Ok on success, or an appropriate error code on failure. + */ + virtual Error copy_device_to_host( + void* dst, + const void* src, + size_t nbytes, + etensor::DeviceIndex index) = 0; + + /** + * Returns the device type this allocator handles. + */ + virtual etensor::DeviceType device_type() const = 0; +}; + +/** + * Registry for device allocators. + * + * Provides a global mapping from DeviceType to DeviceAllocator instances. + * Device allocators register themselves at static initialization time, + * and the runtime queries the registry to find the appropriate allocator + * for a given device type. + * + * Threading contract: + * - Registration is expected to happen once per device type during static + * initialization (single-threaded). The registry itself does not perform + * any locking around register_allocator()/get_allocator(), and concurrent + * registration is not supported. + * - After registration, get_allocator() is safe to call concurrently from + * multiple threads because the underlying array is never mutated again. + * - The DeviceAllocator implementation is responsible for its own + * thread-safety. When multiple Programs are loaded concurrently and each + * needs device memory, the allocator must serialize access to any shared + * state internally (similar to how XNNPACK's weight cache guards its + * internal state). The registry does not provide any synchronization on + * behalf of the allocator. + */ +class DeviceAllocatorRegistry { + public: + /** + * Returns the singleton instance of the registry. + */ + static DeviceAllocatorRegistry& instance(); + + /** + * Register an allocator. The device type is taken from + * alloc->device_type(). Each device type may only be registered once; + * attempting to register a second allocator for the same device type + * will abort. + * + * Not thread-safe. Expected to be called during static initialization. + * + * @param alloc Pointer to the allocator (must have static lifetime). + */ + void register_allocator(DeviceAllocator* alloc); + + /** + * Get the allocator for a specific device type. + * + * Safe to call concurrently with other get_allocator() calls. + * + * @param type The device type. + * @return Pointer to the allocator, or nullptr if not registered. + */ + DeviceAllocator* get_allocator(etensor::DeviceType type); + + private: + DeviceAllocatorRegistry() = default; + + // Singletons must not be copied or moved; instance() returns a reference, + // and silently shallow-copying the registry would lead to confusing bugs + // where modifications to the copy don't affect the real singleton. + DeviceAllocatorRegistry(const DeviceAllocatorRegistry&) = delete; + DeviceAllocatorRegistry& operator=(const DeviceAllocatorRegistry&) = delete; + DeviceAllocatorRegistry(DeviceAllocatorRegistry&&) = delete; + DeviceAllocatorRegistry& operator=(DeviceAllocatorRegistry&&) = delete; + + // Fixed-size array indexed by device type. This avoids dynamic allocation + // and is suitable for embedded environments. + DeviceAllocator* allocators_[etensor::kNumDeviceTypes] = {}; +}; + +// Convenience free functions + +/** + * Register a device allocator. The device type is taken from + * alloc->device_type(). See DeviceAllocatorRegistry::register_allocator() + * for the threading contract. + * + * @param alloc Pointer to the allocator (must have static lifetime). + */ +void register_device_allocator(DeviceAllocator* alloc); + +/** + * Get the device allocator for a specific device type. + * + * @param type The device type. + * @return Pointer to the allocator, or nullptr if not registered. + */ +DeviceAllocator* get_device_allocator(etensor::DeviceType type); + +} // namespace runtime +} // namespace executorch + +namespace torch { +namespace executor { +// TODO(T197294990): Remove these deprecated aliases once all users have moved +// to the new `::executorch` namespaces. +using ::executorch::runtime::DeviceAllocator; +using ::executorch::runtime::DeviceAllocatorRegistry; +using ::executorch::runtime::get_device_allocator; +using ::executorch::runtime::register_device_allocator; +} // namespace executor +} // namespace torch diff --git a/executorch-sys/third-party/executorch/runtime/core/device_memory_buffer.h b/executorch-sys/third-party/executorch/runtime/core/device_memory_buffer.h new file mode 100644 index 00000000..929460ce --- /dev/null +++ b/executorch-sys/third-party/executorch/runtime/core/device_memory_buffer.h @@ -0,0 +1,129 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +#include +#include +#include + +namespace executorch::runtime { + +/** + * RAII wrapper that owns a single device memory allocation. + * + * On destruction, calls DeviceAllocator::deallocate() to free the memory. + * This mirrors the role of std::vector for CPU planned buffers, + * but for device memory (CUDA, etc.). + * + * Move-only: cannot be copied, but can be moved to transfer ownership. + */ +class DeviceMemoryBuffer final { + public: + /** + * Creates a DeviceMemoryBuffer by allocating device memory. + * + * Looks up the DeviceAllocator for the given device type via the + * DeviceAllocatorRegistry. If no allocator is registered for the type, + * returns Error::NotFound. + * + * @param size Number of bytes to allocate. + * @param type The device type (e.g., CUDA). + * @param index The device index (e.g., 0 for cuda:0). + * @param alignment Minimum alignment of the returned pointer in bytes. + * Must be a power of 2. Defaults to DeviceAllocator::kDefaultAlignment. + * @return A Result containing the DeviceMemoryBuffer on success, or an error. + */ + static Result create( + size_t size, + etensor::DeviceType type, + etensor::DeviceIndex index = 0, + size_t alignment = DeviceAllocator::kDefaultAlignment); + + DeviceMemoryBuffer() = default; + + ~DeviceMemoryBuffer() { + if (ptr_ != nullptr && allocator_ != nullptr) { + allocator_->deallocate(ptr_, device_index_); + } + } + + // Move constructor: transfer ownership. + DeviceMemoryBuffer(DeviceMemoryBuffer&& other) noexcept + : ptr_(other.ptr_), + size_(other.size_), + allocator_(other.allocator_), + device_index_(other.device_index_) { + other.ptr_ = nullptr; + other.size_ = 0; + other.allocator_ = nullptr; + } + + // Move assignment: release current, take ownership. + DeviceMemoryBuffer& operator=(DeviceMemoryBuffer&& other) noexcept { + if (this != &other) { + if (ptr_ != nullptr && allocator_ != nullptr) { + allocator_->deallocate(ptr_, device_index_); + } + ptr_ = other.ptr_; + size_ = other.size_; + allocator_ = other.allocator_; + device_index_ = other.device_index_; + other.ptr_ = nullptr; + other.size_ = 0; + other.allocator_ = nullptr; + } + return *this; + } + + // Non-copyable. + DeviceMemoryBuffer(const DeviceMemoryBuffer&) = delete; + DeviceMemoryBuffer& operator=(const DeviceMemoryBuffer&) = delete; + + /// Returns the device pointer, or nullptr if empty/moved-from. + void* data() const { + return ptr_; + } + + /// Returns the size in bytes of the allocation. + size_t size() const { + return size_; + } + + /** + * Returns a Span wrapping the device pointer. + * + * This is intended for use with HierarchicalAllocator, which only performs + * pointer arithmetic on the span data and never dereferences it. Device + * pointers are valid for pointer arithmetic from the CPU side. + */ + Span as_span() const { + return {static_cast(ptr_), size_}; + } + + private: + DeviceMemoryBuffer( + void* ptr, + size_t size, + DeviceAllocator* allocator, + etensor::DeviceIndex device_index) + : ptr_(ptr), + size_(size), + allocator_(allocator), + device_index_(device_index) {} + + void* ptr_ = nullptr; + size_t size_ = 0; + DeviceAllocator* allocator_ = nullptr; + etensor::DeviceIndex device_index_ = 0; +}; + +} // namespace executorch::runtime diff --git a/executorch-sys/third-party/executorch/runtime/core/error.h b/executorch-sys/third-party/executorch/runtime/core/error.h index 80c2ef64..b923604c 100644 --- a/executorch-sys/third-party/executorch/runtime/core/error.h +++ b/executorch-sys/third-party/executorch/runtime/core/error.h @@ -151,8 +151,9 @@ constexpr const char* to_string(const Error error) { return "Error::RegistrationExceedingMaxKernels"; case Error::RegistrationAlreadyRegistered: return "Error::RegistrationAlreadyRegistered"; + default: + return "Error::Unknown"; } - return "Error::Unknown"; } } // namespace runtime diff --git a/executorch-sys/third-party/executorch/runtime/core/exec_aten/exec_aten.h b/executorch-sys/third-party/executorch/runtime/core/exec_aten/exec_aten.h index f539414a..ac4fb9a1 100644 --- a/executorch-sys/third-party/executorch/runtime/core/exec_aten/exec_aten.h +++ b/executorch-sys/third-party/executorch/runtime/core/exec_aten/exec_aten.h @@ -8,6 +8,9 @@ #pragma once +#include +#include + #include // @manual #include // @manual #include // @manual @@ -183,8 +186,7 @@ using quint2x4 = torch::executor::quint2x4; using IntArrayRef = torch::executor::IntArrayRef; template -using OptionalArrayRef = - torch::executor::optional>; +using OptionalArrayRef = std::optional>; using OptionalIntArrayRef = OptionalArrayRef; using torch::executor::compute_numel; diff --git a/executorch-sys/third-party/executorch/runtime/core/exec_aten/util/scalar_type_util.h b/executorch-sys/third-party/executorch/runtime/core/exec_aten/util/scalar_type_util.h index 4470d391..3e8e36b4 100644 --- a/executorch-sys/third-party/executorch/runtime/core/exec_aten/util/scalar_type_util.h +++ b/executorch-sys/third-party/executorch/runtime/core/exec_aten/util/scalar_type_util.h @@ -51,7 +51,7 @@ using ScalarType = at::ScalarType; namespace executorch { namespace aten { using ScalarType = torch::executor::ScalarType; -using string_view = torch::executor::string_view; +using string_view = std::string_view; } // namespace aten } // namespace executorch #endif // USE_ATEN_LIB @@ -916,7 +916,7 @@ struct promote_types { #define ET_INTERNAL_SWITCH(TYPE, CONTEXT, NAME, ...) \ [&] { \ const auto& _st = TYPE; \ - constexpr const char* et_switch_name = NAME; \ + const char* et_switch_name = NAME; \ (void)et_switch_name; /* Suppress unused var */ \ C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wswitch-enum") \ switch (_st) { \ diff --git a/executorch-sys/third-party/executorch/runtime/core/exec_aten/util/tensor_util.h b/executorch-sys/third-party/executorch/runtime/core/exec_aten/util/tensor_util.h index 26b97e5a..b9cbab4c 100644 --- a/executorch-sys/third-party/executorch/runtime/core/exec_aten/util/tensor_util.h +++ b/executorch-sys/third-party/executorch/runtime/core/exec_aten/util/tensor_util.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include #include // std::array #include // PRId64 @@ -932,7 +933,13 @@ inline size_t getLeadingDims( ssize_t(tensor.dim())); size_t dims = 1; for (const auto i : c10::irange(dim)) { - dims *= static_cast(tensor.size(i)); + size_t next_dims; + ET_CHECK_MSG( + !c10::mul_overflows( + dims, static_cast(tensor.size(i)), &next_dims), + "Overflow computing leading dims at dimension %zd", + (ssize_t)i); + dims = next_dims; } return dims; } @@ -949,7 +956,13 @@ inline size_t getTrailingDims( ssize_t(tensor.dim())); size_t dims = 1; for (size_t i = dim + 1; i < static_cast(tensor.dim()); ++i) { - dims *= static_cast(tensor.size(i)); + size_t next_dims; + ET_CHECK_MSG( + !c10::mul_overflows( + dims, static_cast(tensor.size(i)), &next_dims), + "Overflow computing trailing dims at dimension %zu", + i); + dims = next_dims; } return dims; } diff --git a/executorch-sys/third-party/executorch/runtime/core/hierarchical_allocator.h b/executorch-sys/third-party/executorch/runtime/core/hierarchical_allocator.h index d41b98f6..00eb5b3a 100644 --- a/executorch-sys/third-party/executorch/runtime/core/hierarchical_allocator.h +++ b/executorch-sys/third-party/executorch/runtime/core/hierarchical_allocator.h @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -34,6 +35,30 @@ class HierarchicalAllocator final { explicit HierarchicalAllocator(Span> buffers) : buffers_(buffers) {} + /** + * Constructs a new hierarchical allocator with per-buffer device metadata. + * + * @param[in] buffers Same as above. May contain a mix of CPU and device + * pointers — HierarchicalAllocator only does pointer arithmetic, so + * device pointers are valid. + * @param[in] planned_buffer_devices One entry per buffer (same count as + * `buffers`), indicating the `Device` (type + index) for each buffer. + * Different buffers can target the same device type but different + * indices (e.g., `cuda:0` vs `cuda:1`). For CPU-only programs, use the + * single-arg constructor instead. + */ + HierarchicalAllocator( + Span> buffers, + Span planned_buffer_devices) + : buffers_(buffers), planned_buffer_devices_(planned_buffer_devices) { + ET_CHECK_MSG( + planned_buffer_devices.size() == buffers.size(), + "planned_buffer_devices size (%" ET_PRIsize_t + ") must match buffers size (%" ET_PRIsize_t ")", + planned_buffer_devices.size(), + buffers.size()); + } + /** * DEPRECATED: Use spans instead. */ @@ -88,6 +113,17 @@ class HierarchicalAllocator final { return buffer.data() + offset_bytes; } + /** + * Returns per-buffer device metadata. One entry per buffer, same count as + * the `buffers` passed to the constructor. Each entry is a `Device` + * carrying both type and index, so callers can distinguish e.g. `cuda:0` + * from `cuda:1`. Empty if no device metadata was provided (CPU-only + * program). + */ + Span planned_buffer_devices() const { + return planned_buffer_devices_; + } + private: // TODO(T162089316): Remove the span array and to_spans once all users move to // spans. This array is necessary to hold the pointers and sizes that were @@ -113,6 +149,10 @@ class HierarchicalAllocator final { /// The underlying buffers. Span> buffers_; + + /// Per-buffer device metadata. Empty when no device info was provided + /// (CPU-only program). + Span planned_buffer_devices_; }; } // namespace runtime diff --git a/executorch-sys/third-party/executorch/runtime/core/memory_allocator.h b/executorch-sys/third-party/executorch/runtime/core/memory_allocator.h index 001ebd7a..4d8f8da3 100644 --- a/executorch-sys/third-party/executorch/runtime/core/memory_allocator.h +++ b/executorch-sys/third-party/executorch/runtime/core/memory_allocator.h @@ -178,6 +178,22 @@ class MemoryAllocator { return size_; } + // Returns the number of bytes currently allocated from this allocator. The + // default implementation reports the bump cursor's offset from the base + // (cur_ - begin_); subclasses backed by a different allocator should override + // this to match their own accounting. + virtual size_t used_size() const { + return static_cast(cur_ - begin_); + } + + // Returns the number of bytes still available for allocation, not accounting + // for any alignment padding a future allocation may require. The default + // implementation reports end_ - cur_; subclasses should override to stay + // consistent with used_size(). + virtual size_t free_size() const { + return static_cast(end_ - cur_); + } + // Resets the current pointer to the base address. It does nothing to // the contents. virtual void reset() { diff --git a/executorch-sys/third-party/executorch/runtime/core/named_data_map.h b/executorch-sys/third-party/executorch/runtime/core/named_data_map.h index c6b6aa4b..dbd5b21a 100644 --- a/executorch-sys/third-party/executorch/runtime/core/named_data_map.h +++ b/executorch-sys/third-party/executorch/runtime/core/named_data_map.h @@ -31,7 +31,7 @@ class NamedDataMap { * @return Result containing TensorLayout. */ ET_NODISCARD virtual Result get_tensor_layout( - executorch::aten::string_view key) const = 0; + std::string_view key) const = 0; /** * Get data by key. * @@ -39,7 +39,7 @@ class NamedDataMap { * @return Result containing a FreeableBuffer. */ ET_NODISCARD virtual Result get_data( - executorch::aten::string_view key) const = 0; + std::string_view key) const = 0; /** * Loads data corresponding to the key into the provided buffer. @@ -51,10 +51,8 @@ class NamedDataMap { * `size` bytes of memory. * @returns an Error indicating if the load was successful. */ - ET_NODISCARD virtual Error load_data_into( - executorch::aten::string_view key, - void* buffer, - size_t size) const = 0; + ET_NODISCARD virtual Error + load_data_into(std::string_view key, void* buffer, size_t size) const = 0; /** * Get the number of keys in the NamedDataMap. diff --git a/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/c10/util/BFloat16-math.h b/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/c10/util/BFloat16-math.h index 8291cd74..bce89e8a 100644 --- a/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/c10/util/BFloat16-math.h +++ b/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/c10/util/BFloat16-math.h @@ -181,7 +181,7 @@ template < typename T, typename std::enable_if_t, int> = 0> inline T rsqrt(T a) { - return 1.0 / std::sqrt(float(a)); + return 1.0f / std::sqrt(float(a)); } template < typename T, diff --git a/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/c10/util/complex_math.h b/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/c10/util/complex_math.h index d369df50..2b9bbea6 100644 --- a/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/c10/util/complex_math.h +++ b/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/c10/util/complex_math.h @@ -327,7 +327,7 @@ C10_HOST_DEVICE inline c10::complex atanh(const c10::complex& x) { template C10_HOST_DEVICE inline c10::complex log1p(const c10::complex& z) { #if defined(__APPLE__) || defined(__MACOSX) || defined(__CUDACC__) || \ - defined(__HIPCC__) + defined(__HIPCC__) || defined(__SYCL_DEVICE_ONLY__) // For Mac, the new implementation yielded a high relative error. Falling back // to the old version for now. // See https://github.com/numpy/numpy/pull/22611#issuecomment-1667945354 diff --git a/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/torch/headeronly/util/Half.h b/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/torch/headeronly/util/Half.h index a9c0b166..e5aa6226 100644 --- a/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/torch/headeronly/util/Half.h +++ b/executorch-sys/third-party/executorch/runtime/core/portable_type/c10/torch/headeronly/util/Half.h @@ -236,7 +236,7 @@ C10_HOST_DEVICE inline float fp16_ieee_to_fp32_value(uint16_t h) { /* * - Choose either results of conversion of input as a normalized number, or * as a denormalized number, depending on the input exponent. The variable - * two_w contains input exponent in bits 27-31, therefore if its smaller than + * two_w contains input exponent in bits 27-31, therefore if it's smaller than * 2**27, the input is either a denormal number, or zero. * - Combine the result of conversion of exponent and mantissa with the sign * of the input number. diff --git a/executorch-sys/third-party/executorch/runtime/core/portable_type/optional.h b/executorch-sys/third-party/executorch/runtime/core/portable_type/optional.h index 31ad06fd..deff1f1b 100644 --- a/executorch-sys/third-party/executorch/runtime/core/portable_type/optional.h +++ b/executorch-sys/third-party/executorch/runtime/core/portable_type/optional.h @@ -10,16 +10,16 @@ #include +#include + namespace executorch { namespace runtime { namespace etensor { -// NOLINTNEXTLINE(misc-unused-using-decls) -using std::nullopt; -// NOLINTNEXTLINE(misc-unused-using-decls) -using std::nullopt_t; -// NOLINTNEXTLINE(misc-unused-using-decls) -using std::optional; +template +using optional ET_DEPRECATED = std::optional; +using nullopt_t ET_DEPRECATED = std::nullopt_t; +ET_DEPRECATED inline constexpr std::nullopt_t nullopt{std::nullopt}; } // namespace etensor } // namespace runtime @@ -29,8 +29,9 @@ namespace torch { namespace executor { // TODO(T197294990): Remove these deprecated aliases once all users have moved // to the new `::executorch` namespaces. -using ::executorch::runtime::etensor::nullopt; -using ::executorch::runtime::etensor::nullopt_t; -using ::executorch::runtime::etensor::optional; +template +using optional ET_DEPRECATED = std::optional; +using nullopt_t ET_DEPRECATED = std::nullopt_t; +ET_DEPRECATED inline constexpr std::nullopt_t nullopt{std::nullopt}; } // namespace executor } // namespace torch diff --git a/executorch-sys/third-party/executorch/runtime/core/portable_type/string_view.h b/executorch-sys/third-party/executorch/runtime/core/portable_type/string_view.h index 8e28fa02..f1f25f0b 100644 --- a/executorch-sys/third-party/executorch/runtime/core/portable_type/string_view.h +++ b/executorch-sys/third-party/executorch/runtime/core/portable_type/string_view.h @@ -10,11 +10,13 @@ #include +#include + namespace executorch { namespace runtime { namespace etensor { -using std::string_view; +using string_view ET_DEPRECATED = std::string_view; } // namespace etensor } // namespace runtime @@ -24,6 +26,6 @@ namespace torch { namespace executor { // TODO(T197294990): Remove these deprecated aliases once all users have moved // to the new `::executorch` namespaces. -using ::executorch::runtime::etensor::string_view; +using string_view ET_DEPRECATED = std::string_view; } // namespace executor } // namespace torch diff --git a/executorch-sys/third-party/executorch/runtime/core/portable_type/tensor.h b/executorch-sys/third-party/executorch/runtime/core/portable_type/tensor.h index 775bccc1..f4ee2aef 100644 --- a/executorch-sys/third-party/executorch/runtime/core/portable_type/tensor.h +++ b/executorch-sys/third-party/executorch/runtime/core/portable_type/tensor.h @@ -115,6 +115,21 @@ class Tensor { return impl_->shape_dynamism(); } + /// Returns the device where tensor data resides. + Device device() const { + return impl_->device(); + } + + /// Returns the type of device where tensor data resides. + DeviceType device_type() const { + return impl_->device_type(); + } + + /// Returns the device index, or 0 if default/unspecified. + DeviceIndex device_index() const { + return impl_->device_index(); + } + /// Returns a pointer of type T to the constant underlying data blob. template inline const T* const_data_ptr() const { diff --git a/executorch-sys/third-party/executorch/runtime/core/result.h b/executorch-sys/third-party/executorch/runtime/core/result.h index 377573e6..6f8bab86 100644 --- a/executorch-sys/third-party/executorch/runtime/core/result.h +++ b/executorch-sys/third-party/executorch/runtime/core/result.h @@ -216,10 +216,15 @@ using ::executorch::runtime::Result; /** * Unwrap a Result to obtain its value. If the Result contains an error, - * propogate the error via trivial function return. + * propagate the error via trivial function return. * * Note: A function using ET_UNWRAP should itself return a Result or Error. * + * This macro expands to a GNU statement expression and is therefore used as an + * expression (e.g. `auto value = ET_UNWRAP(expr);`). It is NOT portable to + * MSVC, which does not support statement expressions. Code that must compile + * under MSVC should use ET_ASSIGN_OR_RETURN below instead. + * * @param[in] result__ Expression yielding the result to unwrap. * @param[in] ... Optional format string for the log error message and its * arguments. @@ -266,3 +271,63 @@ using ::executorch::runtime::Result; #define ET_INTERNAL_UNWRAP_8 ET_INTERNAL_UNWRAP_2 #define ET_INTERNAL_UNWRAP_9 ET_INTERNAL_UNWRAP_2 #define ET_INTERNAL_UNWRAP_10 ET_INTERNAL_UNWRAP_2 + +/** + * Assign the unwrapped value of a Result to a newly declared variable, or + * return the error via trivial function return. + * + * Unlike ET_UNWRAP (which expands to a GNU statement expression), this macro + * expands to plain statements and is therefore portable to MSVC. Prefer it in + * code that must build with MSVC. + * + * Note: A function using ET_ASSIGN_OR_RETURN should itself return a Result or + * Error. + * + * Usage: + * ET_ASSIGN_OR_RETURN(value, expr); + * ET_ASSIGN_OR_RETURN(value, expr, "log message %d", arg); + * + * @param[in] var__ Name of the variable to declare and assign the unwrapped + * value to. + * @param[in] result__ Expression yielding the result to unwrap. + * @param[in] ... Optional format string for the log error message and its + * arguments. + */ +#define ET_ASSIGN_OR_RETURN(...) \ + ET_INTERNAL_ASSIGN_OR_RETURN_EXPAND(ET_INTERNAL_ASSIGN_OR_RETURN_SELECT( \ + __VA_ARGS__, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1)(__VA_ARGS__)) + +// Internal only: Use ET_ASSIGN_OR_RETURN() instead. +#define ET_INTERNAL_ASSIGN_OR_RETURN_EXPAND(x) x + +// Internal only: Use ET_ASSIGN_OR_RETURN() instead. +#define ET_INTERNAL_ASSIGN_OR_RETURN_SELECT( \ + _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, N, ...) \ + ET_INTERNAL_ASSIGN_OR_RETURN_##N + +// Internal only: Use ET_ASSIGN_OR_RETURN() instead. +#define ET_INTERNAL_ASSIGN_OR_RETURN_2(var__, result__) \ + auto et_assign_result_##var__ = (result__); \ + if (!et_assign_result_##var__.ok()) { \ + return et_assign_result_##var__.error(); \ + } \ + auto var__ = std::move(*et_assign_result_##var__) + +// Internal only: Use ET_ASSIGN_OR_RETURN() instead. +#define ET_INTERNAL_ASSIGN_OR_RETURN_3(var__, result__, message__, ...) \ + auto et_assign_result_##var__ = (result__); \ + if (!et_assign_result_##var__.ok()) { \ + ET_LOG(Error, message__, ##__VA_ARGS__); \ + return et_assign_result_##var__.error(); \ + } \ + auto var__ = std::move(*et_assign_result_##var__) + +// Internal only: Use ET_ASSIGN_OR_RETURN() instead. +#define ET_INTERNAL_ASSIGN_OR_RETURN_4 ET_INTERNAL_ASSIGN_OR_RETURN_3 +#define ET_INTERNAL_ASSIGN_OR_RETURN_5 ET_INTERNAL_ASSIGN_OR_RETURN_3 +#define ET_INTERNAL_ASSIGN_OR_RETURN_6 ET_INTERNAL_ASSIGN_OR_RETURN_3 +#define ET_INTERNAL_ASSIGN_OR_RETURN_7 ET_INTERNAL_ASSIGN_OR_RETURN_3 +#define ET_INTERNAL_ASSIGN_OR_RETURN_8 ET_INTERNAL_ASSIGN_OR_RETURN_3 +#define ET_INTERNAL_ASSIGN_OR_RETURN_9 ET_INTERNAL_ASSIGN_OR_RETURN_3 +#define ET_INTERNAL_ASSIGN_OR_RETURN_10 ET_INTERNAL_ASSIGN_OR_RETURN_3 +#define ET_INTERNAL_ASSIGN_OR_RETURN_11 ET_INTERNAL_ASSIGN_OR_RETURN_3 diff --git a/executorch-sys/third-party/executorch/runtime/executor/memory_manager.h b/executorch-sys/third-party/executorch/runtime/executor/memory_manager.h index 42edd9f0..bc89adcd 100644 --- a/executorch-sys/third-party/executorch/runtime/executor/memory_manager.h +++ b/executorch-sys/third-party/executorch/runtime/executor/memory_manager.h @@ -10,6 +10,8 @@ #include #include +#include +#include namespace executorch { namespace runtime { @@ -42,7 +44,8 @@ class MemoryManager final { * must agree with the corresponding * `MethodMeta::num_memory_planned_buffers()` and * `MethodMeta::memory_planned_buffer_size(N)` values, which are embedded - * in the Program. + * in the Program. For device-aware programs, the per-buffer device + * metadata is owned by the HierarchicalAllocator as well. * @param[in] temp_allocator The allocator to use when allocating temporary * data during kernel or delegate execution. Must outlive the Method that * uses it. May be `nullptr` if the Method does not use kernels or @@ -105,6 +108,29 @@ class MemoryManager final { return temp_allocator_; } + /** + * Returns per-buffer device metadata. One entry per planned memory buffer, + * same count as planned_memory buffers. Empty if no device metadata was + * provided (CPU-only program) or if `planned_memory` is null. + * + * This is a thin wrapper around + * `HierarchicalAllocator::planned_buffer_devices()`. + */ + Span planned_buffer_devices() const { + if (planned_memory_ == nullptr) { + return {}; + } + return planned_memory_->planned_buffer_devices(); + } + + /** + * Returns true if any planned buffer has device metadata attached. + * When false, the memory setup is CPU-only. + */ + bool has_device_memory() const { + return planned_buffer_devices().size() > 0; + } + private: MemoryAllocator* method_allocator_; HierarchicalAllocator* planned_memory_; diff --git a/executorch-sys/third-party/executorch/runtime/executor/merged_data_map.h b/executorch-sys/third-party/executorch/runtime/executor/merged_data_map.h index d5ae9705..aae92d90 100644 --- a/executorch-sys/third-party/executorch/runtime/executor/merged_data_map.h +++ b/executorch-sys/third-party/executorch/runtime/executor/merged_data_map.h @@ -57,7 +57,7 @@ class MergedDataMap final : public NamedDataMap { */ ET_NODISCARD Result get_tensor_layout( - executorch::aten::string_view key) const override { + std::string_view key) const override { auto layout = first_->get_tensor_layout(key); if (layout.ok()) { return layout.get(); @@ -76,8 +76,7 @@ class MergedDataMap final : public NamedDataMap { * @return error if the key is not present or data cannot be loaded. */ ET_NODISCARD - Result get_data( - executorch::aten::string_view key) const override { + Result get_data(std::string_view key) const override { auto data = first_->get_data(key); if (data.error() != Error::NotFound) { return data; @@ -97,7 +96,7 @@ class MergedDataMap final : public NamedDataMap { * @returns an Error indicating if the load was successful. */ ET_NODISCARD Error load_data_into( - ET_UNUSED executorch::aten::string_view key, + ET_UNUSED std::string_view key, ET_UNUSED void* buffer, ET_UNUSED size_t size) const override { return Error::NotImplemented; diff --git a/executorch-sys/third-party/executorch/runtime/executor/pte_data_map.h b/executorch-sys/third-party/executorch/runtime/executor/pte_data_map.h index b4b46a6b..36d33ae3 100644 --- a/executorch-sys/third-party/executorch/runtime/executor/pte_data_map.h +++ b/executorch-sys/third-party/executorch/runtime/executor/pte_data_map.h @@ -79,7 +79,7 @@ class PteDataMap final : public NamedDataMap { */ ET_NODISCARD Result get_tensor_layout( - ET_UNUSED executorch::aten::string_view key) const override { + ET_UNUSED std::string_view key) const override { return Error::NotImplemented; } @@ -91,14 +91,13 @@ class PteDataMap final : public NamedDataMap { * @return error if the key is not present or data cannot be loaded. */ ET_NODISCARD - Result get_data( - executorch::aten::string_view key) const override; + Result get_data(std::string_view key) const override; /** * The PteDataMap currently does not implement load_into. */ ET_NODISCARD Error load_data_into( - ET_UNUSED executorch::aten::string_view key, + ET_UNUSED std::string_view key, ET_UNUSED void* buffer, ET_UNUSED size_t size) const override { return Error::NotImplemented; diff --git a/executorch-sys/third-party/executorch/runtime/platform/compiler.h b/executorch-sys/third-party/executorch/runtime/platform/compiler.h index edd340d1..692d590f 100644 --- a/executorch-sys/third-party/executorch/runtime/platform/compiler.h +++ b/executorch-sys/third-party/executorch/runtime/platform/compiler.h @@ -138,8 +138,14 @@ #define __has_builtin(x) (0) #endif -#if __has_builtin(__builtin_strrchr) +#if defined(__FILE_NAME__) +/// __FILE_NAME__ provides just the filename at +/// compile time, avoiding embedding full paths in the binary +#define ET_SHORT_FILENAME __FILE_NAME__ +#elif __has_builtin(__builtin_strrchr) /// Name of the source file without a directory string. +/// Note: This approach embeds the full path in .rodata even though only the +/// basename is used at runtime. __FILE_NAME__ is preferred when available. #define ET_SHORT_FILENAME (__builtin_strrchr("/" __FILE__, '/') + 1) #else #define ET_SHORT_FILENAME __FILE__ @@ -152,12 +158,17 @@ #define ET_LINE __LINE__ #endif // __has_builtin(__builtin_LINE) -#if __has_builtin(__builtin_FUNCTION) +#if defined(ET_USE_BUILTIN_FUNCTION_NAME) && ET_USE_BUILTIN_FUNCTION_NAME == 0 +/// __FUNCTION__ provides a short undecorated name, saving .rodata space +/// compared to __builtin_FUNCTION() which includes the full signature +/// (namespace, parameters, return type). +#define ET_FUNCTION __FUNCTION__ +#elif __has_builtin(__builtin_FUNCTION) /// Name of the current function as a const char[]. #define ET_FUNCTION __builtin_FUNCTION() #else #define ET_FUNCTION __FUNCTION__ -#endif // __has_builtin(__builtin_FUNCTION) +#endif // As of G3 RJ-2024.3 toolchain, zu format specifier is not supported for Xtensa #if defined(__XTENSA__) diff --git a/executorch-sys/third-party/executorch/runtime/platform/profiler.h b/executorch-sys/third-party/executorch/runtime/platform/profiler.h index d6362781..cb011bd0 100644 --- a/executorch-sys/third-party/executorch/runtime/platform/profiler.h +++ b/executorch-sys/third-party/executorch/runtime/platform/profiler.h @@ -227,8 +227,12 @@ using ::executorch::runtime::track_allocator; #define EXECUTORCH_END_PROF(token_id) \ ::executorch::runtime::end_profiling(token_id); -#define EXECUTORCH_SCOPE_PROF(name) \ - ::executorch::runtime::ExecutorchProfiler profiler(name); +#define EXECUTORCH_SCOPE_PROF_CONCAT_IMPL(a, b) a##b +#define EXECUTORCH_SCOPE_PROF_CONCAT(a, b) \ + EXECUTORCH_SCOPE_PROF_CONCAT_IMPL(a, b) +#define EXECUTORCH_SCOPE_PROF(name) \ + ::executorch::runtime::ExecutorchProfiler EXECUTORCH_SCOPE_PROF_CONCAT( \ + et_profiler_, __LINE__)(name); #define EXECUTORCH_PROFILE_INSTRUCTION_SCOPE(chain_idx, instruction_idx) \ ::executorch::runtime::ExecutorchProfilerInstructionScope \ diff --git a/executorch-sys/third-party/executorch/version.txt b/executorch-sys/third-party/executorch/version.txt index 3a3cd8cc..88c5fb89 100644 --- a/executorch-sys/third-party/executorch/version.txt +++ b/executorch-sys/third-party/executorch/version.txt @@ -1 +1 @@ -1.3.1 +1.4.0 diff --git a/executorch/README.md b/executorch/README.md index 81f432f2..7f76f1da 100644 --- a/executorch/README.md +++ b/executorch/README.md @@ -63,7 +63,7 @@ See `example/hello_world` for a complete example. ## Build To use the library you must compile the C++ executorch library yourself, as there are many configurations that determines which modules, backends, and operations are supported. See the `executorch-sys` crate for more info. -Currently the supported Cpp executorch version is `1.3.1`. +Currently the supported Cpp executorch version is `1.4.0`. ## Cargo Features diff --git a/executorch/src/data_map.rs b/executorch/src/data_map.rs index 65a421e9..af60aa29 100644 --- a/executorch/src/data_map.rs +++ b/executorch/src/data_map.rs @@ -36,13 +36,6 @@ pub trait NamedDataMap { Ok(unsafe { TensorLayout::from_raw(layout) }) } - // ET_NODISCARD virtual Result get_data( - // executorch::aten::string_view key) const = 0; - // ET_NODISCARD virtual Error load_data_into( - // executorch::aten::string_view key, - // void* buffer, - // size_t size) const = 0; - /// Get the number of keys in the NamedDataMap. fn get_num_keys(&self) -> Result { // Safety: sys::executorch_NamedDataMap_get_num_keys writes to the pointer. diff --git a/executorch/src/device.rs b/executorch/src/device.rs index fc4d5ef3..5abeb6e4 100644 --- a/executorch/src/device.rs +++ b/executorch/src/device.rs @@ -2,7 +2,7 @@ use executorch_sys as sys; -use crate::util::IntoRust; +use crate::util::{IntoCpp, IntoRust}; /// An index representing a specific device; e.g. GPU 0 vs GPU 1. pub type DeviceIndex = i8; @@ -32,41 +32,71 @@ impl IntoRust for sys::ET_DeviceType { /// Tensors carry a Device to express where their underlying data resides /// (e.g. CPU host memory vs CUDA device memory). The runtime uses this to /// dispatch memory allocation to the appropriate device allocator. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct Device { - type_: DeviceType, - index: DeviceIndex, -} +#[repr(transparent)] +#[derive(Debug, Clone, Copy)] +pub struct Device(sys::ET_Device); impl Device { /// Construct a Device from its type and index. pub fn new(type_: DeviceType, index: DeviceIndex) -> Self { - Self { type_, index } + Self(sys::ET_Device { + type_: type_.cpp(), + index, + }) } /// Returns the type of device the tensor data resides on. pub fn type_(&self) -> DeviceType { - self.type_ + self.0.type_.rs() } /// Returns the device index. pub fn index(&self) -> DeviceIndex { - self.index + self.0.index } /// Returns true if the device is of CPU type. pub fn is_cpu(&self) -> bool { - self.type_ == DeviceType::Cpu + self.0.type_.rs() == DeviceType::Cpu } } impl IntoRust for sys::ET_Device { type RsType = Device; fn rs(self) -> Device { - Device { - type_: self.type_.rs(), - index: self.index, + Device(self) + } +} +impl IntoCpp for DeviceType { + type CppType = sys::ET_DeviceType; + fn cpp(self) -> sys::ET_DeviceType { + match self { + DeviceType::Cpu => sys::ET_DeviceType::ET_DeviceType_CPU, + DeviceType::Cuda => sys::ET_DeviceType::ET_DeviceType_CUDA, } } } +impl IntoCpp for Device { + type CppType = sys::ET_Device; + fn cpp(self) -> sys::ET_Device { + self.0 + } +} +impl PartialEq for Device { + fn eq(&self, other: &Self) -> bool { + self.type_() == other.type_() && self.index() == other.index() + } +} +impl Eq for Device {} +impl Default for Device { + fn default() -> Self { + Self::new(DeviceType::Cpu, 0) + } +} + +const _: () = { + assert!(size_of::() == size_of::()); + assert!(align_of::() == align_of::()); + assert!(size_of::() == size_of::()); +}; #[cfg(test)] mod tests { diff --git a/executorch/src/lib.rs b/executorch/src/lib.rs index a5fe4985..95f5b92f 100644 --- a/executorch/src/lib.rs +++ b/executorch/src/lib.rs @@ -110,7 +110,7 @@ //! ## Build //! To use the library you must compile the C++ executorch library yourself, as there are many configurations that //! determines which modules, backends, and operations are supported. See the `executorch-sys` crate for more info. -//! Currently the supported Cpp executorch version is `1.3.1`. +//! Currently the supported Cpp executorch version is `1.4.0`. //! //! ## Embedded Systems //! The library is designed to be used both in `std` and `no_std` environments. The `no_std` environment is useful for diff --git a/executorch/src/memory.rs b/executorch/src/memory.rs index bedb9a46..5a3c7190 100644 --- a/executorch/src/memory.rs +++ b/executorch/src/memory.rs @@ -4,6 +4,7 @@ //! This enable using the library in embedded systems where dynamic memory allocation is not allowed, or when allocation //! is a performance bottleneck. +use core::alloc::Layout; use core::ops::Not; use std::cell::UnsafeCell; use std::marker::{PhantomData, PhantomPinned}; @@ -11,6 +12,7 @@ use std::mem::MaybeUninit; use std::pin::Pin; use std::ptr; +use crate::device::Device; use crate::util::Span; use executorch_sys as sys; @@ -51,6 +53,22 @@ pub trait MemoryAllocator<'a> { .not() .then(|| unsafe { std::slice::from_raw_parts_mut(ptr as *mut u8, size) }) } + + /// Returns the number of bytes currently allocated from this allocator. The + /// default implementation reports the bump cursor's offset from the base + /// (cur_ - begin_); subclasses backed by a different allocator should override + /// this to match their own accounting. + fn used_size(&self) -> usize { + unsafe { sys::executorch_MemoryAllocator_used_size(self._cpp_ptr()) } + } + + /// Returns the number of bytes still available for allocation, not accounting + /// for any alignment padding a future allocation may require. The default + /// implementation reports end_ - cur_; subclasses should override to stay + /// consistent with [`used_size`](Self::used_size). + fn free_size(&self) -> usize { + unsafe { sys::executorch_MemoryAllocator_free_size(self._cpp_ptr()) } + } } /// An extension trait for [`MemoryAllocator`] that provides convenient methods to allocate memory for Rust types. pub trait MemoryAllocatorExt<'a>: MemoryAllocator<'a> { @@ -284,6 +302,10 @@ impl<'a> HierarchicalAllocator<'a> { /// `buffers.size()` must be >= `MethodMeta::num_non_const_buffers()`. /// `buffers[N].size()` must be >= `MethodMeta::non_const_buffer_size(N)`. pub fn new(buffers: &'a mut [Span<'a, u8>]) -> Self { + assert_eq!( + Layout::new::>(), + Layout::new::() + ); // Safety: the memory layout of [Span] and [sys::ET_SpanU8] is the same. let buffers = unsafe { std::mem::transmute::<&'a mut [Span<'a, u8>], &'a mut [sys::ET_SpanU8]>(buffers) @@ -297,6 +319,66 @@ impl<'a> HierarchicalAllocator<'a> { PhantomData, ) } + + /// Constructs a new hierarchical allocator with per-buffer device metadata. + /// + /// # Arguments + /// + /// * `buffers` - Same as [`new`](Self::new). May contain a mix of CPU and device pointers — + /// HierarchicalAllocator only does pointer arithmetic, so device pointers are valid. + /// * `planned_buffer_devices` - One entry per buffer (same count as `buffers`), indicating the + /// [`Device`] (type + index) for each buffer. Different buffers can target the same device + /// type but different indices (e.g., `cuda:0` vs `cuda:1`). For CPU-only programs, use + /// [`new`](Self::new) instead. + /// + /// # Panics + /// + /// The underlying Cpp constructor aborts the process if `planned_buffer_devices.len()` does not + /// equal `buffers.len()`. + /// + /// # Safety + /// + /// The caller must ensure each buffer actually resides on the device its + /// `planned_buffer_devices` entry names. + pub unsafe fn new_with_devices( + buffers: &'a mut [Span<'a, u8>], + planned_buffer_devices: &'a [Device], + ) -> Self { + assert_eq!( + Layout::new::>(), + Layout::new::() + ); + // Safety: the memory layout of [Span] and [sys::ET_SpanU8] is the same. + let buffers = unsafe { + std::mem::transmute::<&'a mut [Span<'a, u8>], &'a mut [sys::ET_SpanU8]>(buffers) + }; + let buffers = sys::ET_SpanSpanU8 { + data: buffers.as_mut_ptr(), + len: buffers.len(), + }; + let devices = sys::ET_SpanDevice { + data: planned_buffer_devices.as_ptr().cast(), + len: planned_buffer_devices.len(), + }; + Self( + unsafe { sys::executorch_HierarchicalAllocator_new_with_devices(buffers, devices) }, + PhantomData, + ) + } + + /// Returns per-buffer device metadata. One entry per buffer, same count as the `buffers` passed + /// to the constructor. Each entry is a [`Device`] carrying both type and index, so callers can + /// distinguish e.g. `cuda:0` from `cuda:1`. Empty if no device metadata was provided (CPU-only + /// program). + pub fn planned_buffer_devices(&self) -> &[Device] { + let devices = + unsafe { sys::executorch_HierarchicalAllocator_planned_buffer_devices(&self.0) }; + if devices.len == 0 { + return &[]; // guard against null pointer + } + // Safety: the span points into this allocator, which outlives the returned borrow. + unsafe { std::slice::from_raw_parts(devices.data.cast::(), devices.len) } + } } impl Drop for HierarchicalAllocator<'_> { fn drop(&mut self) { @@ -330,7 +412,8 @@ impl<'a> MemoryManager<'a> { /// Must outlive the Method that uses it. May be [`None`] if the Method does not use any memory-planned tensor data. /// The sizes of the buffers in this HierarchicalAllocator must agree with the corresponding /// `MethodMeta::num_memory_planned_buffers()` and `MethodMeta::memory_planned_buffer_size(N)` values, - /// which are embedded in the Program. + /// which are embedded in the Program. For device-aware programs, the per-buffer device metadata + /// is owned by the [`HierarchicalAllocator`] as well. /// * `temp_allocator` - The allocator to use when allocating temporary data during kernel or delegate execution. /// Must outlive the Method that uses it. May be [`None`] if the Method does not use kernels or delegates that /// allocate temporary data. This allocator will be reset after every kernel or delegate call during execution. @@ -356,6 +439,28 @@ impl<'a> MemoryManager<'a> { PhantomData, ) } + + /// Returns per-buffer device metadata. One entry per planned memory buffer, same count as + /// `planned_memory` buffers. Empty if no device metadata was provided (CPU-only program) or if + /// `planned_memory` is [`None`]. + /// + /// This is a thin wrapper around + /// [`HierarchicalAllocator::planned_buffer_devices`](HierarchicalAllocator::planned_buffer_devices). + pub fn planned_buffer_devices(&self) -> &[Device] { + let devices = unsafe { sys::executorch_MemoryManager_planned_buffer_devices(self.0.get()) }; + if devices.len == 0 { + return &[]; // guard against null pointer + } + // Safety: the span points into the HierarchicalAllocator this manager borrows, which + // outlives the returned borrow. + unsafe { std::slice::from_raw_parts(devices.data.cast::(), devices.len) } + } + + /// Returns true if any planned buffer has device metadata attached. + /// When false, the memory setup is CPU-only. + pub fn has_device_memory(&self) -> bool { + unsafe { sys::executorch_MemoryManager_has_device_memory(self.0.get()) } + } } /// Storage for a non-trivially movable Cpp object. @@ -550,6 +655,34 @@ mod tests { test_memory_allocator(allocator_init, true); } + #[test] + fn buffer_memory_allocator_size_accounting() { + let mut buffer: [u8; 1024] = [0; 1024]; + let allocator = BufferMemoryAllocator::new(&mut buffer); + assert_eq!(allocator.used_size(), 0); + assert_eq!(allocator.free_size(), 1024); + + allocator.allocate_raw(100, 1).unwrap(); + assert_eq!(allocator.used_size(), 100); + assert_eq!(allocator.free_size(), 1024 - 100); + + // Alignment padding counts as used. + allocator.allocate_raw(8, 64).unwrap(); + assert!(allocator.used_size() >= 108); + assert_eq!(allocator.used_size() + allocator.free_size(), 1024); + } + + #[cfg(feature = "std")] + #[test] + fn malloc_memory_allocator_size_accounting_is_zero() { + // MallocMemoryAllocator constructs its base as MemoryAllocator(0, nullptr) and does not + // override used_size/free_size, so both are always 0 regardless of what it hands out. + let allocator = MallocMemoryAllocator::new(); + allocator.allocate_raw(64, 8).unwrap(); + assert_eq!(allocator.used_size(), 0); + assert_eq!(allocator.free_size(), 0); + } + #[cfg(feature = "std")] #[test] fn malloc_memory_allocator() { @@ -624,6 +757,74 @@ mod tests { } } + #[test] + fn hierarchical_allocator_no_devices() { + let mut buf0 = [0_u8; 64]; + let mut buf1 = [0_u8; 32]; + let mut spans = [Span::from_slice(&mut buf0), Span::from_slice(&mut buf1)]; + let allocator = HierarchicalAllocator::new(&mut spans); + assert!(allocator.planned_buffer_devices().is_empty()); + } + + #[test] + fn hierarchical_allocator_with_devices() { + use crate::device::{Device, DeviceType}; + + let mut buf0 = [0_u8; 64]; + let mut buf1 = [0_u8; 32]; + let mut spans = [Span::from_slice(&mut buf0), Span::from_slice(&mut buf1)]; + let devices = [ + Device::new(DeviceType::Cpu, 0), + Device::new(DeviceType::Cuda, 1), + ]; + // Safety: the test never reads through the buffers, so the deliberately-wrong Cuda tag on + // buf1 is never dereferenced. + let allocator = unsafe { HierarchicalAllocator::new_with_devices(&mut spans, &devices) }; + assert_eq!(allocator.planned_buffer_devices(), &devices[..]); + assert_eq!(allocator.planned_buffer_devices()[1].index(), 1); + assert_eq!( + allocator.planned_buffer_devices()[1].type_(), + DeviceType::Cuda + ); + } + + #[test] + fn memory_manager_no_planned_memory() { + let mut method_buf = [0_u8; 1024]; + let method_allocator = BufferMemoryAllocator::new(&mut method_buf); + let manager = MemoryManager::new(&method_allocator, None, None); + assert!(manager.planned_buffer_devices().is_empty()); + assert!(!manager.has_device_memory()); + } + + #[test] + fn memory_manager_planned_memory_without_devices() { + let mut method_buf = [0_u8; 1024]; + let method_allocator = BufferMemoryAllocator::new(&mut method_buf); + let mut buf0 = [0_u8; 64]; + let mut spans = [Span::from_slice(&mut buf0)]; + let mut planned = HierarchicalAllocator::new(&mut spans); + let manager = MemoryManager::new(&method_allocator, Some(&mut planned), None); + assert!(manager.planned_buffer_devices().is_empty()); + assert!(!manager.has_device_memory()); + } + + #[test] + fn memory_manager_planned_memory_with_devices() { + use crate::device::{Device, DeviceType}; + + let mut method_buf = [0_u8; 1024]; + let method_allocator = BufferMemoryAllocator::new(&mut method_buf); + let mut buf0 = [0_u8; 64]; + let mut spans = [Span::from_slice(&mut buf0)]; + let devices = [Device::new(DeviceType::Cuda, 3)]; + // Safety: the test never reads through the buffer, so the Cuda tag is never dereferenced. + let mut planned = unsafe { HierarchicalAllocator::new_with_devices(&mut spans, &devices) }; + let manager = MemoryManager::new(&method_allocator, Some(&mut planned), None); + assert!(manager.has_device_memory()); + assert_eq!(manager.planned_buffer_devices(), &devices[..]); + } + #[test] fn storage_macro() { let _: std::pin::Pin<&mut super::Storage> = storage!(i32); diff --git a/executorch/src/tensor/ptr.rs b/executorch/src/tensor/ptr.rs index d2acb218..d5eacf4c 100644 --- a/executorch/src/tensor/ptr.rs +++ b/executorch/src/tensor/ptr.rs @@ -9,6 +9,7 @@ use super::{ Data, DataMut, DataTyped, DimOrderType, Scalar, SizesType, StridesType, TensorBase, View, ViewMut, }; +use crate::device::Device; use crate::util::{IntoCpp, IntoRust}; use crate::{Error, Result}; @@ -50,6 +51,9 @@ impl TensorPtr<'static, View> { /// /// To create a mutable tensor from an array, use [`TensorPtrBuilder`]. /// + /// The result is always a CPU tensor. To move it to a device, use + /// [`clone_to`](TensorPtr::clone_to). + /// /// # Errors /// /// Returns an error if the array is not dense, i.e. if the strides are not the standard layout strides of some @@ -62,6 +66,9 @@ impl TensorPtr<'static, View> { /// Create a one dimensional [`TensorPtr`] from a vector. /// /// To create a mutable tensor from a vector, use [`TensorPtrBuilder`]. + /// + /// The result is always a CPU tensor. To move it to a device, use + /// [`clone_to`](TensorPtr::clone_to). pub fn from_vec(vec: Vec) -> Self { TensorPtrBuilder::>::from_vec(vec).build().unwrap() } @@ -69,6 +76,14 @@ impl TensorPtr<'static, View> { /// Create a new [`TensorPtr`] with data copied from a tensor. /// /// If the given tensor has a data type different from the target tensor dtype (`S`), the data will be casted. + /// + /// The result is always a CPU tensor. To move it to a device, use + /// [`clone_to`](TensorPtr::clone_to). + /// + /// # Panics + /// + /// The underlying Cpp function aborts the process if the source tensor is not a CPU tensor. + /// Use [`clone_to`](TensorPtr::clone_to) to move a device tensor back to the host first. pub fn copy_of(tensor: &TensorBase<'_, D>) -> Self { let tensor = unsafe { tensor.as_cpp().ptr.cast::().as_ref().unwrap() }; TensorPtr(sys::TensorPtr_clone(tensor, S::TYPE.cpp()), PhantomData) @@ -79,6 +94,9 @@ impl<'a, S: Scalar> TensorPtr<'a, View> { /// /// To create a mutable tensor from an array view, use [`TensorPtrBuilder`]. /// + /// The result is always a CPU tensor. To move it to a device, use + /// [`clone_to`](TensorPtr::clone_to). + /// /// # Errors /// /// Returns an error if the array is not dense, i.e. if the strides are not the standard layout strides of some @@ -93,6 +111,9 @@ impl<'a, S: Scalar> TensorPtr<'a, View> { /// Create a one dimensional [`TensorPtr`] from a slice. /// /// To create a mutable tensor from a slice, use [`TensorPtrBuilder`]. + /// + /// The result is always a CPU tensor. To move it to a device, use + /// [`clone_to`](TensorPtr::clone_to). pub fn from_slice(data: &'a [S]) -> Self { TensorPtrBuilder::>::from_slice(data) .build() @@ -125,6 +146,32 @@ impl TensorPtr<'_, D> { // Safety: the tensor is mutable, and we are the sole borrower. unsafe { TensorBase::from_inner_ref_mut(tensor) } } + + /// Clones a TensorPtr's data onto the given target device, allocating and copying as needed. + /// + /// The transfer direction is inferred from the source and target device: host-to-device when + /// `target` is an accelerator, and device-to-host when `target` is CPU. Copies use the + /// DeviceAllocator registered for the accelerator side; a device-backed result owns its memory + /// and frees it via that allocator when destroyed. + /// + /// Source and target must differ in device domain: for a CPU-to-CPU copy use + /// [`copy_of`](TensorPtr::copy_of), and device-to-device transfers are not supported. + /// + /// # Safety + /// + /// With a non-CPU `target` the returned tensor points at device memory. The rest of this + /// crate's tensor API — [`as_data_ptr`](TensorBase::as_data_ptr), the `as_array*` conversions, + /// the tensor accessors — reads the buffer directly from the host and is only valid for CPU + /// tensors, so the caller must not use those accessors on such a result. + pub unsafe fn clone_to(&self, target: Device) -> TensorPtr<'static, View> + where + D: DataTyped, + { + TensorPtr( + sys::TensorPtr_clone_to(self.0.clone(), target.cpp()), + PhantomData, + ) + } } unsafe impl Send for TensorPtr<'_, D> {} @@ -134,6 +181,7 @@ pub struct TensorPtrBuilder<'a, D: DataTyped> { data: TensorPtrBuilderData<'a, D>, strides: Option>>, dynamism: sys::ET_TensorShapeDynamism, + device: Device, } enum TensorPtrBuilderData<'a, D: DataTyped> { Vec { data: Vec, offset: usize }, @@ -174,6 +222,7 @@ impl TensorPtrBuilder<'static, D> { } }, dynamism: sys::ET_TensorShapeDynamism::ET_TensorShapeDynamism_STATIC, + device: Device::default(), } } @@ -194,6 +243,7 @@ impl TensorPtrBuilder<'static, D> { data: TensorPtrBuilderData::Vec { data, offset: 0 }, strides: None, dynamism: sys::ET_TensorShapeDynamism::ET_TensorShapeDynamism_STATIC, + device: Device::default(), } } } @@ -212,6 +262,7 @@ impl<'a, S: Scalar> TensorPtrBuilder<'a, View> { .map(|&s| s as StridesType), )), dynamism: sys::ET_TensorShapeDynamism::ET_TensorShapeDynamism_STATIC, + device: Device::default(), } } @@ -225,6 +276,7 @@ impl<'a, S: Scalar> TensorPtrBuilder<'a, View> { data: TensorPtrBuilderData::Slice(data), strides: None, dynamism: sys::ET_TensorShapeDynamism::ET_TensorShapeDynamism_STATIC, + device: Device::default(), } } @@ -247,6 +299,7 @@ impl<'a, S: Scalar> TensorPtrBuilder<'a, View> { strides: None, sizes: cxx_vec(sizes), dynamism: sys::ET_TensorShapeDynamism::ET_TensorShapeDynamism_STATIC, + device: Device::default(), } } } @@ -268,6 +321,7 @@ impl<'a, S: Scalar> TensorPtrBuilder<'a, ViewMut> { .map(|&s| s as StridesType), )), dynamism: sys::ET_TensorShapeDynamism::ET_TensorShapeDynamism_STATIC, + device: Device::default(), } } @@ -281,6 +335,7 @@ impl<'a, S: Scalar> TensorPtrBuilder<'a, ViewMut> { data: TensorPtrBuilderData::SliceMut(data), strides: None, dynamism: sys::ET_TensorShapeDynamism::ET_TensorShapeDynamism_STATIC, + device: Device::default(), } } @@ -303,6 +358,7 @@ impl<'a, S: Scalar> TensorPtrBuilder<'a, ViewMut> { strides: None, sizes: cxx_vec(sizes), dynamism: sys::ET_TensorShapeDynamism::ET_TensorShapeDynamism_STATIC, + device: Device::default(), } } } @@ -331,6 +387,23 @@ impl<'a, D: DataTyped> TensorPtrBuilder<'a, D> { self } + /// Set the device on which the tensor's data resides. Defaults to CPU. + /// + /// This sets the Tensor's device location only — no data is allocated or copied. The caller is + /// responsible for ensuring the data buffer already lives on the requested device. To copy CPU + /// data to a device, use [`TensorPtr::clone_to`] instead. + /// + /// # Safety + /// + /// The caller must ensure the data buffer actually resides on `device`. The rest of this + /// crate's tensor API — [`as_data_ptr`](TensorBase::as_data_ptr), the `as_array*` conversions, + /// the tensor accessors — reads the buffer directly from the host and is only valid for CPU + /// tensors, so tagging a tensor with a non-CPU device makes those accessors unsound to call. + pub unsafe fn device(mut self, device: Device) -> Self { + self.device = device; + self + } + /// Build an immutable tensor. /// /// # Errors @@ -397,6 +470,7 @@ impl<'a, D: DataTyped> TensorPtrBuilder<'a, D> { D::Scalar::TYPE.cpp(), self.dynamism, Box::new(sys::util::RustAny::new(Box::new(allocation_vec))), + self.device.cpp(), ) }; Ok(TensorPtr(tensor, PhantomData)) @@ -469,6 +543,7 @@ impl<'a, D: DataTyped> TensorPtrBuilder<'a, D> { D::Scalar::TYPE.cpp(), self.dynamism, Box::new(sys::util::RustAny::new(Box::new(allocation_vec))), + self.device.cpp(), ) }; Ok(TensorPtr(tensor, PhantomData)) @@ -516,6 +591,36 @@ macro_rules! tensor_ptr { mod tests { use super::*; + #[test] + fn builder_device_defaults_to_cpu() { + use crate::device::DeviceType; + + let tensor = TensorPtrBuilder::>::from_vec(vec![1.0, 2.0, 3.0]) + .build() + .unwrap(); + let device = tensor.as_tensor().device(); + assert_eq!(device.type_(), DeviceType::Cpu); + assert_eq!(device.index(), 0); + assert!(device.is_cpu()); + } + + #[test] + fn builder_device_explicit_cpu() { + use crate::device::{Device, DeviceType}; + + let builder = TensorPtrBuilder::>::from_vec(vec![1.0, 2.0, 3.0]); + // Safety: the data is a Rust Vec, which lives in host memory, matching the CPU device. + let builder = unsafe { builder.device(Device::new(DeviceType::Cpu, 0)) }; + let tensor = builder.build().unwrap(); + assert_eq!(tensor.as_tensor().device(), Device::new(DeviceType::Cpu, 0)); + // Tagging the tensor CPU explicitly must not disturb the data. + #[cfg(feature = "ndarray")] + assert_eq!( + tensor.as_tensor().as_array::().as_slice(), + Some(&[1.0_f32, 2.0, 3.0][..]) + ); + } + #[cfg(feature = "ndarray")] #[test] fn fron_array() {