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/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/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/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..f0842ef788a 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=" @@ -353,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.""" @@ -371,10 +408,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,16 +448,55 @@ 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 other_reader is None + ): + # 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: + 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.", - self.skip_h2d_for_method_inputs, + 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): 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..f7d587822c2 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,199 @@ 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. 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): + 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) + + 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], 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, hint) 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), + ) + 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 = [ + 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/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; 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/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/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) { 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 = [