|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +import tempfile |
| 8 | +import unittest |
| 9 | +from pathlib import Path |
| 10 | + |
| 11 | +import torch |
| 12 | +from executorch.backends.cuda.cuda_backend import CudaBackend |
| 13 | +from executorch.exir import ( |
| 14 | + EdgeCompileConfig, |
| 15 | + ExecutorchBackendConfig, |
| 16 | + to_edge_transform_and_lower, |
| 17 | +) |
| 18 | +from executorch.exir.backend.compile_spec_schema import CompileSpec |
| 19 | +from executorch.exir.backend.partitioner import ( |
| 20 | + DelegationSpec, |
| 21 | + Partitioner, |
| 22 | + PartitionResult, |
| 23 | +) |
| 24 | +from executorch.exir.dialects._ops import ops as exir_ops |
| 25 | +from executorch.exir.passes import MemoryPlanningPass |
| 26 | +from executorch.exir.passes.propagate_device_pass import ( |
| 27 | + PropagateDeviceConfig, |
| 28 | + TARGET_DEVICE_COMPILE_SPEC_KEY, |
| 29 | +) |
| 30 | + |
| 31 | + |
| 32 | +class EachAddOnCuda(Partitioner): |
| 33 | + """Lowers each add to its own CUDA delegate. Parameters stay outside, so |
| 34 | + they cross the delegate boundary as in a partial lowering.""" |
| 35 | + |
| 36 | + def partition(self, exported_program) -> PartitionResult: |
| 37 | + spec = DelegationSpec( |
| 38 | + CudaBackend.__name__, |
| 39 | + [ |
| 40 | + CudaBackend.generate_method_name_compile_spec("forward"), |
| 41 | + CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, b"cuda:0"), |
| 42 | + ], |
| 43 | + ) |
| 44 | + tags = {} |
| 45 | + adds = [ |
| 46 | + node |
| 47 | + for node in exported_program.graph_module.graph.nodes |
| 48 | + if node.op == "call_function" |
| 49 | + and node.target == exir_ops.edge.aten.add.Tensor |
| 50 | + ] |
| 51 | + for i, node in enumerate(adds): |
| 52 | + node.meta["delegation_tag"] = f"add{i}" |
| 53 | + tags[f"add{i}"] = spec |
| 54 | + return PartitionResult( |
| 55 | + tagged_exported_program=exported_program, partition_tags=tags |
| 56 | + ) |
| 57 | + |
| 58 | + |
| 59 | +def load(model, inputs, directory, config=None): |
| 60 | + """Lower model with EachAddOnCuda and load its forward method from a file |
| 61 | + in directory.""" |
| 62 | + from executorch.runtime import Runtime |
| 63 | + |
| 64 | + edge = to_edge_transform_and_lower( |
| 65 | + torch.export.export(model, inputs, strict=True), |
| 66 | + partitioner=[EachAddOnCuda()], |
| 67 | + compile_config=EdgeCompileConfig( |
| 68 | + _check_ir_validity=False, _skip_dim_order=True |
| 69 | + ), |
| 70 | + ) |
| 71 | + path = Path(directory) / "model.pte" |
| 72 | + path.write_bytes(edge.to_executorch(config or ExecutorchBackendConfig()).buffer) |
| 73 | + return Runtime.get().load_program(path).load_method("forward") |
| 74 | + |
| 75 | + |
| 76 | +def run(model, inputs) -> list[torch.Tensor]: |
| 77 | + """Lower model with EachAddOnCuda, run it in the runtime, return its outputs.""" |
| 78 | + with tempfile.TemporaryDirectory() as directory: |
| 79 | + method = load(model, inputs, directory) |
| 80 | + return [output.cpu() for output in method.execute(list(inputs))] |
| 81 | + |
| 82 | + |
| 83 | +@unittest.skipUnless(torch.cuda.is_available(), "CUDA required") |
| 84 | +@unittest.skipIf( |
| 85 | + torch.version.hip is not None, |
| 86 | + "loads through the Python runtime, which the ROCm build does not include", |
| 87 | +) |
| 88 | +class CudaDevicePlacementTest(unittest.TestCase): |
| 89 | + def test_transposed_parameter_read_by_a_delegate(self): |
| 90 | + """The program stores a transposed parameter contiguous and copies |
| 91 | + those bytes to the device. to_edge makes the parameter contiguous |
| 92 | + before the delegate is compiled, so the delegate expects that layout.""" |
| 93 | + |
| 94 | + class Model(torch.nn.Module): |
| 95 | + def __init__(self): |
| 96 | + super().__init__() |
| 97 | + weight = torch.arange(6.0).reshape(3, 2).t() |
| 98 | + self.weight = torch.nn.Parameter(weight, requires_grad=False) |
| 99 | + |
| 100 | + def forward(self, x): |
| 101 | + return torch.add(x, self.weight) |
| 102 | + |
| 103 | + model = Model() |
| 104 | + x = torch.arange(6.0).reshape(2, 3) * 10 |
| 105 | + (output,) = run(model, (x,)) |
| 106 | + torch.testing.assert_close(output, model(x)) |
| 107 | + |
| 108 | + def test_constant_read_by_two_delegates(self): |
| 109 | + """Two delegates read a constant stored in the program from outside |
| 110 | + them, so the program copies it to the device for them.""" |
| 111 | + |
| 112 | + class Model(torch.nn.Module): |
| 113 | + def __init__(self): |
| 114 | + super().__init__() |
| 115 | + self.register_buffer("table", torch.arange(6.0).reshape(2, 3)) |
| 116 | + |
| 117 | + def forward(self, x): |
| 118 | + return torch.add(x, self.table), torch.add(x * 2, self.table) |
| 119 | + |
| 120 | + model = Model() |
| 121 | + x = torch.arange(6.0).reshape(2, 3) * 10 |
| 122 | + for output, expected in zip(run(model, (x,)), model(x), strict=True): |
| 123 | + torch.testing.assert_close(output, expected) |
| 124 | + |
| 125 | + def test_input_on_device_also_read_on_host(self): |
| 126 | + """With skip_h2d_for_method_inputs, the caller passes the input on the |
| 127 | + device. A delegate reads it there, and a host op reads a copy of it in |
| 128 | + host memory.""" |
| 129 | + |
| 130 | + class Model(torch.nn.Module): |
| 131 | + def forward(self, x, y): |
| 132 | + return torch.add(x, y), torch.sin(x) |
| 133 | + |
| 134 | + config = ExecutorchBackendConfig( |
| 135 | + propagate_device_config=PropagateDeviceConfig( |
| 136 | + skip_h2d_for_method_inputs=True |
| 137 | + ), |
| 138 | + enable_non_cpu_memory_planning=True, |
| 139 | + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), |
| 140 | + ) |
| 141 | + model = Model() |
| 142 | + inputs = (torch.randn(2, 3), torch.randn(2, 3)) |
| 143 | + with tempfile.TemporaryDirectory() as directory: |
| 144 | + method = load(model, inputs, directory, config) |
| 145 | + outputs = method.execute([t.cuda() for t in inputs]) |
| 146 | + for output, expected in zip(outputs, model(*inputs), strict=True): |
| 147 | + torch.testing.assert_close(output.cpu(), expected) |
0 commit comments