From 7ac0f4d3cad340a506d0d0755ab4dcbdf094e4bf Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 16:44:27 -0700 Subject: [PATCH 1/4] Add DeviceAllocator::synchronize() to wait for device copies The next commit copies the initial data of a tensor planned in device memory there when a method loads, through the device allocator. On CUDA that copy can still be running when copy_host_to_device() returns: it is asynchronous on a caller stream, and from pageable memory even a plain cudaMemcpy may return before the data lands. Two things can then go wrong. The host buffer the data came from may be freed while the copy still reads it. And the first run reads the tensor on its own stream; if that stream is not ordered after the copy, for example when the method loads with one caller stream and runs with another, or runs on a non-blocking stream such as a green context stream, it can read the tensor before its data is there. DeviceAllocator gets synchronize(), which waits for the copies that the allocator made on the calling thread, with the same per-thread settings. The default does nothing, for allocators whose copies finish before they return. CudaAllocator switches to the device and waits for the stream its copies use: the caller stream when one is set, otherwise the default stream of the device. It switches even with a caller stream, because a handle such as nullptr or cudaStreamPerThread names a stream of the current device. Without a caller stream the copy runs on the legacy default stream, so waiting for it can also wait for work that other threads queued on blocking streams of that device. test_cuda_allocator checks that synchronize() waits for the caller stream and for the default stream, and, with two GPUs, for the default stream of device 0 while device 1 is current and the caller stream is nullptr. --- .../cuda/runtime/test/test_cuda_allocator.cpp | 54 +++++++++++++++++++ extension/cuda/cuda_allocator.cpp | 32 +++++++++++ extension/cuda/cuda_allocator.h | 7 +++ runtime/core/device_allocator.h | 17 ++++++ runtime/core/test/mock_cuda_allocator.h | 11 ++++ 5 files changed, 121 insertions(+) diff --git a/backends/cuda/runtime/test/test_cuda_allocator.cpp b/backends/cuda/runtime/test/test_cuda_allocator.cpp index d31ba65e1a4..70a26ec2c5b 100644 --- a/backends/cuda/runtime/test/test_cuda_allocator.cpp +++ b/backends/cuda/runtime/test/test_cuda_allocator.cpp @@ -10,6 +10,8 @@ #include +#include +#include #include #include #include @@ -178,6 +180,58 @@ TEST_F(CudaAllocatorTest, CopyRoundtripWithCallerStream) { ASSERT_EQ(cudaStreamDestroy(s), cudaSuccess); } +// The compatibility header has no HIP names for non-blocking streams or host +// callbacks. +#if !defined(EXECUTORCH_USE_HIP) +// A method can be loaded on one stream and run on another that is not ordered +// after it, so synchronize() has to wait for the stream the copies used: the +// caller stream when one is set, otherwise the default stream. +TEST_F(CudaAllocatorTest, SynchronizeWaitsForTheCopyStream) { + const cudaHostFn_t hold = [](void* flag) { + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + static_cast*>(flag)->store(true); + }; + cudaStream_t caller; + ASSERT_EQ( + cudaStreamCreateWithFlags(&caller, cudaStreamNonBlocking), cudaSuccess); + for (cudaStream_t s : {caller, cudaStream_t{nullptr}}) { + std::optional g; + if (s != nullptr) { + g.emplace(s); + } + std::atomic done{false}; + ASSERT_EQ(cudaLaunchHostFunc(s, hold, &done), cudaSuccess); + EXPECT_EQ(CudaAllocator::instance().synchronize(0), Error::Ok); + EXPECT_TRUE(done.load()) + << (s != nullptr ? "caller" : "default") << " stream"; + // If synchronize() did not wait, the callback must still finish before + // `done` goes out of scope. + ASSERT_EQ(cudaStreamSynchronize(s), cudaSuccess); + } + ASSERT_EQ(cudaStreamDestroy(caller), cudaSuccess); + + // A caller stream of nullptr names the default stream of the current + // device, so with another device current, synchronize(0) still has to wait + // for device 0. + int device_count = 0; + ASSERT_EQ(cudaGetDeviceCount(&device_count), cudaSuccess); + if (device_count < 2) { + return; + } + std::atomic done{false}; + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + ASSERT_EQ(cudaLaunchHostFunc(nullptr, hold, &done), cudaSuccess); + ASSERT_EQ(cudaSetDevice(1), cudaSuccess); + { + executorch::extension::cuda::CallerStreamGuard g(nullptr); + EXPECT_EQ(CudaAllocator::instance().synchronize(0), Error::Ok); + } + EXPECT_TRUE(done.load()) << "default stream of device 0, device 1 current"; + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(nullptr), cudaSuccess); +} +#endif // !EXECUTORCH_USE_HIP + TEST_F(CudaAllocatorTest, CopyHostToDeviceNullDstReturnsInvalidArgument) { CudaAllocator& a = CudaAllocator::instance(); // null dst should fail gracefully not CHECK abort diff --git a/extension/cuda/cuda_allocator.cpp b/extension/cuda/cuda_allocator.cpp index b137bff0a45..c3837a703ae 100644 --- a/extension/cuda/cuda_allocator.cpp +++ b/extension/cuda/cuda_allocator.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -569,6 +570,37 @@ DeviceType CudaAllocator::device_type() const { return DeviceType::CUDA; } +Error CudaAllocator::synchronize(DeviceIndex index) { + ET_CHECK_OR_RETURN_ERROR( + index >= -1, + InvalidArgument, + "CudaAllocator::synchronize: invalid device index %d (must be >= -1)", + static_cast(index)); + // copy_impl switches to the device and copies on the caller stream when one + // is set, and otherwise with cudaMemcpy, on the default stream. Switch here + // too, even with a caller stream: a handle such as nullptr or + // cudaStreamPerThread names a stream of the current device. + const auto caller_stream = getCallerStream(); + std::optional guard; + if (index >= 0) { + Result selected = CUDAGuard::create(index); + ET_CHECK_OK_OR_RETURN_ERROR(selected.error()); + guard.emplace(std::move(selected.get())); + } + const cudaError_t err = + cudaStreamSynchronize(caller_stream.value_or(nullptr)); + if (err != cudaSuccess) { + ET_LOG( + Error, + "CudaAllocator::synchronize: cudaStreamSynchronize failed: %s " + "(device %d)", + cudaGetErrorString(err), + static_cast(index)); + return Error::Internal; + } + return Error::Ok; +} + CudaAllocator& CudaAllocator::instance() { // Registered allocators can be called until the process exits. static auto* allocator = new CudaAllocator(); diff --git a/extension/cuda/cuda_allocator.h b/extension/cuda/cuda_allocator.h index 72e72da68d9..2c55785076a 100644 --- a/extension/cuda/cuda_allocator.h +++ b/extension/cuda/cuda_allocator.h @@ -57,6 +57,13 @@ class EXECUTORCH_EXTENSION_CUDA_API CudaAllocator final executorch::runtime::etensor::DeviceType device_type() const override; + /** + * Waits for the stream that copies use: the caller stream when one is set, + * otherwise the default stream of the device. + */ + executorch::runtime::Error synchronize( + executorch::runtime::etensor::DeviceIndex index) override; + /// Returns the global CudaAllocator singleton. static CudaAllocator& instance(); diff --git a/runtime/core/device_allocator.h b/runtime/core/device_allocator.h index 259bdb9f032..4c5c56153e4 100644 --- a/runtime/core/device_allocator.h +++ b/runtime/core/device_allocator.h @@ -95,6 +95,23 @@ class DeviceAllocator { * Returns the device type this allocator handles. */ virtual etensor::DeviceType device_type() const = 0; + + /** + * Waits until the copies to device `index` that this allocator made on the + * calling thread have finished, so that work on any stream of the device can + * read their results. Call it with the same per-thread settings the copies + * were made with: an allocator that copies on a stream chosen per thread + * waits for the stream chosen when this is called. + * The default does nothing, which suits allocators whose copies finish + * before they return. + * + * @param index The device index. + * @return Error::Ok on success, or an appropriate error code on failure. + */ + virtual Error synchronize(etensor::DeviceIndex index) { + (void)index; + return Error::Ok; + } }; /** diff --git a/runtime/core/test/mock_cuda_allocator.h b/runtime/core/test/mock_cuda_allocator.h index b9639d1021e..5d07624ef74 100644 --- a/runtime/core/test/mock_cuda_allocator.h +++ b/runtime/core/test/mock_cuda_allocator.h @@ -87,6 +87,12 @@ class MockCudaAllocator : public DeviceAllocator { return etensor::DeviceType::CUDA; } + Error synchronize(etensor::DeviceIndex index) override { + synced_h2d_count_ = h2d_count_; + last_sync_index_ = index; + return Error::Ok; + } + /** * Returns true if ptr falls within the most recent allocation range. * Useful for verifying that tensor data_ptrs point to device memory. @@ -115,6 +121,8 @@ class MockCudaAllocator : public DeviceAllocator { last_h2d_src_ = nullptr; last_h2d_size_ = 0; last_h2d_index_ = -1; + synced_h2d_count_ = 0; + last_sync_index_ = -1; last_d2h_dst_ = nullptr; last_d2h_src_ = nullptr; last_d2h_size_ = 0; @@ -140,6 +148,9 @@ class MockCudaAllocator : public DeviceAllocator { const void* last_h2d_src_ = nullptr; size_t last_h2d_size_ = 0; etensor::DeviceIndex last_h2d_index_ = -1; + /// The value of h2d_count_ when synchronize() last ran. + int synced_h2d_count_ = 0; + etensor::DeviceIndex last_sync_index_ = -1; // Device-to-host copy tracking int d2h_count_ = 0; From 17cde2355d9c9f88418762905869ee26615e7a0a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 25 Oct 2026 17:25:02 -0700 Subject: [PATCH 2/4] Load the initial data of a planned tensor into device memory A memory-planned tensor can have initial data, such as a mutable buffer with a starting value. The runtime loads that data into the tensor's planned memory when it loads the method, with the data loader or the named data map. Both write with the CPU, so they cannot fill a planned buffer that lives in device memory, such as CUDA memory. The tensor parser now checks the device of the planned buffer, which the memory manager's planned memory records. For a device buffer it loads the data into host memory and copies it with the device allocator registered for that device, once, at method load. It then waits for the copy with the allocator's synchronize(), so the host buffer can be freed and the first run reads the right data. A host buffer is loaded as before, and the copies made on every call are unchanged. An empty tensor now loads nothing. When all the initial data of a program is empty, the program has no mutable segment at all, so loading such a method failed with "No mutable data segments in program", also for a host buffer. The emitter test now loads such a program. Program gets a private load_mutable_subsegment(), which loads a part of a mutable segment into a new buffer and checks the same bounds as load_mutable_subsegment_into(). Both share the code that finds the part. The program test loads the same weights both ways and expects the same bytes, and expects both to fail for bad indices and sizes. --- exir/emit/test/test_emit.py | 20 +++++ runtime/executor/program.cpp | 35 ++++++-- runtime/executor/program.h | 27 ++++++ runtime/executor/targets.bzl | 1 + runtime/executor/tensor_parser_exec_aten.cpp | 89 +++++++++++++++++++- runtime/executor/test/program_test.cpp | 26 ++++++ 6 files changed, 189 insertions(+), 9 deletions(-) diff --git a/exir/emit/test/test_emit.py b/exir/emit/test/test_emit.py index 78480d3b55c..e36f075a7b0 100644 --- a/exir/emit/test/test_emit.py +++ b/exir/emit/test/test_emit.py @@ -297,6 +297,26 @@ def forward(self, x): torch.ones(10, dtype=torch.int64), ) + # An empty buffer has no initial data, so the program has no mutable + # segment, and loading the method must not look for one. + class EmptyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("cache_pos", torch.zeros(0, 2)) + + def forward(self, x): + self.cache_pos.add_(x) + return self.cache_pos + 1 + + x = torch.zeros(0, 2) + et_program_empty = to_edge( + torch.export.export(EmptyModule(), (x,), strict=True) + ).to_executorch(config=et_config) + method_empty = runtime.load_program(et_program_empty.buffer).load_method( + "forward" + ) + self.assertEqual(method_empty.execute((x,))[0].shape, (0, 2)) + def test_int_list_input(self): class M(torch.nn.Module): def forward(self, x, y, z): diff --git a/runtime/executor/program.cpp b/runtime/executor/program.cpp index 538e091a3c2..d9136a49771 100644 --- a/runtime/executor/program.cpp +++ b/runtime/executor/program.cpp @@ -634,12 +634,10 @@ Result Program::LoadSegment( static_cast(absolute_offset), segment->size(), segment_info); } -Error Program::load_mutable_subsegment_into( +Result Program::locate_mutable_subsegment( size_t mutable_data_segments_index, size_t offset_index, - size_t size, - void* buffer) const { - EXECUTORCH_SCOPE_PROF("Program::load_subsegment_into"); + size_t size) const { // Check that the program has segments. if (loader_ == nullptr || segment_base_offset_ == 0) { ET_LOG(Error, "No segments in program"); @@ -744,7 +742,34 @@ Error Program::load_mutable_subsegment_into( "segment base+offset %zu + subsegment offset %zu overflows", base_plus_seg, offset); - return loader_->load_into(total_offset, size, info, buffer); + return MutableSubsegment{total_offset, info}; +} + +Error Program::load_mutable_subsegment_into( + size_t mutable_data_segments_index, + size_t offset_index, + size_t size, + void* buffer) const { + EXECUTORCH_SCOPE_PROF("Program::load_subsegment_into"); + Result subsegment = locate_mutable_subsegment( + mutable_data_segments_index, offset_index, size); + if (!subsegment.ok()) { + return subsegment.error(); + } + return loader_->load_into(subsegment->offset, size, subsegment->info, buffer); +} + +Result Program::load_mutable_subsegment( + size_t mutable_data_segments_index, + size_t offset_index, + size_t size) const { + EXECUTORCH_SCOPE_PROF("Program::load_mutable_subsegment"); + Result subsegment = locate_mutable_subsegment( + mutable_data_segments_index, offset_index, size); + if (!subsegment.ok()) { + return subsegment.error(); + } + return loader_->load(subsegment->offset, size, subsegment->info); } } // namespace ET_RUNTIME_NAMESPACE diff --git a/runtime/executor/program.h b/runtime/executor/program.h index 613137c4983..02d08f6ff40 100644 --- a/runtime/executor/program.h +++ b/runtime/executor/program.h @@ -284,6 +284,33 @@ class Program final { size_t size, void* buffer) const; + /** + * Loads a portion of a mutable segment into a new buffer, for a destination + * such as device memory that the data loader cannot write. + * + * Takes the same arguments as load_mutable_subsegment_into(), except the + * buffer, and fails in the same cases. + * + * @returns The data as a FreeableBuffer of `size` bytes. + */ + ET_NODISCARD Result load_mutable_subsegment( + size_t mutable_data_segments_index, + size_t offset_index, + size_t size) const; + + /// Where a mutable subsegment starts in the data loader's data, and the + /// segment info to load it with. + struct MutableSubsegment { + size_t offset; + DataLoader::SegmentInfo info; + }; + + /// Checks the indices and size of a mutable subsegment and locates it. + ET_NODISCARD Result locate_mutable_subsegment( + size_t mutable_data_segments_index, + size_t offset_index, + size_t size) const; + private: Program( DataLoader* loader, diff --git a/runtime/executor/targets.bzl b/runtime/executor/targets.bzl index 81d0a58667f..20a6c31c51b 100644 --- a/runtime/executor/targets.bzl +++ b/runtime/executor/targets.bzl @@ -128,6 +128,7 @@ def define_common_targets(): ], deps = [ "//executorch/schema:program", + "//executorch/runtime/core:device_allocator", "//executorch/runtime/core/exec_aten/util:tensor_dimension_limit", "//executorch/runtime/core/portable_type/c10/c10:c10", ], diff --git a/runtime/executor/tensor_parser_exec_aten.cpp b/runtime/executor/tensor_parser_exec_aten.cpp index e8b8bf5b6ab..1de898e160f 100644 --- a/runtime/executor/tensor_parser_exec_aten.cpp +++ b/runtime/executor/tensor_parser_exec_aten.cpp @@ -8,6 +8,7 @@ #include +#include #include #include #include @@ -32,6 +33,15 @@ class TensorParser final { return program->load_mutable_subsegment_into( mutable_data_segments_index, offset_index, size, buffer); } + + ET_NODISCARD static Result load_mutable_subsegment( + const Program* program, + size_t mutable_data_segments_index, + size_t offset_index, + size_t size) { + return program->load_mutable_subsegment( + mutable_data_segments_index, offset_index, size); + } }; ET_NODISCARD Result> parseTensorList( @@ -184,6 +194,58 @@ ET_NODISCARD Result getMemPlannedPtr( return allocator->get_offset_address(memory_id, memory_offset, nbytes); } +namespace { + +/// The device of the planned buffer that `allocation_info` points into. A +/// buffer without device metadata is host memory. +::executorch::runtime::etensor::Device getPlannedBufferDevice( + const executorch_flatbuffer::AllocationDetails* allocation_info, + const HierarchicalAllocator* allocator) { + const auto devices = allocator->planned_buffer_devices(); + const uint32_t memory_id = allocation_info->memory_id() - 1; + if (memory_id < devices.size()) { + return devices[memory_id]; + } + return ::executorch::runtime::etensor::DeviceType::CPU; +} + +/// Copies `nbytes` of host `data` to `dst` in the memory of `device`. A data +/// loader writes with the CPU, so it cannot load into device memory directly. +ET_NODISCARD Error copyToDevice( + Result data, + void* dst, + size_t nbytes, + ::executorch::runtime::etensor::Device device) { + if (!data.ok()) { + return data.error(); + } + // A data loader returns no buffer for zero bytes, which allocators reject. + if (nbytes == 0) { + return Error::Ok; + } + ET_CHECK_OR_RETURN_ERROR( + data->size() >= nbytes, + InvalidProgram, + "Loaded %zu bytes for a tensor of %zu bytes", + data->size(), + nbytes); + ::executorch::runtime::DeviceAllocator* device_allocator = + ::executorch::runtime::get_device_allocator(device.type()); + ET_CHECK_OR_RETURN_ERROR( + device_allocator != nullptr, + NotFound, + "No allocator registered for device type %d", + static_cast(device.type())); + ET_CHECK_OK_OR_RETURN_ERROR(device_allocator->copy_host_to_device( + dst, data->data(), nbytes, device.index())); + // The first run may read the tensor on a stream the copy is not ordered + // with, for example when a caller sets its stream only around execute. This + // also lets `data` be freed when this returns. + return device_allocator->synchronize(device.index()); +} + +} // namespace + ET_NODISCARD Result getTensorDataPtr( const executorch_flatbuffer::Tensor* s_tensor, const Program* program, @@ -246,8 +308,14 @@ ET_NODISCARD Result getTensorDataPtr( if (!planned_ptr.ok()) { return planned_ptr.error(); } - auto load_error = - named_data_map->load_data_into(fqn, planned_ptr.get(), nbytes); + const auto device = getPlannedBufferDevice(allocation_info, allocator); + auto load_error = device.is_cpu() + ? named_data_map->load_data_into(fqn, planned_ptr.get(), nbytes) + : copyToDevice( + named_data_map->get_data(fqn), + planned_ptr.get(), + nbytes, + device); if (load_error != Error::Ok) { return load_error; } @@ -273,8 +341,21 @@ ET_NODISCARD Result getTensorDataPtr( if (!planned_ptr.ok()) { return planned_ptr.error(); } - auto err = TensorParser::load_mutable_subsegment_into( - program, 0, s_tensor->data_buffer_idx(), nbytes, planned_ptr.get()); + // A program whose initial data is all empty has no mutable segment to + // load an empty tensor from. + if (nbytes == 0) { + return planned_ptr; + } + const auto device = getPlannedBufferDevice(allocation_info, allocator); + auto err = device.is_cpu() + ? TensorParser::load_mutable_subsegment_into( + program, 0, data_buffer_idx, nbytes, planned_ptr.get()) + : copyToDevice( + TensorParser::load_mutable_subsegment( + program, 0, data_buffer_idx, nbytes), + planned_ptr.get(), + nbytes, + device); if (err != Error::Ok) { return err; diff --git a/runtime/executor/test/program_test.cpp b/runtime/executor/test/program_test.cpp index 5787aaecabb..8d349e53387 100644 --- a/runtime/executor/test/program_test.cpp +++ b/runtime/executor/test/program_test.cpp @@ -111,6 +111,15 @@ class ProgramTestFriend final { mutable_data_segments_index, offset_index, size, buffer); } + ET_NODISCARD static Result load_mutable_subsegment( + const Program* program, + size_t mutable_data_segments_index, + size_t offset_index, + size_t size) { + return program->load_mutable_subsegment( + mutable_data_segments_index, offset_index, size); + } + const static executorch_flatbuffer::Program* GetInternalProgram( const Program* program) { return program->internal_program_; @@ -861,6 +870,13 @@ TEST_F(ProgramTest, LoadFromMutableSegment) { EXPECT_EQ(std::memcmp(buffer, buffer2, kWeightSize), 0) << "Loading the same segment twice should return identical data"; + // Loading into a new buffer, as for device memory, returns the same data. + Result loaded = ProgramTestFriend::load_mutable_subsegment( + &program.get(), 0, 1, kWeightSize); + ASSERT_EQ(loaded.error(), Error::Ok); + ASSERT_EQ(loaded->size(), kWeightSize); + EXPECT_EQ(std::memcmp(loaded->data(), buffer, kWeightSize), 0); + // Verify that the loaded data is non-trivial (not all zeros). bool has_nonzero = false; for (size_t i = 0; i < kWeightSize; ++i) { @@ -921,6 +937,16 @@ TEST_F(ProgramTest, LoadFromMutableSegment) { err = ProgramTestFriend::load_mutable_subsegment_into( &program.get(), 500, 1, 1, buffer); EXPECT_NE(err, Error::Ok); + + // The new buffer variant checks the same bounds. + EXPECT_NE( + ProgramTestFriend::load_mutable_subsegment(&program.get(), 0, 1, 500) + .error(), + Error::Ok); + EXPECT_NE( + ProgramTestFriend::load_mutable_subsegment(&program.get(), 0, 500, 1) + .error(), + Error::Ok); } TEST_F(ProgramTest, LoadAndCheckPTESize) { From f631ac76ce1aa435d3051af3a2d269fefb650a50 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 25 Oct 2026 17:25:02 -0700 Subject: [PATCH 3/4] Add skip_h2d_for_constants to keep constants in device memory With skip_h2d_for_method_inputs, a method input that a device delegate reads stays in device memory. A constant, such as a weight or a buffer the method does not change, was still copied to the device on every call, because its data is stored in the program, on the host. The new opt-in PropagateDeviceConfig option skip_h2d_for_constants keeps a constant that only delegates on one device read in device memory. PropagateDevicePass tags it with that device instead of adding a copy. Memory planning gives it planned device memory for the whole method, also when graph inputs are not planned, because the method reads it on every call. The emitter stores its data as the initial state of that memory, as it does for a mutable buffer with an initial value, and the runtime loads it there once, when it loads the method. Each method that reads the constant keeps its own copy, so this is opt-in. A constant that anything else reads, such as a host op or a delegate on another device, is still copied, with the same warning, and error_on_per_call_copies turns that copy into an error under this option too. A constant that is changed on the host before a delegate reads it, such as by a view, a transpose or a clone, is not read by the delegate directly, so it is still copied; the docs say so. They also say that share_mutable_buffers, which needs every planned tensor in host memory, does not work with this option. The CUDA guide now says that every skip flag needs enable_non_cpu_memory_planning, and that skip_h2d_for_constants works on its own, without skip_h2d_for_method_inputs or unplanned inputs. The new Python test checks that a buffer or parameter read only by delegates on one device has no copy and no warning, has planned device memory for the whole method with its data as the initial state, with and without planned graph inputs, and that a constant also read on the host or on a second device, or any constant without the option, is still copied. It fails without this change. export_program_with_device_info.py now also exports such a program, on cuda:1, once with the constant in the program and once in a separate .ptd file. Two new C++ tests parse the constant of each into planned memory on a mock CUDA device and check that its data was copied there, to device 1, and that the copy was waited for. tensor_parser_device_test now also builds and runs with CMake. A new CUDA test runs a program whose constant is kept in device memory twice, with different inputs, and compares with eager PyTorch, with error_on_per_call_copies set so it fails if the constant is copied. --- backends/cuda/tests/test_device_placement.py | 34 ++++ docs/source/backends/cuda/cuda-overview.md | 28 ++- exir/emit/_emitter.py | 4 +- exir/memory_planning.py | 22 ++- exir/passes/propagate_device_config.py | 30 ++- exir/passes/propagate_device_pass.py | 44 ++++- exir/program/_program.py | 1 + exir/tests/test_propagate_device_pass.py | 172 +++++++++++++++++- runtime/executor/test/CMakeLists.txt | 26 +++ runtime/executor/test/targets.bzl | 4 + .../test/tensor_parser_device_test.cpp | 126 ++++++++++++- .../models/export_program_with_device_info.py | 53 +++++- test/models/targets.bzl | 3 + 13 files changed, 516 insertions(+), 31 deletions(-) diff --git a/backends/cuda/tests/test_device_placement.py b/backends/cuda/tests/test_device_placement.py index b54242802a1..95764a8338e 100644 --- a/backends/cuda/tests/test_device_placement.py +++ b/backends/cuda/tests/test_device_placement.py @@ -145,3 +145,37 @@ def forward(self, x, y): outputs = method.execute([t.cuda() for t in inputs]) for output, expected in zip(outputs, model(*inputs), strict=True): torch.testing.assert_close(output.cpu(), expected) + + def test_constant_kept_in_device_memory(self): + """With skip_h2d_for_constants, a constant that only delegates read is + loaded into device memory once, when the method loads, and later + delegates must not reuse that memory, so every call reads it. + error_on_per_call_copies makes the export fail if the constant is + copied on every call instead.""" + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("table", torch.arange(6.0).reshape(2, 3)) + + def forward(self, x): + a = torch.add(x, self.table) + b = torch.add(a, a) + return torch.add(b, b) + + config = ExecutorchBackendConfig( + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True, + skip_h2d_for_constants=True, + error_on_per_call_copies=True, + ), + enable_non_cpu_memory_planning=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + ) + model = Model() + x = torch.arange(6.0).reshape(2, 3) + with tempfile.TemporaryDirectory() as directory: + method = load(model, (x,), directory, config) + for scale in (1.0, 10.0): + (output,) = method.execute([(x * scale).cuda()]) + torch.testing.assert_close(output.cpu(), model(x * scale)) diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md index 32effc7e520..518073101d9 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -164,13 +164,14 @@ overwrite, unless you pass `clone_outputs=False`. ### Things worth knowing -**The two flags are independent.** Skip only the input copies if you produce data on the GPU but -want results back on the host, or only the output copies for the reverse. +**`skip_h2d_for_method_inputs` and `skip_d2h_for_method_outputs` are independent.** Skip only the +input copies if you produce data on the GPU but want results back on the host, or only the output +copies for the reverse. -**Both flags need `enable_non_cpu_memory_planning=True`.** It defaults to `True`, so leaving it -out works, and the snippet above sets it explicitly only to make the requirement visible. Setting -it to `False` while asking for either skip raises a `ValueError`, because copy insertion happens -during device-aware memory planning. +**Every skip flag needs `enable_non_cpu_memory_planning=True`.** It defaults to `True`, so leaving +it out works, and the snippet above sets it explicitly only to make the requirement visible. +Setting it to `False` while asking for any skip raises a `ValueError`, because copy insertion +happens during device-aware memory planning. **Inputs must be unplanned**, via `MemoryPlanningPass(alloc_graph_input=False)`. Without it the program reserves its own input buffer on the device, and the runtime refuses to copy the caller's @@ -213,7 +214,7 @@ the graph in the delegate. `skip_d2h_for_method_outputs`, read a copy in host memory. A delegate on another device then gets a second copy, from the host. - A constant, such as a weight, that a delegate reads from outside is copied to each device that - reads it. + reads it. `skip_h2d_for_constants` (below) can remove this copy. - A returned output that an operator outside the delegate writes in place, such as a mutable custom operator left on the host, is copied to the host and returned there. The method metadata gives the device of each output. @@ -221,6 +222,19 @@ the graph in the delegate. Add `error_on_per_call_copies=True` to `PropagateDeviceConfig` to make `to_executorch` fail instead of warning. It fails only on a copy that one of the skip flags you set was meant to remove. +**Set `skip_h2d_for_constants=True` to keep constants in device memory.** If only delegates on +one device read a constant, it is loaded into that device's memory once, when the method loads, +and not copied on every call. The cost: it takes device memory for as long as the method is +loaded, and each method that reads it has its own copy. Its data is stored as the initial state +of that memory, so if the same data is also stored as an ordinary constant, for example for a +method that reads it on the host, the `.pte` holds it twice. A constant that a host operator or a +second device also reads is still copied on every call. So is a constant that is changed on the +host before the delegate reads it, such as by a view, a transpose or a clone, and +`error_on_per_call_copies` does not catch that copy. Have the partitioner take that change into +the delegate, so that the delegate reads the constant itself. `share_mutable_buffers` needs every +planned tensor in host memory, so it does not work with this flag. This flag works on its own: it +needs neither `skip_h2d_for_method_inputs` nor unplanned inputs. + **`skip_h2d_for_method_inputs` refuses inputs that the method changes.** `to_executorch` fails if the method changes an input, even when the delegate makes the change. This includes a mutable buffer, such as a KV cache, even though the caller does not pass it. For a buffer, put it and every diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py index efdfaab00cd..8bbeedd59f2 100644 --- a/exir/emit/_emitter.py +++ b/exir/emit/_emitter.py @@ -2057,11 +2057,13 @@ def _is_buffer(node: Node, graph_signature: ExportGraphSignature) -> bool: # overall be considered a mutable buffer if it has already been memory planned. This would # suggest that the same abstract buffer is mutable in another entry point so we should # compel it to be considered mutable in all entry points at emission just as the user did with - # memory planning. + # memory planning. A device resident constant is planned too, but + # it is not mutable. is_mutable_buffer |= ( _is_buffer(self.node, self.exported_program.graph_signature) and spec.mem_id is not None and spec.mem_offset is not None + and not self.node.meta.get("device_resident_constant", False) ) # If the placeholder has a constant_tag, it is external to the PTE file diff --git a/exir/memory_planning.py b/exir/memory_planning.py index 79642091a35..ee04ab9ddac 100644 --- a/exir/memory_planning.py +++ b/exir/memory_planning.py @@ -351,8 +351,11 @@ def update_tensor_lifetime( else: start = node_idx if start is None or start > node_idx else start - if node.op == "placeholder" and _is_mutable_buffer(node, gs): - # mutable buffers are never freed + if node.op == "placeholder" and ( + _is_mutable_buffer(node, gs) or _is_device_resident_constant(node) + ): + # Mutable buffers and device resident constants keep their data across + # calls, so they are never freed. end = max_node_idx else: end = node_idx if end is None or end < node_idx else end @@ -386,6 +389,14 @@ def _is_mutable_buffer( return False +def _is_device_resident_constant(node: Node) -> bool: + """Whether node is a constant that PropagateDevicePass placed in device + memory. The runtime copies its data from the program or its data file + into its planned memory once, at load, so it is planned like a mutable + buffer instead of being used in place.""" + return node.op == "placeholder" and node.meta.get("device_resident_constant", False) + + def _get_mutable_buffer_specs( nodes: Iterable[Node], graph_signature: Optional[ExportGraphSignature] ) -> Set[TensorSpec]: @@ -425,7 +436,11 @@ def get_graph_input_tensors( ) -> Set[TensorSpec]: graph_input_tensors = set() for node in nodes: - if node.op == "placeholder" and not _is_mutable_buffer(node, graph_signature): + if ( + node.op == "placeholder" + and not _is_mutable_buffer(node, graph_signature) + and not _is_device_resident_constant(node) + ): for spec in get_node_tensor_specs(node): graph_input_tensors.add(spec) @@ -535,6 +550,7 @@ def collect_specs_from_nodes( # noqa: C901 ignore_const and spec.const and not node.meta.get("weight_has_gradient", False) + and not _is_device_resident_constant(node) ): continue if dedup: diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py index 3d6dfeeefda..208a0db3b01 100644 --- a/exir/passes/propagate_device_config.py +++ b/exir/passes/propagate_device_config.py @@ -25,7 +25,7 @@ @compatibility(is_backward_compatible=False) @dataclass class PropagateDeviceConfig: - """Controls whether the runtime copies method inputs and outputs across the device boundary. + """Controls whether the runtime copies method inputs, outputs and constants across the device boundary. Memory planning reserves buffers for graph inputs and outputs by default. A program that skips the input copy must use ``alloc_graph_input=False``. Unplanning the outputs as well, as below, @@ -46,6 +46,8 @@ class PropagateDeviceConfig: already owns and refuses to copy into it, so every call fails with ``NotSupported``. Planned outputs work, but they live in the method's memory, which the next run overwrites, and the Python ``execute()`` copies them unless it gets ``clone_outputs=False``. + + ``skip_h2d_for_constants`` does not need ``alloc_graph_input=False``: it works with planned inputs. """ # When True, method-level input tensors that feed directly into a device @@ -61,7 +63,8 @@ class PropagateDeviceConfig: # by the caller, so they are still copied: their data stays on the host, # in the program or in its external data file. That copy runs on every # call, once for each device that reads the constant, and stays in device - # memory until the last delegate there has read it. + # memory until the last delegate there has read it. Set + # skip_h2d_for_constants to keep constants on the device instead. # # Pair with MemoryPlanningPass(alloc_graph_input=False), or the runtime reserves its own buffer # for the input and refuses the caller's tensor. @@ -85,21 +88,37 @@ class PropagateDeviceConfig: # propagate_device_config of ExecutorchBackendConfig. skip_d2h_for_method_outputs: Union[bool, Dict[str, bool]] = False - # When True, the copies that the skip flags above cannot remove make + # When True, the copies that the skip flags cannot remove make # to_executorch raise instead of logging a warning: with - # skip_h2d_for_method_inputs, a constant copied to the device and a method - # input copied to the host for a reader off its device, and with + # skip_h2d_for_method_inputs or skip_h2d_for_constants, a constant copied + # to the device; with skip_h2d_for_method_inputs, a method input copied to + # the host for a reader off its device; and with # skip_d2h_for_method_outputs, a returned output copied to the host for a # host op that writes it. Each runs on every call. Use it to make sure a # program that skips its boundary copies has none left. error_on_per_call_copies: bool = False + # When True, a constant that only delegates on one device read is kept in + # that device's memory instead of being copied there on every call. The + # runtime loads its data there once, when it loads the method, through + # host memory. The constant then uses device memory for as long as the + # method is loaded, and each method that reads it keeps its own copy. A + # constant that anything else reads is still copied on every call. Only a + # constant that a delegate reads directly is kept: one changed on the host + # first, such as by a view, a transpose or a clone, is copied on every + # call, and error_on_per_call_copies does not catch that copy. To keep + # it, have the partitioner take that change into the delegate too. + # MemoryPlanningPass(share_mutable_buffers=True) needs every planned tensor + # in host memory, so it does not work with this option. + skip_h2d_for_constants: bool = False + def __hash__(self) -> int: return hash( ( str(self.skip_h2d_for_method_inputs), str(self.skip_d2h_for_method_outputs), self.error_on_per_call_copies, + self.skip_h2d_for_constants, ) ) @@ -110,4 +129,5 @@ def __eq__(self, other: object) -> bool: self.skip_h2d_for_method_inputs == other.skip_h2d_for_method_inputs and self.skip_d2h_for_method_outputs == other.skip_d2h_for_method_outputs and self.error_on_per_call_copies == other.error_on_per_call_copies + and self.skip_h2d_for_constants == other.skip_h2d_for_constants ) diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py index 4e3ed5035bc..235ab7ddbd1 100644 --- a/exir/passes/propagate_device_pass.py +++ b/exir/passes/propagate_device_pass.py @@ -303,7 +303,13 @@ class PropagateDevicePass(PassBase): instead (user provides GPU tensor at runtime). Every other reader of the placeholder reads one D2H copy of it. A placeholder that an op writes in place on the host is refused. Constants are placeholders too, but their - data is stored on the host, so they still get an H2D copy. + data is stored on the host, so they still get an H2D copy unless + skip_h2d_for_constants keeps them on the device. + - skip_h2d_for_constants: A constant that only delegates on one device + read is tagged with that device instead of getting an H2D copy. Memory + planning keeps it in device memory for the whole method, and the runtime + loads its data there once, when it loads the method. Any other constant + still gets an H2D copy. - skip_d2h_for_method_outputs: The graph output reads a delegate output on device, not its D2H copy; the copy stays only for other readers. When one of them writes the copy in place, the graph output reads the written @@ -320,18 +326,23 @@ def __init__( skip_d2h_for_method_outputs: bool = False, enable_non_cpu_memory_planning: bool = False, error_on_per_call_copies: bool = False, + skip_h2d_for_constants: bool = False, ) -> None: super().__init__() self.skip_h2d_for_method_inputs = skip_h2d_for_method_inputs self.skip_d2h_for_method_outputs = skip_d2h_for_method_outputs self.enable_non_cpu_memory_planning = enable_non_cpu_memory_planning self.error_on_per_call_copies = error_on_per_call_copies + self.skip_h2d_for_constants = skip_h2d_for_constants if ( - skip_h2d_for_method_inputs or skip_d2h_for_method_outputs + skip_h2d_for_method_inputs + or skip_d2h_for_method_outputs + or skip_h2d_for_constants ) and not enable_non_cpu_memory_planning: raise ValueError( - "skip_h2d_for_method_inputs and skip_d2h_for_method_outputs are " + "skip_h2d_for_method_inputs, skip_d2h_for_method_outputs and " + "skip_h2d_for_constants are " "only meaningful when enable_non_cpu_memory_planning=True, since " "they control host/device copy insertion which only happens during " "device-aware memory planning. Set enable_non_cpu_memory_planning=" @@ -371,10 +382,10 @@ def _insert_h2d_copies( ) -> bool: """Make each tensor input of a delegate call read device memory. - A method input that is not a constant moves to the device under - skip_h2d_for_method_inputs. - Any other input reads an H2D copy, reused from an earlier delegate on - the same device when possible. + Under skip_h2d_for_method_inputs, a method input moves to the device, + and under skip_h2d_for_constants, so does a constant that only + delegates on that device read. Any other input reads an H2D copy, + reused from an earlier delegate on the same device when possible. """ changed = False new_args = list(node.args) @@ -411,6 +422,23 @@ def _insert_h2d_copies( f"it is created, for example in the pass that adds it." ) + if ( + self.skip_h2d_for_constants + and arg_spec.const + and self._is_placeholder(arg) + and all( + self._delegate_device(graph_module, user) + == (target_device_type, device_index) + for user in arg.users + ) + ): + # Memory planning keeps it in device memory for the whole + # method, and the runtime loads its data there once. + _set_device_on_spec(arg_spec, target_device_type, device_index) + arg.meta["device_resident_constant"] = True + changed = True + continue + h2d_node = _reusable_copy(arg, node, target_device_type, device_index) if h2d_node is None: if arg_spec.const: @@ -420,7 +448,7 @@ def _insert_h2d_copies( f"{device_index} from outside it, so it is copied to the " f"device on every call. Have the partitioner take it into " f"the delegate to avoid that copy.", - self.skip_h2d_for_method_inputs, + self.skip_h2d_for_method_inputs or self.skip_h2d_for_constants, ) with graph_module.graph.inserting_before(node): h2d_node = graph_module.graph.call_function( diff --git a/exir/program/_program.py b/exir/program/_program.py index b1d896281c0..118fc5a5860 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -794,6 +794,7 @@ def edge_to_executorch_passes( skip_d2h_for_method_outputs=device_cfg.skip_d2h_for_method_outputs, enable_non_cpu_memory_planning=config.enable_non_cpu_memory_planning, error_on_per_call_copies=device_cfg.error_on_per_call_copies, + skip_h2d_for_constants=device_cfg.skip_h2d_for_constants, ), EdgeToBackendOpsPass(), RemoveGraphAssertsPass(), diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py index 612e8a06290..26dcb04032e 100644 --- a/exir/tests/test_propagate_device_pass.py +++ b/exir/tests/test_propagate_device_pass.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import itertools import operator import unittest from copy import deepcopy @@ -241,6 +242,19 @@ def _host_writer_config(write_view: bool, **kwargs) -> ExecutorchBackendConfig: ) +def _skip_h2d_for_constants_config( + alloc_graph_input: bool, skip_constants: bool = True +) -> ExecutorchBackendConfig: + return ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True, skip_h2d_for_constants=skip_constants + ), + enable_non_cpu_memory_planning=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=alloc_graph_input), + ) + + class TestPropagateDevicePass(unittest.TestCase): @staticmethod def _collect_tensor_specs(node: torch.fx.Node) -> List[TensorSpec]: @@ -1881,9 +1895,10 @@ def forward(self, a, b): ) def test_constant_feeding_device_delegate_is_copied(self): - """A constant feeding a device delegate gets an _h2d_copy into planned device - memory, with or without skip_h2d_for_method_inputs: its data lives in the - program on the host, so it is not a method input the caller can place.""" + """Without skip_h2d_for_constants, a constant feeding a device delegate + gets an _h2d_copy into planned device memory, with or without + skip_h2d_for_method_inputs: its data lives in the program on the host, + so it is not a method input the caller can place.""" class Model(torch.nn.Module): def __init__(self): @@ -1922,6 +1937,157 @@ def forward(self, a): nodes.h2d_nodes, program, DeviceType.CUDA, pipeline, "H2D copy" ) + def test_skip_h2d_for_constants_keeps_constant_in_device_memory(self): + """With skip_h2d_for_constants=True, a buffer or parameter that only + delegates on one device read, one or several, is planned in that + device's memory for the whole method, with its data as the initial + state, so it needs no copy and no warning. Graph inputs being + unplanned does not change that.""" + + class Buffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("offset", torch.randn(2, 2)) + + def forward(self, a): + return torch.add(torch.add(a, self.offset), a) + + class Parameter(torch.nn.Module): + def __init__(self): + super().__init__() + self.offset = torch.nn.Parameter(torch.randn(2, 2)) + + def forward(self, a): + return torch.add(a, self.offset) + + class TwoDelegates(Buffer): + def forward(self, a): + return torch.add(torch.add(a, self.offset), self.offset) + + for alloc_graph_input, model in itertools.product( + (True, False), (Buffer(), Parameter(), TwoDelegates()) + ): + with self.assertNoLogs(propagate_device_pass.logger, "WARNING"): + results = self._get_executorch_program( + model, + (torch.randn(2, 2),), + PerAddDeviceDelegatePartitioner(["cuda:0"]), + _skip_h2d_for_constants_config(alloc_graph_input), + ) + for pipeline, program, gm in results: + with self.subTest( + pipeline=pipeline, + alloc_input=alloc_graph_input, + model=type(model).__name__, + ): + self.assertEqual(_collect_device_copy_nodes(gm).h2d_nodes, []) + (offset,) = [ + n + for n in gm.graph.nodes + if n.op == "placeholder" and n.name != "a" + ] + spec = offset.meta["spec"] + self.assertTrue(spec.const) + self.assertEqual(spec.lifetime, [0, len(gm.graph.nodes) - 1]) + self._assert_buffer_device( + spec, program, DeviceType.CUDA, "device resident constant" + ) + (tensor,) = [ + value.val + for value in program.execution_plan[0].values + if getattr(value.val, "allocation_info", None) is not None + and value.val.data_buffer_idx > 0 + ] + self.assertEqual( + tensor.extra_tensor_info.device_type, DeviceType.CUDA + ) + + def test_skip_h2d_for_constants_copies_constants_read_elsewhere(self): + """A constant that a host op, or a delegate on another device, reads + too is still copied with skip_h2d_for_constants=True, and so is every + constant without the option.""" + + class Buffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("offset", torch.randn(2, 2)) + + def forward(self, a): + return torch.add(a, self.offset) + + class AlsoReadOnHost(Buffer): + def forward(self, a): + return torch.add(a, self.offset), self.offset * 2 + + class ReadByTwoDelegates(Buffer): + def forward(self, a): + return torch.add(a, self.offset), torch.add(a * 2, self.offset) + + cases = { + "host reader": (AlsoReadOnHost(), ["cuda:0"], True, [0]), + "two devices": (ReadByTwoDelegates(), ["cuda:0", "cuda:1"], True, [0, 1]), + "without the option": (Buffer(), ["cuda:0"], False, [0]), + } + for case, (model, devices, skip_constants, copied_to) in cases.items(): + with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs: + results = self._get_executorch_program( + model, + (torch.randn(2, 2),), + PerAddDeviceDelegatePartitioner(devices), + _skip_h2d_for_constants_config(False, skip_constants), + ) + self.assertIn("Constant 'b_offset'", logs.output[0]) + for pipeline, _program, gm in results: + with self.subTest(pipeline=pipeline, case=case): + copies = [ + n.meta["spec"].device_index + for n in _collect_device_copy_nodes(gm).h2d_nodes + if n.args[0].name == "b_offset" + ] + self.assertEqual(sorted(copies), copied_to) + + def test_skip_h2d_for_constants_on_its_own(self): + """skip_h2d_for_constants works without skip_h2d_for_method_inputs: + the method input is still copied to the device and the constant is + not. A constant that is still copied raises with + error_on_per_call_copies.""" + + class Buffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("offset", torch.randn(2, 2)) + + def forward(self, a): + return torch.add(a, self.offset) + + class AlsoReadOnHost(Buffer): + def forward(self, a): + return torch.add(a, self.offset), self.offset * 2 + + def config(error_on_per_call_copies): + return ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_constants=True, + error_on_per_call_copies=error_on_per_call_copies, + ), + enable_non_cpu_memory_planning=True, + ) + + inputs = (torch.randn(2, 2),) + for pipeline, _program, gm in self._get_executorch_program( + Buffer(), inputs, DeviceAwarePartitioner("cuda:0"), config(False) + ): + with self.subTest(pipeline=pipeline): + (copy,) = _collect_device_copy_nodes(gm).h2d_nodes + self.assertEqual(copy.args[0].name, "a") + with self.assertRaisesRegex( + RuntimeError, "Constant 'b_offset'.*error_on_per_call_copies=True" + ): + self._get_executorch_program( + AlsoReadOnHost(), inputs, DeviceAwarePartitioner("cuda:0"), config(True) + ) + def test_error_on_per_call_copies(self): """error_on_per_call_copies turns the warning about a copy that a skip flag cannot remove into an error: a constant copied to the device and diff --git a/runtime/executor/test/CMakeLists.txt b/runtime/executor/test/CMakeLists.txt index af81f4acb4a..562a9aa5ce1 100644 --- a/runtime/executor/test/CMakeLists.txt +++ b/runtime/executor/test/CMakeLists.txt @@ -86,6 +86,9 @@ add_custom_target( add_custom_command( OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddWithDevice.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddConstantWithDevice.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddConstantWithDeviceProgram.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddConstantWithDeviceProgram.ptd" COMMAND ${PYTHON_EXECUTABLE} -m test.models.export_program_with_device_info --outdir "${CMAKE_CURRENT_BINARY_DIR}" WORKING_DIRECTORY ${EXECUTORCH_ROOT} @@ -94,6 +97,9 @@ add_custom_command( add_custom_target( method_meta_test_pte_files DEPENDS "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddWithDevice.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddConstantWithDevice.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddConstantWithDeviceProgram.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddConstantWithDeviceProgram.ptd" ) set(test_env @@ -110,6 +116,9 @@ set(test_env "ET_MODULE_STATEFUL_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleStateful.pte" "ET_MODULE_ADD_MUL_DELEGATED_PATH=${CMAKE_CURRENT_BINARY_DIR}/delegated/ModuleAddMul.pte" "ET_MODULE_ADD_WITH_DEVICE_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddWithDevice.pte" + "ET_MODULE_ADD_CONSTANT_WITH_DEVICE_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddConstantWithDevice.pte" + "ET_MODULE_ADD_CONSTANT_WITH_DEVICE_PROGRAM_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddConstantWithDeviceProgram.pte" + "ET_MODULE_ADD_CONSTANT_WITH_DEVICE_DATA_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddConstantWithDeviceProgram.ptd" "ET_MODULE_ADD_MUL_DELEGATED_SCRATCH_PATH=${CMAKE_CURRENT_BINARY_DIR}/delegated/ModuleAddMul-scratch.pte" ) @@ -163,6 +172,23 @@ target_include_directories( "${EXECUTORCH_ROOT}/third-party/flatbuffers/include" ) +et_cxx_test( + tensor_parser_device_test + SOURCES + tensor_parser_device_test.cpp + EXTRA_LIBS + extension_data_loader + extension_flat_tensor + program_schema +) +add_dependencies(tensor_parser_device_test method_meta_test_pte_files) +set_property(TEST tensor_parser_device_test PROPERTY ENVIRONMENT ${test_env}) +target_include_directories( + tensor_parser_device_test + PRIVATE "${CMAKE_INSTALL_PREFIX}/schema/include" + "${EXECUTORCH_ROOT}/third-party/flatbuffers/include" +) + et_cxx_test( program_test SOURCES program_test.cpp EXTRA_LIBS extension_data_loader program_schema diff --git a/runtime/executor/test/targets.bzl b/runtime/executor/test/targets.bzl index 314eaae3c8a..cdb53c155b2 100644 --- a/runtime/executor/test/targets.bzl +++ b/runtime/executor/test/targets.bzl @@ -358,9 +358,13 @@ def define_common_targets(is_fbcode = False): "//executorch/runtime/core:device_memory_buffer", "//executorch/runtime/core/test:mock_cuda_allocator", "//executorch/extension/data_loader:file_data_loader", + "//executorch/extension/flat_tensor:flat_tensor_data_map", "//executorch/schema:program", ], env = { "ET_MODULE_ADD_WITH_DEVICE_PATH": "$(location fbcode//executorch/test/models:exported_program_with_device_info[ModuleAddWithDevice.pte])", + "ET_MODULE_ADD_CONSTANT_WITH_DEVICE_PATH": "$(location fbcode//executorch/test/models:exported_program_with_device_info[ModuleAddConstantWithDevice.pte])", + "ET_MODULE_ADD_CONSTANT_WITH_DEVICE_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_program_with_device_info[ModuleAddConstantWithDeviceProgram.pte])", + "ET_MODULE_ADD_CONSTANT_WITH_DEVICE_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_with_device_info[ModuleAddConstantWithDeviceProgram.ptd])", }, ) diff --git a/runtime/executor/test/tensor_parser_device_test.cpp b/runtime/executor/test/tensor_parser_device_test.cpp index aafa6d58f8a..097e338f2c3 100644 --- a/runtime/executor/test/tensor_parser_device_test.cpp +++ b/runtime/executor/test/tensor_parser_device_test.cpp @@ -8,15 +8,18 @@ /** * Tests that device info (device_type) is correctly parsed from serialized - * tensors in .pte files into TensorImpl at runtime. + * tensors in .pte files into TensorImpl at runtime, and that the initial data + * of a tensor planned in device memory is copied there. * - * Uses a .pte exported with DeviceAwarePartitioner (CUDA device annotation) - * so that delegate output tensors carry device_type=CUDA in ExtraTensorInfo. + * Uses .pte files exported with DeviceAwarePartitioner (CUDA device + * annotation) so that delegate tensors carry device_type=CUDA in + * ExtraTensorInfo. */ #include #include +#include #include #include #include @@ -28,12 +31,14 @@ #include using executorch::aten::Tensor; +using executorch::extension::FlatTensorDataMap; using executorch::runtime::DeviceMemoryBuffer; using executorch::runtime::Error; using executorch::runtime::HierarchicalAllocator; using executorch::runtime::MemoryAllocator; using executorch::runtime::MemoryManager; using executorch::runtime::MethodMeta; +using executorch::runtime::NamedDataMap; using executorch::runtime::Program; using executorch::runtime::register_device_allocator; using executorch::runtime::Result; @@ -283,3 +288,118 @@ TEST_F(TensorParserDeviceTest, CudaTensorDataPtrPointsToDeviceMemory) { EXPECT_EQ(cuda_with_device_memory, 3) << "All 3 CUDA tensors should have data_ptr in device memory"; } + +namespace { + +/// Parses the constant of ModuleAddConstantWithDevice, which is planned in +/// cuda:1 memory with its data as the initial state, into planned memory that +/// records which buffer is on which device, as Module and the Python bindings +/// set it up. +void expectConstantCopiedToDevice( + const char* program_env, + const char* data_env) { + const char* program_path = std::getenv(program_env); + ASSERT_NE(program_path, nullptr) << program_env << " env var not set"; + Result loader = FileDataLoader::from(program_path); + ASSERT_EQ(loader.error(), Error::Ok); + Result program = + Program::load(&loader.get(), Program::Verification::Minimal); + ASSERT_EQ(program.error(), Error::Ok); + + std::unique_ptr data_loader; + std::unique_ptr data_map; + if (data_env != nullptr) { + const char* data_path = std::getenv(data_env); + ASSERT_NE(data_path, nullptr) << data_env << " env var not set"; + Result loaded = FileDataLoader::from(data_path); + ASSERT_EQ(loaded.error(), Error::Ok); + data_loader = std::make_unique(std::move(loaded.get())); + Result map = FlatTensorDataMap::load(data_loader.get()); + ASSERT_EQ(map.error(), Error::Ok); + data_map = std::make_unique(std::move(map.get())); + } + + Result method_meta = program->method_meta("forward"); + ASSERT_EQ(method_meta.error(), Error::Ok); + std::vector> planned_spans; + std::vector planned_devices; + std::vector> cpu_buffers; + std::vector device_buffers; + for (size_t i = 0; i < method_meta->num_memory_planned_buffers(); ++i) { + auto size = method_meta->memory_planned_buffer_size(i); + ASSERT_TRUE(size.ok()); + auto device = method_meta->memory_planned_buffer_device(i); + ASSERT_TRUE(device.ok()); + planned_devices.push_back(device.get()); + if (device->is_cpu()) { + cpu_buffers.emplace_back(size.get()); + planned_spans.emplace_back( + cpu_buffers.back().data(), cpu_buffers.back().size()); + } else { + auto dmb = DeviceMemoryBuffer::create( + size.get(), device->type(), device->index()); + ASSERT_TRUE(dmb.ok()); + planned_spans.emplace_back(dmb->as_span()); + device_buffers.push_back(std::move(dmb.get())); + } + } + ASSERT_EQ(device_buffers.size(), 1); + HierarchicalAllocator planned_memory( + {planned_spans.data(), planned_spans.size()}, + {planned_devices.data(), planned_devices.size()}); + constexpr size_t kMethodAllocBytes = 32 * 1024U; + auto method_alloc_pool = std::make_unique(kMethodAllocBytes); + MemoryAllocator method_allocator(kMethodAllocBytes, method_alloc_pool.get()); + MemoryManager memory_manager(&method_allocator, &planned_memory); + + // The constant is the only tensor with both initial data and planned memory. + auto* values = ProgramTestFriend::GetInternalProgram(&program.get()) + ->execution_plan() + ->Get(0) + ->values(); + const executorch_flatbuffer::Tensor* constant = nullptr; + for (size_t i = 0; i < values->size(); ++i) { + auto* s_tensor = values->Get(i)->val_as_Tensor(); + bool has_data = s_tensor != nullptr && + (s_tensor->data_buffer_idx() > 0 || + (s_tensor->extra_tensor_info() != nullptr && + s_tensor->extra_tensor_info()->location() == + executorch_flatbuffer::TensorDataLocation::EXTERNAL)); + if (has_data && s_tensor->allocation_info() != nullptr) { + ASSERT_EQ(constant, nullptr) << "more than one such tensor"; + constant = s_tensor; + } + } + ASSERT_NE(constant, nullptr); + + const int h2d_count = g_mock_cuda.h2d_count_; + Result tensor = parseTensor( + &program.get(), + &memory_manager, + constant, + static_cast(data_map.get())); + ASSERT_EQ(tensor.error(), Error::Ok); + EXPECT_EQ(g_mock_cuda.h2d_count_, h2d_count + 1); + EXPECT_EQ(g_mock_cuda.last_h2d_index_, 1); + EXPECT_EQ(g_mock_cuda.synced_h2d_count_, h2d_count + 1) + << "the copy was not waited for before parsing returned"; + EXPECT_EQ(g_mock_cuda.last_sync_index_, g_mock_cuda.last_h2d_index_); + EXPECT_TRUE(g_mock_cuda.is_device_ptr(tensor->const_data_ptr())); + const float* data = tensor->const_data_ptr(); + for (int i = 0; i < 4; ++i) { + EXPECT_EQ(data[i], static_cast(i)); + } +} + +} // namespace + +TEST_F(TensorParserDeviceTest, ConstantInProgramCopiedToDeviceMemory) { + expectConstantCopiedToDevice( + "ET_MODULE_ADD_CONSTANT_WITH_DEVICE_PATH", /*data_env=*/nullptr); +} + +TEST_F(TensorParserDeviceTest, ConstantInDataFileCopiedToDeviceMemory) { + expectConstantCopiedToDevice( + "ET_MODULE_ADD_CONSTANT_WITH_DEVICE_PROGRAM_PATH", + "ET_MODULE_ADD_CONSTANT_WITH_DEVICE_DATA_PATH"); +} diff --git a/test/models/export_program_with_device_info.py b/test/models/export_program_with_device_info.py index 9e895205935..87fb16e0b7e 100644 --- a/test/models/export_program_with_device_info.py +++ b/test/models/export_program_with_device_info.py @@ -6,10 +6,13 @@ # pyre-unsafe -"""Exports a simple model with device-annotated tensors for C++ testing. +"""Exports simple models with device-annotated tensors for C++ testing. Uses DeviceAwarePartitioner (BackendWithCompilerDemo + target_device=cuda:0) so that delegate output tensors are annotated with CUDA device in the .pte. +ModuleAddConstantWithDevice is exported with skip_h2d_for_constants and +target_device=cuda:1, so its constant is planned in cuda:1 memory with its data +as the initial state, once in the program and once in a separate .ptd file. """ import argparse @@ -18,6 +21,8 @@ import torch from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge from executorch.exir.backend.test.device_util import DeviceAwarePartitioner +from executorch.exir.passes import MemoryPlanningPass +from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig from torch import nn from torch.export import export @@ -32,6 +37,49 @@ def get_random_inputs(self): return (torch.randn(2, 2), torch.randn(2, 2)) +class ModuleAddConstantWithDevice(nn.Module): + """Adds a constant buffer to the input in a delegate on CUDA.""" + + offset: torch.Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("offset", torch.arange(4.0).reshape(2, 2)) + + def forward(self, a: torch.Tensor) -> torch.Tensor: + return torch.add(a, self.offset) + + def get_random_inputs(self): + return (torch.randn(2, 2),) + + +def export_constant_on_device(outdir: str, external_constants: bool) -> None: + model = ModuleAddConstantWithDevice() + lowered = to_edge( + export(model, model.get_random_inputs()), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).to_backend(DeviceAwarePartitioner("cuda:1")) + et_prog = lowered.to_executorch( + ExecutorchBackendConfig( # type: ignore[call-arg] + emit_stacktrace=False, + enable_non_cpu_memory_planning=True, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True, skip_h2d_for_constants=True + ), + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + external_constants=external_constants, + ) + ) + name = "ModuleAddConstantWithDevice" + ("Program" if external_constants else "") + with open(os.path.join(outdir, f"{name}.pte"), "wb") as fp: + fp.write(et_prog.buffer) + if external_constants: + et_prog._tensor_data[name] = et_prog._tensor_data.pop( + "_default_external_constant" + ) + et_prog.write_tensor_data_to_file(outdir) + + def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--outdir", type=str, required=True) @@ -60,6 +108,9 @@ def main() -> None: fp.write(et_prog.buffer) print(f"Exported ModuleAddWithDevice to {outfile}") + for external_constants in (False, True): + export_constant_on_device(args.outdir, external_constants) + if __name__ == "__main__": main() diff --git a/test/models/targets.bzl b/test/models/targets.bzl index bd4df80a50a..dd1e2bd0c7f 100644 --- a/test/models/targets.bzl +++ b/test/models/targets.bzl @@ -226,6 +226,9 @@ def define_common_targets(): cmd = "$(exe :export_program_with_device_info) --outdir $OUT", outs = { "ModuleAddWithDevice.pte": ["ModuleAddWithDevice.pte"], + "ModuleAddConstantWithDevice.pte": ["ModuleAddConstantWithDevice.pte"], + "ModuleAddConstantWithDeviceProgram.pte": ["ModuleAddConstantWithDeviceProgram.pte"], + "ModuleAddConstantWithDeviceProgram.ptd": ["ModuleAddConstantWithDeviceProgram.ptd"], }, default_outs = ["."], visibility = [ From 0cf1823d152ce66aa352cad607943263b1d79264 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 26 Oct 2026 12:16:01 -0700 Subject: [PATCH 4/4] Point the constant copy warning at skip_h2d_for_constants The warning for a constant copied to the device on every call only said to move it into the delegate. Without skip_h2d_for_constants, it now also suggests the option when only delegates on that device read the constant, unless the constant is empty and costs nothing to copy. With the option, it names the other reader that keeps the constant off the device. --- exir/passes/propagate_device_pass.py | 62 +++++++++++++++++++++--- exir/tests/test_propagate_device_pass.py | 54 ++++++++++++++++++--- 2 files changed, 103 insertions(+), 13 deletions(-) diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py index 235ab7ddbd1..f0842ef788a 100644 --- a/exir/passes/propagate_device_pass.py +++ b/exir/passes/propagate_device_pass.py @@ -364,6 +364,32 @@ def _delegate_device( return None return _get_target_device_from_compile_specs(lowered_module) + def _constant_reader_off_device( + self, + graph_module: torch.fx.GraphModule, + constant: torch.fx.Node, + device_type: schema.DeviceType, + device_index: int, + ) -> Optional[torch.fx.Node]: + """A reader of *constant* other than a delegate on this device, or None. + + A delegate on another device reads it through its _h2d_copy, so the + reader behind such a copy counts. + """ + for user in constant.users: + readers = ( + user.users + if user.target == torch.ops.et_copy._h2d_copy.default + else [user] + ) + for reader in readers: + if self._delegate_device(graph_module, reader) != ( + device_type, + device_index, + ): + return reader + return None + def _report_per_call_copy(self, message: str, skip_requested: bool) -> None: """Warn about a copy that runs on every call. When a skip flag asked for no such copy and error_on_per_call_copies is set, raise instead.""" @@ -422,15 +448,18 @@ def _insert_h2d_copies( f"it is created, for example in the pass that adds it." ) + other_reader = ( + self._constant_reader_off_device( + graph_module, arg, target_device_type, device_index + ) + if arg_spec.const + else None + ) if ( self.skip_h2d_for_constants and arg_spec.const and self._is_placeholder(arg) - and all( - self._delegate_device(graph_module, user) - == (target_device_type, device_index) - for user in arg.users - ) + and other_reader is None ): # Memory planning keeps it in device memory for the whole # method, and the runtime loads its data there once. @@ -442,12 +471,31 @@ def _insert_h2d_copies( h2d_node = _reusable_copy(arg, node, target_device_type, device_index) if h2d_node is None: if arg_spec.const: + remedy = ( + "Have the partitioner take it into the delegate to avoid " + "that copy." + ) + # The option saves nothing for an empty constant. + if ( + not self.skip_h2d_for_constants + and other_reader is None + and arg_spec.nbytes() > 0 + ): + remedy = ( + "Have the partitioner take it into the delegate, or set " + "PropagateDeviceConfig(skip_h2d_for_constants=True), to " + "avoid that copy." + ) + elif self.skip_h2d_for_constants and other_reader is not None: + remedy = ( + f"skip_h2d_for_constants cannot keep it on the device, " + f"because '{other_reader.name}' also reads it. {remedy}" + ) self._report_per_call_copy( f"Constant '{arg.name}' ({arg_spec.nbytes()} bytes) feeds a " f"delegate on {target_device_type.name.lower()}:" f"{device_index} from outside it, so it is copied to the " - f"device on every call. Have the partitioner take it into " - f"the delegate to avoid that copy.", + f"device on every call. {remedy}", self.skip_h2d_for_method_inputs or self.skip_h2d_for_constants, ) with graph_module.graph.inserting_before(node): diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py index 26dcb04032e..f7d587822c2 100644 --- a/exir/tests/test_propagate_device_pass.py +++ b/exir/tests/test_propagate_device_pass.py @@ -2005,7 +2005,8 @@ def forward(self, a): def test_skip_h2d_for_constants_copies_constants_read_elsewhere(self): """A constant that a host op, or a delegate on another device, reads too is still copied with skip_h2d_for_constants=True, and so is every - constant without the option.""" + constant without the option. The warning says why the option did not + keep it, or suggests the option when it is off.""" class Buffer(torch.nn.Module): def __init__(self): @@ -2023,12 +2024,49 @@ class ReadByTwoDelegates(Buffer): def forward(self, a): return torch.add(a, self.offset), torch.add(a * 2, self.offset) + class EmptyBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("offset", torch.zeros(2, 0)) + + def forward(self, a): + return torch.add(a[:, :0], self.offset) + + kept_elsewhere = "skip_h2d_for_constants cannot keep it on the device" cases = { - "host reader": (AlsoReadOnHost(), ["cuda:0"], True, [0]), - "two devices": (ReadByTwoDelegates(), ["cuda:0", "cuda:1"], True, [0, 1]), - "without the option": (Buffer(), ["cuda:0"], False, [0]), + "host reader": (AlsoReadOnHost(), ["cuda:0"], True, [0], kept_elsewhere), + "two devices": ( + ReadByTwoDelegates(), + ["cuda:0", "cuda:1"], + True, + [0, 1], + kept_elsewhere, + ), + "without the option": ( + Buffer(), + ["cuda:0"], + False, + [0], + "or set PropagateDeviceConfig(skip_h2d_for_constants=True)", + ), + # The option would not help, so it is not suggested. + "host reader without the option": ( + AlsoReadOnHost(), + ["cuda:0"], + False, + [0], + "take it into the delegate to avoid that copy", + ), + # Nor for an empty constant, which costs nothing to copy. + "empty without the option": ( + EmptyBuffer(), + ["cuda:0"], + False, + [0], + "take it into the delegate to avoid that copy", + ), } - for case, (model, devices, skip_constants, copied_to) in cases.items(): + for case, (model, devices, skip_constants, copied_to, hint) in cases.items(): with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs: results = self._get_executorch_program( model, @@ -2036,7 +2074,11 @@ def forward(self, a): PerAddDeviceDelegatePartitioner(devices), _skip_h2d_for_constants_config(False, skip_constants), ) - self.assertIn("Constant 'b_offset'", logs.output[0]) + warnings = [w for w in logs.output if "Constant 'b_offset'" in w] + self.assertTrue(warnings) + for warning in warnings: + self.assertIn(hint, warning) + self.assertNotIn("_h2d_copy", warning) for pipeline, _program, gm in results: with self.subTest(pipeline=pipeline, case=case): copies = [