diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index b490e28c95f..25684f8b32e 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -26,6 +26,8 @@ on: - .ci/scripts/export_model_artifact.sh - .ci/scripts/test_model_e2e.sh - examples/models/muse-glimmer/** + - extension/pybindings/** + - runtime/__init__.py workflow_dispatch: concurrency: @@ -160,6 +162,8 @@ jobs: unittest-cuda: name: unittest-cuda needs: [changed-files, run-decision] + # This job runs the pybindings device test, so it also has to fire on the + # two pybindings paths the workflow filter lists. if: | contains(needs.changed-files.outputs.changed-files, 'backends/cuda') || contains(needs.changed-files.outputs.changed-files, 'backends/aoti') || @@ -167,6 +171,8 @@ jobs: contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test-cuda-build.sh') || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/export_model_artifact.sh') || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e.sh') || + contains(needs.changed-files.outputs.changed-files, 'extension/pybindings') || + contains(needs.changed-files.outputs.changed-files, 'runtime/__init__.py') || needs.run-decision.outputs.is-full-run == 'true' uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: @@ -190,6 +196,29 @@ jobs: conda install -y -c conda-forge 'libstdcxx-ng>=12' export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH + # The pybindings loader only asks for device memory on a build that has + # a device allocator registered, so this test cannot run in the CPU + # jobs. It needs the install above and nothing built below, so run it + # here and fail before the long builds. + # + # unittest exits 0 when a test skips, and this test skips when the CUDA + # backend is missing or no GPU is visible. Both are the prerequisites + # this job exists to provide, so treat a skip as a failure rather than + # let a packaging regression pass silently. `tee` would otherwise hide + # the real exit status behind its own, so read it back explicitly. + set +e + python -m unittest -v executorch.extension.pybindings.test.test_pybindings.PybindingsTest.test_device_planned_method_allocates_on_the_device 2>&1 | tee /tmp/pybindings_device_test.log + pybindings_device_status=${PIPESTATUS[0]} + set -e + if [ "${pybindings_device_status}" -ne 0 ]; then + echo "::error::the pybindings device test failed" + exit "${pybindings_device_status}" + fi + if grep -qE "^OK \(skipped=|skipped=[1-9]" /tmp/pybindings_device_test.log; then + echo "::error::the pybindings device test skipped in the CUDA job, which means the CUDA backend or a visible GPU is missing" + exit 1 + fi + # Build ExecuTorch with CUDA support cmake --workflow --preset llm-release-cuda diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 6905c0553f9..27255b3b8ac 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -7,6 +7,7 @@ */ #include +#include #include #include #include @@ -31,6 +32,7 @@ #include #include #include +#include #include #include #include @@ -81,6 +83,7 @@ using ::executorch::ET_RUNTIME_NAMESPACE::get_num_registered_backends; using ::executorch::ET_RUNTIME_NAMESPACE::get_registered_kernels; using ::executorch::ET_RUNTIME_NAMESPACE::Kernel; using ::executorch::ET_RUNTIME_NAMESPACE::Method; +using ::executorch::ET_RUNTIME_NAMESPACE::MethodMeta; using ::executorch::ET_RUNTIME_NAMESPACE::Program; using ::executorch::extension::BufferDataLoader; using ::executorch::extension::MallocMemoryAllocator; @@ -88,6 +91,7 @@ using ::executorch::extension::MmapDataLoader; using ::executorch::extension::ET_BUNDLED_MODULE_NAMESPACE::BundledModule; using ::executorch::extension::pybindings::PyDataLoader; using ::executorch::runtime::DataLoader; +using ::executorch::runtime::DeviceMemoryBuffer; using ::executorch::runtime::Error; using ::executorch::runtime::EValue; using ::executorch::runtime::EventTracerDebugLogLevel; @@ -1077,18 +1081,33 @@ inline std::shared_ptr load_program( /// A wrapper/util class for executorch memory allocations/manager. class ProgramMemory { public: - explicit ProgramMemory(std::vector>&& non_const_buffers) + /// `devices` is empty when every buffer is on the host, which keeps + /// `MemoryManager::has_device_memory()` false for CPU-only programs. + /// Otherwise it holds one entry per buffer, indexed like `sizes`. + /// + /// Members initialize in declaration order and each one reads the members + /// declared before it, so that order is load-bearing. Device buffers come + /// first so that a device that is missing or out of memory throws before the + /// host arenas are allocated and zero-filled, rather than after. + ProgramMemory( + std::vector&& sizes, + std::vector&& devices) : runtime_allocator_(), - non_const_buffers_(std::move(non_const_buffers)), + planned_sizes_(std::move(sizes)), + planned_devices_(std::move(devices)), + device_buffers_(allocate_device_buffers()), + non_const_buffers_(allocate_host_buffers()), non_const_spans_(create_non_const_spans()), - non_const_allocator_( - {non_const_spans_.data(), non_const_spans_.size()}), + non_const_allocator_(create_non_const_allocator()), mem_manager_( &const_allocator_, &non_const_allocator_, &runtime_allocator_, &temp_allocator_) {} + explicit ProgramMemory(std::vector&& sizes) + : ProgramMemory(std::move(sizes), {}) {} + /// Returns a pointer to the internal memory manager, the Memory instance /// must outlive this pointer. MemoryManager* mem_manager() { @@ -1105,6 +1124,16 @@ class ProgramMemory { MallocMemoryAllocator temp_allocator_{}; + std::vector planned_sizes_; + + std::vector planned_devices_; + + // Backs device-tagged buffers; the entry is empty for a CPU-tagged buffer. + // Parallel to non_const_buffers_ so both index by planned buffer id. Empty + // for an all-host program. + std::vector device_buffers_; + + // Backs CPU-tagged buffers; the entry is empty for a device-tagged buffer. std::vector> non_const_buffers_; std::vector> non_const_spans_; @@ -1113,16 +1142,125 @@ class ProgramMemory { MemoryManager mem_manager_; + bool is_device_buffer(size_t index) const { + return index < planned_devices_.size() && !planned_devices_[index].is_cpu(); + } + + std::vector> allocate_host_buffers() { + std::vector> result; + result.reserve(planned_sizes_.size()); + for (size_t i = 0; i < planned_sizes_.size(); ++i) { + if (is_device_buffer(i)) { + result.emplace_back(); + } else { + result.emplace_back(planned_sizes_[i]); + } + } + return result; + } + + std::vector allocate_device_buffers() { + std::vector result; + if (planned_devices_.empty()) { + return result; + } + // Both vectors are filled in lockstep today, so this only fires if a + // future caller breaks that. HierarchicalAllocator aborts on a mismatch, + // so check here instead, where a Python caller can catch it. + THROW_IF_ERROR( + planned_devices_.size() == planned_sizes_.size() + ? Error::Ok + : Error::InvalidArgument, + "Have %zu planned buffer sizes but %zu device tags", + planned_sizes_.size(), + planned_devices_.size()); + result.reserve(planned_sizes_.size()); + for (size_t i = 0; i < planned_sizes_.size(); ++i) { + if (!is_device_buffer(i)) { + result.emplace_back(); + continue; + } + auto buffer = DeviceMemoryBuffer::create( + planned_sizes_[i], + planned_devices_[i].type(), + planned_devices_[i].index()); + THROW_IF_ERROR( + buffer.error(), + "Failed to allocate %" PRId64 " bytes for buffer %zu on device %d:%d", + planned_sizes_[i], + i, + static_cast(planned_devices_[i].type()), + static_cast(planned_devices_[i].index())); + result.emplace_back(std::move(buffer.get())); + } + return result; + } + std::vector> create_non_const_spans() { std::vector> result; - for (size_t i = 0; i < non_const_buffers_.size(); i++) { - result.push_back( - {non_const_buffers_[i].data(), non_const_buffers_[i].size()}); + result.reserve(planned_sizes_.size()); + for (size_t i = 0; i < planned_sizes_.size(); ++i) { + if (is_device_buffer(i)) { + result.push_back(device_buffers_[i].as_span()); + } else { + result.push_back( + {non_const_buffers_[i].data(), non_const_buffers_[i].size()}); + } } return result; } + + HierarchicalAllocator create_non_const_allocator() { + Span> buffers( + non_const_spans_.data(), non_const_spans_.size()); + return planned_devices_.empty() + ? HierarchicalAllocator(buffers) + : HierarchicalAllocator( + buffers, {planned_devices_.data(), planned_devices_.size()}); + } }; +/// True if any of the method's memory-planned buffers must live off the host. +bool has_device_buffers(const MethodMeta& method_meta) { + for (size_t i = 0; i < method_meta.num_memory_planned_buffers(); ++i) { + auto device = method_meta.memory_planned_buffer_device(i); + THROW_IF_ERROR( + device.error(), "Failed to get device of planned buffer %zu", i); + if (!device.get().is_cpu()) { + return true; + } + } + return false; +} + +/// Arenas sized and placed for a single method, used when that method's +/// buffers cannot come from the program-wide host arenas. Returns nullptr when +/// every buffer is on the host, so one pass over the metadata answers both +/// whether the method needs its own arenas and how big they are. +std::shared_ptr make_method_memory( + const MethodMeta& method_meta) { + const size_t num_buffers = method_meta.num_memory_planned_buffers(); + std::vector sizes; + std::vector devices; + sizes.reserve(num_buffers); + devices.reserve(num_buffers); + bool needs_device_memory = false; + for (size_t i = 0; i < num_buffers; ++i) { + auto size = method_meta.memory_planned_buffer_size(i); + THROW_IF_ERROR(size.error(), "Failed to get size of planned buffer %zu", i); + auto device = method_meta.memory_planned_buffer_device(i); + THROW_IF_ERROR( + device.error(), "Failed to get device of planned buffer %zu", i); + needs_device_memory |= !device.get().is_cpu(); + sizes.push_back(size.get()); + devices.push_back(device.get()); + } + if (!needs_device_memory) { + return nullptr; + } + return std::make_shared(std::move(sizes), std::move(devices)); +} + struct PyMethod final { explicit PyMethod( std::shared_ptr memory, @@ -1423,8 +1561,17 @@ struct PyProgram final { for (size_t i = 0; i < state_->program_->num_methods(); ++i) { auto name = state_->program_->get_method_name(i).get(); auto method_meta = state_->program_->method_meta(name).get(); - for (size_t j = 0; j < method_meta.num_non_const_buffers(); j++) { - int64_t buffer_size = method_meta.non_const_buffer_size(j).get(); + // A device-planned method gets its own arenas in load_method and never + // reads these, so letting its sizes in would only grow the host arenas + // the other methods share. + if (has_device_buffers(method_meta)) { + continue; + } + for (size_t j = 0; j < method_meta.num_memory_planned_buffers(); ++j) { + auto size = method_meta.memory_planned_buffer_size(j); + THROW_IF_ERROR( + size.error(), "Failed to get size of planned buffer %zu", j); + int64_t buffer_size = size.get(); if (non_const_buffer_sizes.find(j) == non_const_buffer_sizes.end()) { non_const_buffer_sizes.insert({j, buffer_size}); } else { @@ -1434,16 +1581,14 @@ struct PyProgram final { } } - // Allocate the arenas. Using vector because we need to remember the size as - // well, so vector is easier then unique_ptr. - std::vector> non_const_buffers_; - for (std::map::iterator i = non_const_buffer_sizes.begin(); - i != non_const_buffer_sizes.end(); - i++) { - non_const_buffers_.push_back(std::vector(i->second)); + // Allocate the shared host arenas. + std::vector planned_sizes; + planned_sizes.reserve(non_const_buffer_sizes.size()); + for (const auto& entry : non_const_buffer_sizes) { + planned_sizes.push_back(entry.second); } - memory_ = std::make_shared(std::move(non_const_buffers_)); + memory_ = std::make_shared(std::move(planned_sizes)); if (event_tracer_ && debug_buffer_size > 0) { // If a debug buffer was requested for the ETDump, allocate it and make // sure its lifetime is as long as the event_tracer. @@ -1508,9 +1653,22 @@ struct PyProgram final { } std::unique_ptr load_method(const std::string& method_name) { + Result meta = + state_->program_->method_meta(method_name.c_str()); + THROW_IF_ERROR( + meta.error(), + "Failed to get method meta for method %s, error: 0x:%" PRIx32, + method_name.c_str(), + static_cast(meta.error())); + // Device memory is claimed here rather than at program load so that one + // accelerator method cannot make the rest of the program unloadable. A + // host-only method keeps sharing the program-wide arenas, so its planned + // memory is not isolated from the other host-only methods of this program. + auto method_memory = make_method_memory(meta.get()); + auto memory = method_memory ? std::move(method_memory) : memory_; Result res = state_->program_->load_method( method_name.c_str(), - memory_->mem_manager(), + memory->mem_manager(), event_tracer_.get(), state_->data_map_.get()); THROW_IF_ERROR( @@ -1519,7 +1677,9 @@ struct PyProgram final { method_name.c_str(), static_cast(res.error())); return std::make_unique( - memory_, state_, std::make_unique(std::move(res.get()))); + std::move(memory), + state_, + std::make_unique(std::move(res.get()))); } Span get_etdump_debug_buffer() { diff --git a/extension/pybindings/test/BUCK b/extension/pybindings/test/BUCK index a18abac2382..e90e837732d 100644 --- a/extension/pybindings/test/BUCK +++ b/extension/pybindings/test/BUCK @@ -26,6 +26,7 @@ fbcode_target( "//executorch/exir:pass_manager", "//executorch/exir:scalar_type", "//executorch/exir/_serialize:lib", + "//executorch/exir/backend:partitioner", "//executorch/exir/emit:lib", "//executorch/exir/passes:lib", "//executorch/runtime/core:core", @@ -39,6 +40,8 @@ fbcode_target( preload_deps = ["//executorch/kernels/quantized:aot_lib"], deps = [ ":make_test", + "//executorch/exir:schema", + "//executorch/exir/backend/test:device_util", "//executorch/extension/pybindings:portable_lib", ], ) @@ -50,6 +53,8 @@ fbcode_target( preload_deps = ["//executorch/kernels/quantized:aot_lib"], deps = [ ":make_test", + "//executorch/exir:schema", + "//executorch/exir/backend/test:device_util", "//executorch/extension/pybindings:aten_lib", "//executorch/kernels/quantized:aot_lib", ], @@ -61,6 +66,8 @@ fbcode_target( srcs = ["test_pybindings.py"], deps = [ ":make_test", + "//executorch/exir:schema", + "//executorch/exir/backend/test:device_util", ], ) diff --git a/extension/pybindings/test/make_test.py b/extension/pybindings/test/make_test.py index a1bf4b980e0..1e47dd9b6c6 100644 --- a/extension/pybindings/test/make_test.py +++ b/extension/pybindings/test/make_test.py @@ -6,10 +6,11 @@ # pyre-unsafe -from typing import Any, Optional, Tuple +from typing import Any, Dict, Optional, Tuple import torch from executorch.exir import ExecutorchBackendConfig, ExecutorchProgramManager, to_edge +from executorch.exir.backend.partitioner import Partitioner from torch.export import export @@ -151,6 +152,7 @@ def get_inputs(self): def create_program( eager_module: torch.nn.Module, et_config: Optional[ExecutorchBackendConfig] = None, + partitioner: Optional[Dict[str, Partitioner]] = None, ) -> Tuple[ExecutorchProgramManager, Tuple[Any, ...]]: """Returns an executorch program based on ModuleAdd, along with inputs.""" @@ -179,7 +181,10 @@ def forward(self, *args, **kwargs): wrapped_mod = WrapperModule(getattr(eager_module, method_name)) exported_methods[method_name] = export(wrapped_mod, method_input, strict=True) - exec_prog = to_edge(exported_methods).to_executorch(config=et_config) + edge_prog = to_edge(exported_methods) + if partitioner is not None: + edge_prog = edge_prog.to_backend(partitioner) + exec_prog = edge_prog.to_executorch(config=et_config) # Create the ExecuTorch program from the graph. exec_prog.dump_executorch_program(verbose=True) diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 19b3620cef3..7e0d3dfce25 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -6,14 +6,18 @@ # pyre-unsafe +import os import sys +import tempfile import unittest from io import StringIO import torch from executorch.exir import ExecutorchBackendConfig, to_edge +from executorch.exir.backend.test.device_util import DeviceAwarePartitioner from executorch.exir.passes import MemoryPlanningPass +from executorch.exir.schema import DeviceType from executorch.extension.pybindings.test.make_test import ( create_program, ModuleAdd, @@ -786,3 +790,153 @@ def test_program_method_rejects_input_on_unrepresentable_device(self): message = str(caught.exception) self.assertIn("is on device meta", message) self.assertIn("only CPU and CUDA tensors", message) + + def test_program_loads_when_one_method_is_device_planned(self): + # Linking the CUDA backend registers a CUDA allocator at static init, and the + # registry has no way to drop one, so there the device load would succeed with + # or without the fix and this test could not tell them apart. + if "CudaBackend" in self.runtime._get_registered_backend_names(): + self.skipTest("a registered CUDA allocator satisfies the device load") + + exported_program, inputs = create_program( + ModuleMulti(), + et_config=ExecutorchBackendConfig(enable_non_cpu_memory_planning=True), + partitioner={"forward2": DeviceAwarePartitioner()}, + ) + + # Without this the test would quietly become a second copy of + # test_method_multiple_entry if the planner stopped tagging devices. + planned_devices = { + plan.name: [ + buffer.device_type for buffer in (plan.non_const_buffer_device or []) + ] + for plan in exported_program.executorch_program.execution_plan + } + self.assertIn(DeviceType.CUDA, planned_devices["forward2"]) + self.assertNotIn(DeviceType.CUDA, planned_devices["forward"]) + + program = self.load_prog_fn(exported_program.buffer) + self.assertEqual(program.num_methods(), 2) + + method = program.load_method("forward") + self.assertTrue(torch.allclose(method.call(inputs)[0], torch.ones(2, 2) * 2)) + + # Asking for device memory at all is what this asserts. Before the fix the + # loader put every planned buffer in host memory, so the device-planned method + # never asked and simply loaded. Now it asks, and with no device allocator + # registered the request is refused. + with self.assertRaises(RuntimeError) as caught: + program.load_method("forward2") + self.assertIn("on device", str(caught.exception)) + + # A failed load must leave the method that already loaded usable. + self.assertTrue(torch.allclose(method.call(inputs)[0], torch.ones(2, 2) * 2)) + + def test_device_planned_method_allocates_on_the_device(self): + # The other device test covers the refusal. This one covers what the + # refusal is protecting: on a build that does have a device allocator, + # the arena has to come off the device, not out of host memory. It needs + # a real accelerator, so it only runs where one is present. + if "CudaBackend" not in self.runtime._get_registered_backend_names(): + self.skipTest("needs a build with the CUDA backend linked in") + if not torch.cuda.is_available(): + self.skipTest("needs a visible CUDA device") + + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from executorch.exir import to_edge_transform_and_lower + from executorch.exir.backend.compile_spec_schema import CompileSpec + + # Large enough that the arena is far bigger than the noise in + # mem_get_info, which moves by a few MiB as contexts are created. + side = 2048 + inputs = (torch.ones(side, side), torch.ones(side, side)) + + class HostOnly(torch.nn.Module): + def forward(self, x, y): + return x + y + + class Delegated(torch.nn.Module): + def forward(self, x, y): + return (x + y) * 2.0 + + edge = to_edge_transform_and_lower( + { + "forward": export(HostOnly(), inputs, strict=True), + "forward2": export(Delegated(), inputs, strict=True), + }, + partitioner={ + "forward2": [CudaPartitioner([CompileSpec("method_name", b"forward2")])] + }, + ) + exported_program = edge.to_executorch( + config=ExecutorchBackendConfig(enable_non_cpu_memory_planning=True) + ) + + plans = { + plan.name: plan + for plan in exported_program.executorch_program.execution_plan + } + device_buffers = { + buffer.buffer_idx + for buffer in (plans["forward2"].non_const_buffer_device or []) + } + self.assertTrue(device_buffers, "the planner tagged nothing for the device") + self.assertFalse(plans["forward"].non_const_buffer_device or []) + device_bytes = sum( + size + for index, size in enumerate(plans["forward2"].non_const_buffer_sizes) + if index in device_buffers + ) + + # The CUDA backend keeps its weights in a separate file, so the program + # has to be loaded from disk with that file alongside it. + with tempfile.TemporaryDirectory() as directory: + pte_path = os.path.join(directory, "program.pte") + with open(pte_path, "wb") as pte_file: + exported_program.write_to_file(pte_file) + data_names = sorted(exported_program._tensor_data or {}) + exported_program.write_tensor_data_to_file(directory) + data_path = ( + os.path.join(directory, data_names[0] + ".ptd") if data_names else None + ) + + torch.cuda.init() + program = self.runtime._load_program(pte_path, data_path=data_path) + + torch.cuda.synchronize() + free_before, _ = torch.cuda.mem_get_info() + + # The host-only method must not touch the device at all. + host_method = program.load_method("forward") + torch.cuda.synchronize() + free_after_host, _ = torch.cuda.mem_get_info() + self.assertLess(free_before - free_after_host, device_bytes // 2) + + device_method = program.load_method("forward2") + torch.cuda.synchronize() + free_after_device, _ = torch.cuda.mem_get_info() + self.assertGreaterEqual( + free_after_host - free_after_device, device_bytes * 0.9 + ) + + # Before the fix this ran on a host pointer and the backend rejected + # it, so getting the right answer back is itself part of the check. + expected = (inputs[0] + inputs[1]) * 2.0 + self.assertTrue( + torch.allclose(device_method.call(inputs)[0].cpu(), expected) + ) + self.assertTrue( + torch.allclose(host_method.call(inputs)[0].cpu(), inputs[0] + inputs[1]) + ) + # The device arenas are private, so running the host method in + # between must not disturb them. + self.assertTrue( + torch.allclose(device_method.call(inputs)[0].cpu(), expected) + ) + + del device_method + del host_method + del program + torch.cuda.synchronize() + free_at_end, _ = torch.cuda.mem_get_info() + self.assertGreaterEqual(free_at_end - free_after_device, device_bytes * 0.9) diff --git a/runtime/__init__.py b/runtime/__init__.py index 011a15163cb..ba83ff54856 100644 --- a/runtime/__init__.py +++ b/runtime/__init__.py @@ -183,6 +183,27 @@ def method_names(self) -> Set[str]: def load_method(self, name: str) -> Optional[Method]: """Loads a method from the program. + Memory-planned buffers are allocated differently depending on where the + program placed them. A method whose buffers are all on the host shares + one set of arenas with every other host-only method of this program, + sized to the largest of them, so two such methods overwrite each + other's intermediate values. Outputs are copied out on every call, so + that sharing does not change what execute returns. A method with any + buffer placed on an accelerator gets its own arenas instead, claimed + when that method is first loaded and owned by that method until the + method itself is released. This program caches every method it loads, + so in normal use those arenas live as long as the program does. A + missing or exhausted accelerator therefore fails that one load rather + than the whole program, but whatever it does claim stays claimed for + every method loaded after it. + + A program exported with ``share_mutable_buffers=True`` relies on every + method reading its shared mutable state from one allocation. A method + with accelerator-placed buffers gets its own arenas instead, so it also + gets its own copy of that state and will not observe writes made through + another method. Export without ``share_mutable_buffers`` if a method + needs both accelerator memory and shared state. + Args: name: The name of the method to load. diff --git a/shim_et/xplat/executorch/extension/pybindings/pybindings.bzl b/shim_et/xplat/executorch/extension/pybindings/pybindings.bzl index 98ab3c7cf94..e6b0319ca97 100644 --- a/shim_et/xplat/executorch/extension/pybindings/pybindings.bzl +++ b/shim_et/xplat/executorch/extension/pybindings/pybindings.bzl @@ -10,6 +10,7 @@ MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB = [ PORTABLE_MODULE_DEPS = [ "//executorch/runtime/kernel:operator_registry", "//executorch/runtime/executor:program", + "//executorch/runtime/core:device_memory_buffer", "//executorch/devtools/bundled_program/schema:bundled_program_schema_fbs", "//executorch/extension/aten_util:aten_bridge", "//executorch/devtools/bundled_program:runtime", @@ -27,6 +28,7 @@ PORTABLE_MODULE_DEPS = [ ATEN_MODULE_DEPS = [ "//executorch/runtime/kernel:operator_registry_aten", "//executorch/runtime/executor:program_aten", + "//executorch/runtime/core:device_memory_buffer", "//executorch/runtime/core/exec_aten:lib_aten", "//executorch/devtools/bundled_program/schema:bundled_program_schema_fbs", "//executorch/extension/data_loader:buffer_data_loader",