diff --git a/kernels/portable/cpu/op__device_copy.cpp b/kernels/portable/cpu/op__device_copy.cpp index 01fadd084ef..b636dbbb489 100644 --- a/kernels/portable/cpu/op__device_copy.cpp +++ b/kernels/portable/cpu/op__device_copy.cpp @@ -9,16 +9,24 @@ /** * Runtime kernels for et_copy._h2d_copy and et_copy._d2h_copy ops. * - * These ops transfer tensor data between CPU and device memory using - * the DeviceAllocator interface. The device type is inferred from the - * tensor metadata (out.device_type() for H2D, self.device_type() for D2H), - * which was set during AOT serialization by PropagateDevicePass. + * These ops transfer tensor data between CPU and device memory using the + * DeviceAllocator interface. + * + * In portable (non-ATen) mode the device type/index are read from the tensor + * metadata (out for H2D, self for D2H), which PropagateDevicePass set during + * AOT serialization. In ATen mode the planned at::Tensor does not carry that + * metadata, so the direction comes from the op identity (_h2d vs _d2h) and the + * copy targets the registered CUDA DeviceAllocator on the current device. */ #include #include #include +#ifdef USE_ATEN_LIB +#include // TORCH_CHECK (catchable, throws std::exception) +#endif + namespace torch { namespace executor { namespace native { @@ -26,6 +34,169 @@ namespace native { using Tensor = executorch::aten::Tensor; using DeviceAllocator = executorch::runtime::DeviceAllocator; using Error = executorch::runtime::Error; +using RuntimeDeviceIndex = executorch::runtime::etensor::DeviceIndex; +using RuntimeDeviceType = executorch::runtime::etensor::DeviceType; + +namespace { + +#ifdef USE_ATEN_LIB +// In ATen mode the memory-planned `at::Tensor` does not carry the runtime's +// planned device metadata (it is constructed CPU-side by the ATen tensor +// parser), so the copy direction is taken from the op identity (_h2d vs _d2h) +// rather than from tensor device metadata. +// +// The non-CPU device is CUDA: the DeviceAllocator registry is keyed by +// DeviceType with no "the one non-CPU allocator" lookup, and CUDA is the only +// non-CPU accelerator wired up for ATen-mode device copies today. If another +// ATen-mode accelerator is added, this (and the registry lookup) must learn to +// resolve its type. +// +// `index == -1` requests the allocator's current device. This is a convention +// of the CUDA allocator implementation (see CudaAllocator::copy_* / +// CudaAllocator::allocate in backends/cuda/runtime/cuda_allocator.cpp), not a +// documented guarantee on the DeviceAllocator interface. It assumes the planned +// buffer lives on the current device, i.e. a single-GPU ATen-mode setup; +// multi-GPU ATen mode (planned buffer not on the current device) is not handled +// because the planned device index is not recoverable from the at::Tensor here. +// TODO: support multi-GPU / non-CUDA ATen-mode device copy by recovering the +// planned device type/index (e.g. via the tensor_parser_aten path) instead of +// assuming the current CUDA device. +// +// Unlike the non-ATen branch below, these kernels cannot re-check tensor device +// placement (h2d source-must-be-CPU / dest-must-be-non-CPU, etc.): that +// metadata is not on the at::Tensor. They trust the graph wiring, which the +// device-placement pass controls when it inserts these ops. +constexpr RuntimeDeviceType kAtenDeviceType = RuntimeDeviceType::CUDA; +constexpr RuntimeDeviceIndex kCurrentDeviceIndex = -1; + +DeviceAllocator* require_device_allocator( + RuntimeDeviceType device_type, + const char* op_name) { + DeviceAllocator* allocator = + executorch::runtime::get_device_allocator(device_type); + TORCH_CHECK( + allocator != nullptr, + op_name, + ": no device allocator registered for device_type=", + static_cast(device_type)); + return allocator; +} +#else +RuntimeDeviceType runtime_device_type(const Tensor& tensor) { + return tensor.unsafeGetTensorImpl()->device_type(); +} + +RuntimeDeviceIndex runtime_device_index(const Tensor& tensor) { + return tensor.unsafeGetTensorImpl()->device_index(); +} +#endif // USE_ATEN_LIB + +} // namespace + +#ifdef USE_ATEN_LIB +namespace { + +// Bytes writable in `t` starting at its data_ptr(), accounting for a possible +// non-zero storage offset (a view may start partway into its storage). +size_t writable_nbytes(const Tensor& t) { + auto* impl = t.unsafeGetTensorImpl(); + const size_t storage_bytes = impl->storage().nbytes(); + const size_t offset_bytes = + static_cast(impl->storage_offset()) * impl->itemsize(); + return offset_bytes >= storage_bytes ? 0 : storage_bytes - offset_bytes; +} + +// Shared body for both directions. In ATen mode the tensors carry no planned +// device metadata, so the direction comes from the caller (`to_device`) and the +// device is the registered CUDA allocator's current device. Validates every +// assumption a raw byte copy makes before touching the allocator: +// - `out` resizes to `self`'s shape (ATen resize_tensor is metadata-only and +// does not grow storage, so a too-small `out` would otherwise overrun), +// - `self` and `out` share dtype and are contiguous, +// - `out`'s offset-aware writable bytes can hold `self.nbytes()`. +// A zero-byte copy returns early (the allocator rejects null pointers, which +// empty tensors can have). Violations raise via TORCH_CHECK (a catchable +// exception derived from std::exception) rather than aborting: these kernels +// run inside a libtorch process in ATen mode, so a bad boundary tensor or a +// transient copy fault should be recoverable by the host rather than crash the +// process. The ATen custom-op ABI binds the contextless overload, so there is +// no KernelRuntimeContext to return a portable-style Error through. +Tensor& device_copy(const Tensor& self, Tensor& out, bool to_device) { + const char* op_name = to_device ? "_h2d_copy" : "_d2h_copy"; + + // Check preconditions (dtype / contiguity / capacity) BEFORE resizing `out`, + // so a bad-argument failure leaves `out` untouched. Note the allocator copy + // below runs after resize_tensor(), so if the copy itself fails `out` may + // already be resized when we throw; that is a device/transport fault, not a + // caller-argument error, and the host is expected to treat the copy as failed + // rather than reuse `out`. + TORCH_CHECK( + self.scalar_type() == out.scalar_type(), + op_name, + ": self/out dtype mismatch (", + self.scalar_type(), + " vs ", + out.scalar_type(), + ")"); + TORCH_CHECK( + self.is_contiguous() && out.is_contiguous(), + op_name, + ": self and out must be contiguous"); + const size_t nbytes = self.nbytes(); + const size_t out_writable = writable_nbytes(out); + TORCH_CHECK( + nbytes <= out_writable, + op_name, + ": out too small (self.nbytes()=", + nbytes, + ", out writable=", + out_writable, + ")"); + TORCH_CHECK( + resize_tensor(out, self.sizes()) == Error::Ok, + op_name, + ": cannot resize out to self sizes (self.nbytes()=", + self.nbytes(), + ")"); + + if (nbytes == 0) { + return out; + } + + DeviceAllocator* allocator = + require_device_allocator(kAtenDeviceType, op_name); + const Error err = to_device ? allocator->copy_host_to_device( + out.mutable_data_ptr(), + self.const_data_ptr(), + nbytes, + kCurrentDeviceIndex) + : allocator->copy_device_to_host( + out.mutable_data_ptr(), + self.const_data_ptr(), + nbytes, + kCurrentDeviceIndex); + TORCH_CHECK( + err == Error::Ok, + op_name, + ": device copy failed (", + nbytes, + " bytes, ", + to_device ? "host->device" : "device->host", + ")"); + return out; +} + +} // namespace + +Tensor& _h2d_copy_out(const Tensor& self, Tensor& out) { + return device_copy(self, out, /*to_device=*/true); +} + +Tensor& _d2h_copy_out(const Tensor& self, Tensor& out) { + return device_copy(self, out, /*to_device=*/false); +} + +#else /** * Copies tensor data from host (CPU) memory to device memory. @@ -37,21 +208,20 @@ using Error = executorch::runtime::Error; */ Tensor& _h2d_copy_out(KernelRuntimeContext& ctx, const Tensor& self, Tensor& out) { - auto device_type = out.unsafeGetTensorImpl()->device_type(); - auto device_index = out.unsafeGetTensorImpl()->device_index(); + auto device_type = runtime_device_type(out); + auto device_index = runtime_device_index(out); ET_KERNEL_CHECK_MSG( ctx, - self.unsafeGetTensorImpl()->device_type() == - executorch::runtime::etensor::DeviceType::CPU, + runtime_device_type(self) == RuntimeDeviceType::CPU, InvalidArgument, out, "_h2d_copy: source tensor must be on CPU, got device_type=%d", - static_cast(self.unsafeGetTensorImpl()->device_type())); + static_cast(runtime_device_type(self))); ET_KERNEL_CHECK_MSG( ctx, - device_type != executorch::runtime::etensor::DeviceType::CPU, + device_type != RuntimeDeviceType::CPU, InvalidArgument, out, "_h2d_copy: destination tensor must be on a non-CPU device"); @@ -98,24 +268,23 @@ _h2d_copy_out(KernelRuntimeContext& ctx, const Tensor& self, Tensor& out) { */ Tensor& _d2h_copy_out(KernelRuntimeContext& ctx, const Tensor& self, Tensor& out) { - auto device_type = self.unsafeGetTensorImpl()->device_type(); - auto device_index = self.unsafeGetTensorImpl()->device_index(); + auto device_type = runtime_device_type(self); + auto device_index = runtime_device_index(self); ET_KERNEL_CHECK_MSG( ctx, - device_type != executorch::runtime::etensor::DeviceType::CPU, + device_type != RuntimeDeviceType::CPU, InvalidArgument, out, "_d2h_copy: source tensor must be on a non-CPU device"); ET_KERNEL_CHECK_MSG( ctx, - out.unsafeGetTensorImpl()->device_type() == - executorch::runtime::etensor::DeviceType::CPU, + runtime_device_type(out) == RuntimeDeviceType::CPU, InvalidArgument, out, "_d2h_copy: destination tensor must be on CPU, got device_type=%d", - static_cast(out.unsafeGetTensorImpl()->device_type())); + static_cast(runtime_device_type(out))); ET_KERNEL_CHECK_MSG( ctx, @@ -149,6 +318,8 @@ _d2h_copy_out(KernelRuntimeContext& ctx, const Tensor& self, Tensor& out) { return out; } +#endif // USE_ATEN_LIB + } // namespace native } // namespace executor } // namespace torch diff --git a/kernels/portable/et_copy_ops.yaml b/kernels/portable/et_copy_ops.yaml new file mode 100644 index 00000000000..786aae91b22 --- /dev/null +++ b/kernels/portable/et_copy_ops.yaml @@ -0,0 +1,16 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This yaml file contains ExecuTorch device-copy ops for ATen-mode runtime +# registration. These ops are inserted by device partitioning and are not ATen +# native operators, so ATen-mode codegen must consume them through the custom-op +# path. + +- func: et_copy::_h2d_copy.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: torch::executor::_h2d_copy_out + +- func: et_copy::_d2h_copy.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: torch::executor::_d2h_copy_out diff --git a/kernels/portable/targets.bzl b/kernels/portable/targets.bzl index b80ce347768..952077e3f2c 100644 --- a/kernels/portable/targets.bzl +++ b/kernels/portable/targets.bzl @@ -37,6 +37,11 @@ def define_common_targets(): visibility = ["PUBLIC"], ) + runtime.export_file( + name = "et_copy_ops.yaml", + visibility = ["PUBLIC"], + ) + et_operator_library( name = "executorch_all_ops", include_all_operators = True, @@ -96,3 +101,43 @@ def define_common_targets(): visibility = ["PUBLIC"], define_static_targets = True, ) + + et_operator_library( + name = "et_copy_ops", + ops_schema_yaml_target = "//executorch/kernels/portable:et_copy_ops.yaml", + define_static_targets = True, + visibility = ["PUBLIC"], + ) + + # ATen-mode registration for the device-copy ops + # (`et_copy::_h2d_copy.out` / `_d2h_copy.out`) that the ExecuTorch + # device-placement pass inserts at CPU<->accelerator boundaries (e.g. + # around CUDA-delegated subgraphs). These ops live in `functions.yaml`, + # so `generated_lib` already registers them for the portable runtime. + # ATen-mode codegen, however, reads ATen's `native_functions.yaml` for + # `functions_yaml_target` entries and cannot register non-ATen ops that + # way, so `generated_lib_aten` does not cover them. This dedicated lib + # registers them for ATen-mode runtimes via the custom-ops path, reusing + # the same `op__device_copy` kernel (its `_aten` variant). `et_copy` is + # intentionally in its own `et_copy_ops.yaml` (not `custom_ops.yaml`) so + # it is not double-registered in the portable `generated_lib`. + # + # Note: because `op__device_copy` is in the shared CUSTOM_OPS list, its + # `_aten` kernel is also compiled into `operators_aten` and thus links + # (as unregistered code) into the general `generated_lib_aten`. This lib + # is what actually *registers* the ops for ATen-mode consumers; it does + # not isolate the kernel object from other ATen consumers. + executorch_generated_lib( + name = "device_copy_ops_aten_lib", + deps = [ + ":et_copy_ops", + ], + kernel_deps = [ + "//executorch/kernels/portable/cpu:op__device_copy_aten", + ], + custom_ops_yaml_target = "//executorch/kernels/portable:et_copy_ops.yaml", + custom_ops_requires_aot_registration = False, + aten_mode = True, + visibility = ["PUBLIC"], + define_static_targets = True, + ) diff --git a/kernels/test/op__device_copy_aten_test.cpp b/kernels/test/op__device_copy_aten_test.cpp new file mode 100644 index 00000000000..1c5ffeb3a12 --- /dev/null +++ b/kernels/test/op__device_copy_aten_test.cpp @@ -0,0 +1,175 @@ +/* + * 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. + */ + +/** + * ATen-mode tests for et_copy._h2d_copy.out / et_copy._d2h_copy.out. + * + * In ATen mode the memory-planned at::Tensor does not carry the runtime's + * planned device metadata, so the kernels take the copy direction from the op + * identity (_h2d vs _d2h) and route through the registered non-CPU + * DeviceAllocator. These tests use a MockCudaAllocator (host memory) so they + * exercise direction-from-op-identity, the current-device index (-1), and the + * resize-before-copy behavior without requiring a GPU. + * + * The portable-mode kernels are covered separately in op__device_copy_test.cpp, + * whose TensorImpl-based construction does not compile in ATen mode. + */ + +#include + +#include + +#include + +#include +#include +#include +#include + +using executorch::aten::Tensor; +using executorch::runtime::get_device_allocator; +using executorch::runtime::register_device_allocator; +using executorch::runtime::etensor::DeviceType; +using executorch::runtime::testing::MockCudaAllocator; + +// The ATen-mode custom-op kernels, as declared by the codegen `kernel_name` +// (torch::executor::_h2d_copy_out -> torch::executor::native::_h2d_copy_out). +// The 2-arg (contextless) overloads are the ones bound in ATen mode. +namespace torch { +namespace executor { +namespace native { +Tensor& _h2d_copy_out(const Tensor& self, Tensor& out); +Tensor& _d2h_copy_out(const Tensor& self, Tensor& out); +} // namespace native +} // namespace executor +} // namespace torch + +namespace { + +MockCudaAllocator& mock_cuda() { + static MockCudaAllocator instance; + return instance; +} + +class OpDeviceCopyAtenTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + executorch::runtime::runtime_init(); + if (get_device_allocator(DeviceType::CUDA) == nullptr) { + register_device_allocator(&mock_cuda()); + } + // The registry has no unregister/replace. If a real CUDA allocator was + // registered first (e.g. this test is linked into a binary that pulls in + // the CUDA backend, whose static init registers CudaAllocator::instance()), + // the mock is not installed and these tests cannot observe the mock + // counters. Skip rather than fail in that configuration; in this CPU-only + // target the mock is always the registered allocator. + if (get_device_allocator(DeviceType::CUDA) != &mock_cuda()) { + GTEST_SKIP() + << "a non-mock CUDA allocator is registered; these tests require " + << "MockCudaAllocator and the registry cannot be replaced"; + } + } + + void SetUp() override { + mock_cuda().reset(); + } +}; + +// H2D takes direction from the op identity: it must call copy_host_to_device +// (not d2h) and forward the current-device index (-1), regardless of the +// at::Tensor device metadata (which is CPU in ATen mode). +TEST_F(OpDeviceCopyAtenTest, H2dCopiesHostToDeviceByOpIdentity) { + const at::Tensor self = at::tensor({1.0f, 2.0f, 3.0f, 4.0f}); + at::Tensor out = at::zeros({4}); + + Tensor& result = torch::executor::native::_h2d_copy_out(self, out); + + EXPECT_EQ(mock_cuda().h2d_count_, 1); + EXPECT_EQ(mock_cuda().d2h_count_, 0); + EXPECT_EQ(mock_cuda().last_h2d_size_, 4 * sizeof(float)); + EXPECT_EQ(mock_cuda().last_h2d_index_, -1); + EXPECT_TRUE(at::equal(out, self)); + EXPECT_EQ(&result, &out); +} + +// D2H is the mirror: it must call copy_device_to_host and forward index -1. +TEST_F(OpDeviceCopyAtenTest, D2hCopiesDeviceToHostByOpIdentity) { + const at::Tensor self = at::tensor({5.0f, 6.0f, 7.0f, 8.0f}); + at::Tensor out = at::zeros({4}); + + Tensor& result = torch::executor::native::_d2h_copy_out(self, out); + + EXPECT_EQ(mock_cuda().d2h_count_, 1); + EXPECT_EQ(mock_cuda().h2d_count_, 0); + EXPECT_EQ(mock_cuda().last_d2h_size_, 4 * sizeof(float)); + EXPECT_EQ(mock_cuda().last_d2h_index_, -1); + EXPECT_TRUE(at::equal(out, self)); + EXPECT_EQ(&result, &out); +} + +// out is resized to self's shape before the copy: start out at a different +// shape so the resize is actually observable. +TEST_F(OpDeviceCopyAtenTest, H2dResizesOutToSelf) { + const at::Tensor self = at::tensor({1.0f, 2.0f, 3.0f}); + at::Tensor out = at::zeros({8}); + ASSERT_NE(out.sizes(), self.sizes()); + + torch::executor::native::_h2d_copy_out(self, out); + + EXPECT_EQ(out.sizes(), self.sizes()); + EXPECT_EQ(mock_cuda().h2d_count_, 1); + EXPECT_EQ(mock_cuda().last_h2d_size_, 3 * sizeof(float)); + EXPECT_TRUE(at::equal(out, self)); +} + +// If self is larger than out's backing storage, the kernel must reject the copy +// rather than overrun out's storage (ATen resize_tensor only updates sizes, it +// does not grow storage). It raises a catchable exception (not an abort) so a +// libtorch host can recover. +TEST_F(OpDeviceCopyAtenTest, H2dRejectsWhenSelfExceedsOutStorage) { + const at::Tensor self = at::tensor({1.0f, 2.0f, 3.0f, 4.0f}); + at::Tensor out = at::zeros({2}); // storage smaller than self + EXPECT_THROW( + torch::executor::native::_h2d_copy_out(self, out), std::exception); + EXPECT_EQ(mock_cuda().h2d_count_, 0); +} + +// Same overrun guard for the d2h direction. Both directions share device_copy, +// but assert d2h explicitly so the guard survives any future un-sharing. +TEST_F(OpDeviceCopyAtenTest, D2hRejectsWhenSelfExceedsOutStorage) { + const at::Tensor self = at::tensor({1.0f, 2.0f, 3.0f, 4.0f}); + at::Tensor out = at::zeros({2}); // storage smaller than self + EXPECT_THROW( + torch::executor::native::_d2h_copy_out(self, out), std::exception); + EXPECT_EQ(mock_cuda().d2h_count_, 0); +} + +// A dtype mismatch would make the raw byte copy meaningless, so the kernel must +// reject rather than copy. +TEST_F(OpDeviceCopyAtenTest, H2dRejectsDtypeMismatch) { + const at::Tensor self = at::tensor({1.0f, 2.0f}); + at::Tensor out = at::zeros({2}, at::kInt); + EXPECT_THROW( + torch::executor::native::_h2d_copy_out(self, out), std::exception); + EXPECT_EQ(mock_cuda().h2d_count_, 0); +} + +// A zero-size copy is a valid no-op: it must not touch the allocator (whose +// null-pointer checks would otherwise reject empty tensors). +TEST_F(OpDeviceCopyAtenTest, H2dZeroSizeIsNoOp) { + const at::Tensor self = at::zeros({0}); + at::Tensor out = at::zeros({0}); + + torch::executor::native::_h2d_copy_out(self, out); + + EXPECT_EQ(mock_cuda().h2d_count_, 0); + EXPECT_EQ(out.numel(), 0); +} + +} // namespace diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl index 431ec96b447..837c7327c4f 100644 --- a/kernels/test/targets.bzl +++ b/kernels/test/targets.bzl @@ -186,6 +186,22 @@ def define_common_targets(): "//executorch/runtime/platform:platform", ], ) + + # ATen-mode kernels for et_copy use at::Tensor construction, which the + # portable TensorImpl-based op__device_copy_test above cannot express, so + # the ATen path has its own dedicated test. + runtime.cxx_test( + name = "op__device_copy_aten_test", + srcs = ["op__device_copy_aten_test.cpp"], + preprocessor_flags = ["-DUSE_ATEN_LIB"], + deps = [ + "//executorch/kernels/portable/cpu:op__device_copy_aten", + "//executorch/runtime/core:device_allocator", + "//executorch/runtime/core/exec_aten:lib_aten", + "//executorch/runtime/core/test:mock_cuda_allocator", + "//executorch/runtime/platform:platform", + ], + ) _common_op_test("op_abs_test", ["aten", "portable"]) _common_op_test("op_acos_test", ["aten", "portable"]) _common_op_test("op_acosh_test", ["aten", "portable"]) diff --git a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl index 479f3913f8f..65016cd1ca1 100644 --- a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl @@ -1405,19 +1405,21 @@ ATEN_OPS = ( "//executorch/kernels/portable/cpu/util:copy_ops_util", ], ), - op_target( - name = "op__device_copy", - deps = [ - "//executorch/runtime/core:device_allocator", - ], - ), ) -# Operators that are not listed in `functions.yaml` (i.e., operators listed in -# `custom_ops.yaml`), which are not compatible with the core ATen operators. -# Every entry here will be backed by a cxx_library target with the given name -# and deps, as well as a similar `_aten` target that uses at::Tensor and -# related types. +# Operators that need a `_aten` target (using at::Tensor and related +# types) in addition to the lean `` target. Every entry here is backed by +# a cxx_library target with the given name and deps, plus a similar +# `_aten` target. +# +# Most entries are custom ops listed only in `custom_ops.yaml` (not in +# `functions.yaml`) because they are not core ATen operators. `op__device_copy` +# is an exception: it stays registered for the portable runtime via +# `functions.yaml` (`et_copy::_h2d_copy.out` / `_d2h_copy.out`), but is listed +# here so that an `op__device_copy_aten` target also exists. That `_aten` target +# backs `:device_copy_ops_aten_lib`, which registers these device-copy ops for +# ATen-mode runtimes (ATen-mode codegen consumes `custom_ops.yaml`-style +# schemas, not `functions.yaml`, so `generated_lib_aten` cannot register them). # # Note that a single target (or single .cpp file) can't mix ATen and non-ATen # ops, and must be split. They can, however, share common code via a library dep @@ -1426,6 +1428,12 @@ CUSTOM_OPS = ( op_target( name = "op_allclose", ), + op_target( + name = "op__device_copy", + deps = [ + "//executorch/runtime/core:device_allocator", + ], + ), ) def portable_source_list():