diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index a069db4c1fc..785d89d9c9f 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -331,9 +331,9 @@ def preprocess( custom_pass(device_edge_program.graph_module) # ``run_decompositions`` retraces the complete ExportedProgram even - # when none of the table's operators occur in the graph. The in-place - # passes above preserve graph inputs and outputs, so the existing graph - # signature remains valid when the decomposition table cannot apply. + # when none of the table's operators occur in the graph. Nothing else + # rebuilds the graph signature, so the in-place passes above must keep + # it valid when the decomposition table cannot apply. if contains_any_call_fn_target_op( device_edge_program.graph_module, decomposition_table ): diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 7e8d4c88675..0b924242754 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -80,6 +80,7 @@ fbcode_target( "passes/move_cond_predicate_to_cpu.py", "passes/lower_offgraph_kv.py", "passes/remove_assert_async.py", + "passes/replace_complex_with_real.py", "passes/replace_fused_sdpa.py", "passes/replace_int64_floordiv.py", "passes/replace_masked_copy_with_where.py", diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 3f368cdd37e..b35a6373dd4 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -44,6 +44,9 @@ MoveCondPredicateToCpuPass, ) from executorch.backends.cuda.passes.remove_assert_async import RemoveAssertAsyncPass +from executorch.backends.cuda.passes.replace_complex_with_real import ( + ReplaceComplexWithRealPass, +) from executorch.backends.cuda.passes.replace_fused_sdpa import ( ReplaceFusedSdpaWithSdpaPass, ) @@ -1266,6 +1269,7 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any] triton_kernel_mode = mode passes = [ RemoveAssertAsyncPass(), + ReplaceComplexWithRealPass(), MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass(), ReplaceFusedSdpaWithSdpaPass(), diff --git a/backends/cuda/passes/replace_complex_with_real.py b/backends/cuda/passes/replace_complex_with_real.py new file mode 100644 index 00000000000..baeb1183474 --- /dev/null +++ b/backends/cuda/passes/replace_complex_with_real.py @@ -0,0 +1,398 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Graph Transformation Pass for Complex Numbers. + +Inductor generates no code for complex numbers. It calls an ATen kernel for +every op that reads or writes a complex tensor, and this backend ships none of +those kernels. So this pass keeps each complex tensor as a real tensor with a +last dim of size 2 that holds the real and imaginary parts, the layout of +``torch.view_as_real``, and rewrites the ops on it as ops on real numbers: + + view_as_complex(x) -> x + view_as_real(z) -> z + polar(r, theta) -> stack([r * cos(theta), r * sin(theta)], -1) + z * w -> stack([zr * wr - zi * wi, zr * wi + zi * wr], -1) + +Views, copies, indexing, repeats, splits, concatenation and stacking act on +both parts alike, so they keep their op, with the dims they name counted in +the complex tensor. A cast of a real tensor to complex gets zero imaginary +parts. A complex constant, parameter or buffer is stored in the real layout +too. A decomposition table cannot do this, since a decomposition returns the +dtype of the op it replaces. + +The pass rewrites a program only when it can rewrite every op that reads or +writes a complex tensor in it. Otherwise it leaves the program as it is and +logs why, and the backend reports the ops it has no kernels for. It never +rewrites control flow branches, such as those of torch.cond, or complex32 +values. +""" + +import logging +import operator +from typing import Any, Callable, Dict, List, Optional, Set, Tuple + +import torch +import torch.utils._pytree as pytree +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import ExportedProgram +from torch.fx import Graph, Node + +logger = logging.getLogger(__name__) + +_aten = exir_ops.edge.aten + + +def _complex_vals(node: Node) -> List[torch.Tensor]: + val = node.meta.get("val") + vals = val if isinstance(val, (list, tuple)) else (val,) + return [v for v in vals if isinstance(v, torch.Tensor) and v.is_complex()] + + +# The pair's extra dim comes last, so a dim of the complex tensor names the +# same dim of the pair once it is counted from the front. +def _dim(dim: int, ndim: int) -> int: + return dim + ndim if dim < 0 else dim + + +def _dims(dims: List[int], ndim: int) -> List[int]: + return [_dim(d, ndim) for d in dims] + + +def _new_dim(dim: int, ndim: int) -> int: + return _dim(dim, ndim + 1) + + +def _permutation(dims: List[int], ndim: int) -> List[int]: + return _dims(dims, ndim) + [ndim] + + +def _repeats(repeats: List[Any], ndim: int) -> List[Any]: + return [*repeats, 1] + + +def _size(size: List[Any], ndim: int) -> List[Any]: + return [*size, 2] + + +# Ops that act on the real and imaginary parts alike. For each argument that +# names dims or sizes of the complex tensor, the entry gives the argument's +# position, and how to name the same dims of the pair. These arguments' +# default values never need changing. +_SAME_ON_BOTH_PARTS: Dict[Any, Tuple[Tuple[int, Callable[..., Any]], ...]] = { + # Torch-TensorRT turns every view into _reshape_copy. + _aten._reshape_copy.default: ((1, _size),), + _aten.alias_copy.default: (), + _aten.cat.default: ((1, _dim),), + _aten.clone.default: (), + _aten.expand_copy.default: ((1, _size),), + _aten.index.Tensor: (), + _aten.permute_copy.default: ((1, _permutation),), + _aten.repeat.default: ((1, _repeats),), + _aten.select.int: ((1, _dim),), + _aten.select_copy.int: ((1, _dim),), + _aten.slice.Tensor: ((1, _dim),), + _aten.slice_copy.Tensor: ((1, _dim),), + _aten.split.Tensor: ((2, _dim),), + _aten.split_copy.Tensor: ((2, _dim),), + _aten.split_with_sizes_copy.default: ((2, _dim),), + _aten.squeeze_copy.default: (), + _aten.squeeze_copy.dim: ((1, _dim),), + _aten.squeeze_copy.dims: ((1, _dims),), + _aten.stack.default: ((1, _new_dim),), + _aten.transpose_copy.int: ((1, _dim), (2, _dim)), + _aten.unbind_copy.int: ((1, _dim),), + _aten.unsqueeze_copy.default: ((1, _new_dim),), + _aten.view_copy.default: ((1, _size),), +} + +_LIST_INPUTS = (_aten.cat.default, _aten.stack.default) + +# Copies that may change the dtype. The pair gets the precision of the +# complex dtype, and the complex tensor's dim order with its extra dim last. +_COPIES = ( + _aten._to_copy.default, + exir_ops.edge.dim_order_ops._clone_dim_order.default, + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, +) +_DENSE_FORMATS = (None, torch.preserve_format, torch.contiguous_format) +_PROVENANCE = ("from_node", "nn_module_stack", "source_fn_stack", "stack_trace") + +# The fake kernels of some view copies, such as permute_copy, make dynamic +# sizes static or fail on them, so a new view copy gets its value from its +# view, laid out as the copy's would be. view_copy keeps its own kernel, which +# keeps dynamic sizes, because view fails on some strides that view_copy +# accepts. +_VIEWS: Dict[Any, Any] = { + _aten.alias_copy.default: torch.ops.aten.alias.default, + _aten.expand_copy.default: torch.ops.aten.expand.default, + _aten.permute_copy.default: torch.ops.aten.permute.default, + _aten.select_copy.int: torch.ops.aten.select.int, + _aten.slice_copy.Tensor: torch.ops.aten.slice.Tensor, + _aten.split_copy.Tensor: torch.ops.aten.split.Tensor, + _aten.split_with_sizes_copy.default: torch.ops.aten.split_with_sizes.default, + _aten.squeeze_copy.default: torch.ops.aten.squeeze.default, + _aten.squeeze_copy.dim: torch.ops.aten.squeeze.dim, + _aten.squeeze_copy.dims: torch.ops.aten.squeeze.dims, + _aten.transpose_copy.int: torch.ops.aten.transpose.int, + _aten.unbind_copy.int: torch.ops.aten.unbind.int, + _aten.unsqueeze_copy.default: torch.ops.aten.unsqueeze.default, +} + + +class _Rewriter: + """Builds, in graph order, the real replacement of every node that reads + or writes a complex tensor.""" + + def __init__(self, graph: Graph, complex_nodes: Set[Node]) -> None: + self.graph = graph + self.complex_nodes = complex_nodes + self.ndims = { + node: node.meta["val"].dim() + for node in complex_nodes + if isinstance(node.meta["val"], torch.Tensor) + } + self.replacements: Dict[Node, Node] = {} + self.source: Optional[Node] = None + + def get(self, value: Any) -> Any: + """``value`` with each complex node in it replaced by its pair.""" + return torch.fx.map_arg( + value, + lambda n: self.replacements[n] if n in self.complex_nodes else n, + ) + + def call(self, target: Any, *args: Any, **kwargs: Any) -> Node: + node = self.graph.call_function(target, args, kwargs) + # Errors and profiles then point at the op that this one replaces. + node.meta.update( + {k: v for k, v in self.source.meta.items() if k in _PROVENANCE} + ) + fake_args, fake_kwargs = torch.fx.map_arg( + (args, kwargs), lambda n: n.meta["val"] + ) + value = _VIEWS.get(target, target)(*fake_args, **fake_kwargs) + if target in _VIEWS: + value = pytree.tree_map_only( + torch.Tensor, + lambda t: t.clone(memory_format=torch.contiguous_format), + value, + ) + node.meta["val"] = value + return node + + def parts(self, value: Node) -> List[Node]: + pair = self.get(value) + return [self.call(_aten.select.int, pair, -1, i) for i in (0, 1)] + + +def _view_as_complex(rw: _Rewriter, node: Node) -> Node: + return node.args[0] + + +def _view_as_real(rw: _Rewriter, node: Node) -> Node: + return rw.get(node.args[0]) + + +def _polar(rw: _Rewriter, node: Node) -> Node: + magnitude, angle = node.args + parts = [ + rw.call(_aten.mul.Tensor, magnitude, rw.call(trig, angle)) + for trig in (_aten.cos.default, _aten.sin.default) + ] + return rw.call(_aten.stack.default, parts, -1) + + +def _mul(rw: _Rewriter, node: Node) -> Node: + (ar, ai), (br, bi) = (rw.parts(arg) for arg in node.args) + mul = _aten.mul.Tensor + real = rw.call(_aten.sub.Tensor, rw.call(mul, ar, br), rw.call(mul, ai, bi)) + imag = rw.call(_aten.add.Tensor, rw.call(mul, ar, bi), rw.call(mul, ai, br)) + return rw.call(_aten.stack.default, [real, imag], -1) + + +def _copy(rw: _Rewriter, node: Node) -> Node: + source, kwargs = node.args[0], dict(node.kwargs) + if kwargs.get("dtype") is not None: + kwargs["dtype"] = kwargs["dtype"].to_real() + if source not in rw.complex_nodes: + # Edge programs cast a real operand to complex before complex math + # with it. The cast gives zero imaginary parts. + real = rw.call(node.target, source, **kwargs) + imag = rw.call(_aten.zeros_like.default, real) + return rw.call(_aten.stack.default, [real, imag], -1) + if kwargs.get("dim_order") is not None: + kwargs["dim_order"] = [*kwargs["dim_order"], len(kwargs["dim_order"])] + return rw.call(node.target, rw.get(source), **kwargs) + + +def _same_on_both_parts(rw: _Rewriter, node: Node) -> Node: + args = list(node.args) + tensor = args[0][0] if node.target in _LIST_INPUTS else args[0] + for index, pair_value in _SAME_ON_BOTH_PARTS[node.target]: + if index < len(args): + args[index] = pair_value(args[index], rw.ndims[tensor]) + return rw.call(node.target, *rw.get(args), **rw.get(dict(node.kwargs))) + + +def _getitem(rw: _Rewriter, node: Node) -> Node: + return rw.call(operator.getitem, rw.get(node.args[0]), node.args[1]) + + +_REWRITES: Dict[Any, Callable[[_Rewriter, Node], Node]] = { + _aten.view_as_complex_copy.default: _view_as_complex, + _aten.view_as_real_copy.default: _view_as_real, + _aten.polar.default: _polar, + _aten.mul.Tensor: _mul, + operator.getitem: _getitem, + **{op: _copy for op in _COPIES}, + **{op: _same_on_both_parts for op in _SAME_ON_BOTH_PARTS}, +} + + +def _keeps_dims_in_order(node: Node) -> bool: + # With the pair's dim after them, the dims of a channels-last layout are + # in an order that no memory format describes, and the dim order ops + # accept only those that one does. + dim_order = node.kwargs.get("dim_order") + return node.kwargs.get("memory_format") in _DENSE_FORMATS and ( + dim_order is None or list(dim_order) == list(range(len(dim_order))) + ) + + +def _can_rewrite(node: Node, complex_nodes: Set[Node]) -> bool: + if node.op != "call_function" or node.target not in _REWRITES: + return False + # The rewrites read by position the arguments that exported programs + # pass by position. + schema = getattr(node.target, "_schema", None) + if schema is not None and any( + not arg.kwarg_only for arg in schema.arguments if arg.name in node.kwargs + ): + return False + if node.target in _COPIES: + # A cast from complex to real has no rewrite here. + return node in complex_nodes and _keeps_dims_in_order(node) + if node.target is _aten.clone.default: + return _keeps_dims_in_order(node) + if node.target is _aten.mul.Tensor: + # Edge programs cast a real factor to complex first. + return all(arg in complex_nodes for arg in node.args) + if node.target in _LIST_INPUTS: + return all(tensor in complex_nodes for tensor in node.args[0]) + return True + + +def _constant_name(exported_program: ExportedProgram, node: Node) -> Optional[str]: + """The name of the constant, parameter or buffer that ``node`` reads, + unless it is a buffer that the graph writes.""" + signature = exported_program.graph_signature + for inputs in ( + signature.inputs_to_lifted_tensor_constants, + signature.inputs_to_parameters, + ): + if node.name in inputs: + return inputs[node.name] + buffer = signature.inputs_to_buffers.get(node.name) + if buffer in signature.buffers_to_mutate.values(): + return None + return buffer + + +def _store_as_real(exported_program: ExportedProgram, node: Node) -> None: + name = _constant_name(exported_program, node) + tensors = ( + exported_program._constants + if name in exported_program.constants + else exported_program._state_dict + ) + tensor = tensors[name] + # Partitioning can give one tensor several placeholders, and the first of + # them stores it. + if tensor.is_complex(): + pair = torch.view_as_real(tensor.detach().resolve_conj()).clone() + if isinstance(tensor, torch.nn.Parameter): + pair = torch.nn.Parameter(pair, requires_grad=False) + tensors[name] = pair + node.meta["val"] = torch.view_as_real(node.meta["val"].resolve_conj()) + + +def _reason_to_keep( + exported_program: ExportedProgram, complex_nodes: Set[Node], affected: Set[Node] +) -> Optional[str]: + """Why the pass must leave the program as it is, if it must.""" + graph_module = exported_program.graph_module + for module in graph_module.modules(): + if module is not graph_module and isinstance(module, torch.fx.GraphModule): + for node in module.graph.nodes: + if _complex_vals(node): + return f"it does not rewrite control flow branches: {node.format_node()}" + for node in graph_module.graph.nodes: + if node not in affected: + continue + if node.op == "placeholder": + supported = _constant_name(exported_program, node) is not None + else: + supported = _can_rewrite(node, complex_nodes) + if not supported: + return f"it cannot rewrite {node.format_node()}" + # PyTorch computes complex32 math in float32, and its pair would + # compute in float16. + if any(val.dtype == torch.complex32 for val in _complex_vals(node)): + return f"it does not rewrite complex32: {node.format_node()}" + return None + + +class ReplaceComplexWithRealPass: + """ + Replace complex tensors with real tensors that hold their real and + imaginary parts in a last dim of size 2, and the ops on them with ops on + real numbers, so that Inductor can generate code for them. See the module + docstring for the ops it handles. + """ + + requires_exported_program = True + + def __call__(self, exported_program: ExportedProgram) -> None: + graph_module = exported_program.graph_module + graph = graph_module.graph + complex_nodes = {node for node in graph.nodes if _complex_vals(node)} + affected = complex_nodes | { + user for node in complex_nodes for user in node.users + } + reason = _reason_to_keep(exported_program, complex_nodes, affected) + if reason is not None: + logger.warning( + "The CUDA backend cannot lower the complex numbers in this " + "program, since %s", + reason, + ) + return + if not complex_nodes: + return + + rw = _Rewriter(graph, complex_nodes) + order = [node for node in graph.nodes if node in affected] + for node in order: + if node.op == "placeholder": + _store_as_real(exported_program, node) + rw.replacements[node] = node + continue + rw.source = node + with graph.inserting_before(node): + rw.replacements[node] = _REWRITES[node.target](rw, node) + signature = exported_program.graph_signature + for node in reversed(order): + replacement = rw.replacements[node] + if replacement is not node: + node.replace_all_uses_with(replacement) + # The signature names the program's outputs by node, and a + # view_as_real can be one of them. + signature.replace_all_uses(node.name, replacement.name) + graph.erase_node(node) + graph.lint() + graph_module.recompile() diff --git a/backends/cuda/passes/tests/test_replace_complex_with_real.py b/backends/cuda/passes/tests/test_replace_complex_with_real.py new file mode 100644 index 00000000000..793fe16555c --- /dev/null +++ b/backends/cuda/passes/tests/test_replace_complex_with_real.py @@ -0,0 +1,159 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +import torch.utils._pytree as pytree +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.passes.replace_complex_with_real import ( + _VIEWS, + ReplaceComplexWithRealPass, +) +from executorch.backends.cuda.passes.tests.test_replace_fused_sdpa import ( + _Fn, + _program_for_cuda_passes, +) +from executorch.exir.backend.canonical_partitioners.duplicate_constant_node_pass import ( + duplicate_constant_node, +) + + +class TestReplaceComplexWithRealPass(unittest.TestCase): + def test_registered(self): + """The rewrite must stay in the CUDA backend's passes.""" + self.assertTrue( + any( + isinstance(p, ReplaceComplexWithRealPass) + for p in CudaBackend.get_custom_passes([]) + ) + ) + + def test_rewrite(self): + """The rewrite of complex numbers gives eager's results.""" + angles = torch.outer(torch.arange(8.0), torch.rand(3)) + + class Rotate(torch.nn.Module): + def __init__(self): + super().__init__() + # A plain attribute becomes a constant of the program. + self.freqs = torch.polar(torch.ones_like(angles), angles) + self.phase = torch.nn.Parameter( + torch.polar(torch.ones(2, 6), torch.rand(2, 6)), requires_grad=False + ) + + def forward(self, x, positions, scale): + f = torch.stack(self.freqs[positions].unbind(-1), -1) + f = f.unsqueeze(-2).expand(-1, 2, -1).flatten(-2) + # Torch-TensorRT turns views into _reshape_copy. + f = torch.ops.aten._reshape_copy(f, [-1, 1, 6]) + pairs = torch.view_as_complex(x.unflatten(-1, (-1, 2))) + # The product keeps the permuted layout, so export flattens + # the transpose back with a view. + scaled = pairs.unflatten(-1, (2, 3)).permute(0, 1, -1, -2) * scale + scaled = scaled.transpose(-1, -2).flatten(-2) + phase = self.phase.repeat(2, 1).select(-2, 1).detach() + # The cast sets the precision of everything after it. + rotated = scaled * f.to(torch.complex128) * phase + first, second = rotated.chunk(2, dim=-1) + second = torch.cat(second.split([1, 2], dim=-1), -1) + joined = torch.cat([first, second], -1) + return ( + torch.view_as_real(joined.squeeze(-2).clone()), + torch.view_as_real(joined.unsqueeze(0).squeeze((0, -2))), + ) + + model = Rotate() + scale = torch.randn(3, 2) + x = torch.randn(5, 1, 12) + length = torch.export.Dim("length", min=2, max=64) + program = _program_for_cuda_passes( + model, + (x, torch.tensor([3, 0, 7, 7, 1]), scale), + dynamic_shapes=({0: length}, {0: length}, None), + ) + ReplaceComplexWithRealPass()(program) + for node in program.graph.nodes: + val = node.meta.get("val") + self.assertFalse(isinstance(val, torch.Tensor) and val.is_complex()) + if node.target in _VIEWS: + # A view copy returns new tensors laid out in order. + self.assertTrue(all(t.is_contiguous() for t in pytree.tree_leaves(val))) + # The backend retraces the program when its decomposition table applies. + program = program.run_decompositions({}) + for positions in (torch.tensor([3, 0, 7, 7, 1]), torch.tensor([6, 2, 4])): + inputs = (x[: len(positions)], positions, scale) + torch.testing.assert_close(program.module()(*inputs), model(*inputs)) + + def test_tensor_with_two_placeholders(self): + """Partitioning gives a tensor a placeholder per op that reads it when + another backend takes ops between those, and the pass stores the + tensor once.""" + + class Rotate(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("freqs", torch.polar(torch.ones(3), torch.rand(3))) + + def forward(self, x, y): + return tuple( + torch.view_as_real(torch.view_as_complex(t) * self.freqs) + for t in (x, y) + ) + + model = Rotate() + inputs = (torch.randn(3, 2), torch.randn(3, 2)) + program = _program_for_cuda_passes(model, inputs) + self.assertEqual( + duplicate_constant_node(program, "b_freqs"), {"b_freqs_copy_0"} + ) + ReplaceComplexWithRealPass()(program) + for node in program.graph.nodes: + val = node.meta.get("val") + self.assertFalse(isinstance(val, torch.Tensor) and val.is_complex()) + torch.testing.assert_close(program.module()(*inputs), model(*inputs)) + + def test_leaves_unsupported_programs(self): + """A program with a complex op that the rewrite does not handle keeps + all its complex numbers, so no program is left half rewritten.""" + + class MutatedBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("freqs", torch.randn(3, dtype=torch.cfloat)) + + def forward(self, x): + self.freqs.mul_(torch.view_as_complex(x)) + return torch.view_as_real(self.freqs.clone()) + + def square(x): + z = torch.view_as_complex(x) + return torch.view_as_real(z * z) + + def channels_last(x): + z = torch.view_as_complex(x).clone(memory_format=torch.channels_last) + return torch.view_as_real(z) + + def in_branch(x): + return torch.cond(x.sum() > 0, square, torch.neg, (square(x),)) + + pairs = torch.randn(1, 2, 3, 4, 2) + for model, inputs in ( + (_Fn(lambda z: torch.view_as_real(z * z)), (torch.randn(3).cfloat(),)), + (_Fn(lambda x: torch.view_as_real(torch.view_as_complex(x) + 1)), (pairs,)), + (_Fn(square), (pairs.half(),)), + (_Fn(lambda x: torch.view_as_complex(x).float()), (pairs,)), + (_Fn(channels_last), (pairs,)), + (_Fn(in_branch), (pairs,)), + (MutatedBuffer(), (torch.randn(3, 2),)), + ): + program = _program_for_cuda_passes(model, inputs) + code = program.graph_module.code + tensors = {**program.state_dict, **program.constants} + ReplaceComplexWithRealPass()(program) + self.assertEqual(program.graph_module.code, code) + for name, tensor in {**program.state_dict, **program.constants}.items(): + self.assertIs(tensor, tensors[name]) diff --git a/backends/cuda/runtime/shims/memory.cpp b/backends/cuda/runtime/shims/memory.cpp index a9a39a7c7b1..f36b283b4c0 100644 --- a/backends/cuda/runtime/shims/memory.cpp +++ b/backends/cuda/runtime/shims/memory.cpp @@ -46,6 +46,29 @@ using c10::ScalarType; namespace { +// A graph capture cannot record a synchronous copy to or from device memory, +// and such a copy aborts the process when it fails, so a shim that makes one +// refuses first. Info, not Error: the backend then retries the call without a +// graph. An empty tensor copies nothing, so it needs no check. +AOTITorchError refuse_inside_capture( + const SlimTensor* tensor, + const char* what) { +#if defined(CUDA_AVAILABLE) + if (tensor->is_cuda() && tensor->numel() > 0) { + auto stream = getCurrentCUDAStream(tensor->device_index()); + ET_CHECK_OK_OR_RETURN_ERROR(stream.error()); + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + ET_CUDA_CHECK_OR_RETURN_ERROR( + cudaStreamIsCapturing(stream.get(), &capture)); + if (capture != cudaStreamCaptureStatusNone) { + ET_LOG(Info, "%s: not supported inside a CUDA graph capture", what); + return Error::NotSupported; + } + } +#endif + return Error::Ok; +} + // Reads the one element as T. Dispatch on the dtype the tensor actually holds: // item() copies sizeof(T) bytes and does not check, so asking for the wrong // width reads past a one-element allocation. @@ -67,25 +90,8 @@ AOTITorchError narrow_to(From value, To* ret_value) { template AOTITorchError read_one_element(const SlimTensor* tensor, T* ret_value) { -#if defined(CUDA_AVAILABLE) - // A graph capture cannot record the copy to the host below, and the copy - // aborts the process when it fails, so refuse first. Info, not Error: the - // backend then retries the call without a graph. - if (tensor->is_cuda()) { - auto stream = getCurrentCUDAStream(tensor->device_index()); - ET_CHECK_OK_OR_RETURN_ERROR(stream.error()); - cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; - ET_CUDA_CHECK_OR_RETURN_ERROR( - cudaStreamIsCapturing(stream.get(), &capture)); - if (capture != cudaStreamCaptureStatusNone) { - ET_LOG( - Info, - "reading a single element: a device value cannot be read inside a " - "CUDA graph capture"); - return Error::NotSupported; - } - } -#endif + ET_CHECK_OK_OR_RETURN_ERROR( + refuse_inside_capture(tensor, "reading a single element")); switch (tensor->dtype()) { case ScalarType::Byte: return narrow_to(tensor->item(), ret_value); @@ -352,6 +358,8 @@ AOTITorchError aoti_torch_clone_preserve_strides( ret_new_tensor != nullptr, InvalidArgument, "aoti_torch_clone_preserve_strides: ret_new_tensor is null"); + ET_CHECK_OK_OR_RETURN_ERROR( + refuse_inside_capture(self, "aoti_torch_clone_preserve_strides")); *ret_new_tensor = new SlimTensor(self->clone()); return Error::Ok; @@ -365,6 +373,7 @@ AOTITorchError aoti_torch_clone(SlimTensor* self, SlimTensor** ret_new_tensor) { ret_new_tensor != nullptr, InvalidArgument, "aoti_torch_clone: ret_new_tensor is null"); + ET_CHECK_OK_OR_RETURN_ERROR(refuse_inside_capture(self, "aoti_torch_clone")); *ret_new_tensor = new SlimTensor(self->clone()); return Error::Ok; @@ -437,6 +446,9 @@ aoti_torch_copy_(SlimTensor* self, SlimTensor* src, int32_t non_blocking) { ET_CHECK_OR_RETURN_ERROR( src != nullptr, InvalidArgument, "aoti_torch_copy_: src is null"); + ET_CHECK_OK_OR_RETURN_ERROR(refuse_inside_capture(self, "aoti_torch_copy_")); + ET_CHECK_OK_OR_RETURN_ERROR(refuse_inside_capture(src, "aoti_torch_copy_")); + // SlimTensor::copy_() handles: // - Same numel validation // - Same dtype validation diff --git a/backends/cuda/runtime/shims/memory.h b/backends/cuda/runtime/shims/memory.h index ef51ff5eef0..4c4d67b7a53 100644 --- a/backends/cuda/runtime/shims/memory.h +++ b/backends/cuda/runtime/shims/memory.h @@ -150,6 +150,10 @@ aoti_torch_new_tensor_handle(SlimTensor* orig_handle, SlimTensor** new_handle); /** * Clone APIs. + * + * A clone of a CUDA tensor is a synchronous copy, which a CUDA graph capture + * cannot record, so both return Error::NotSupported while the current stream + * is capturing. */ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_clone_preserve_strides( SlimTensor* self, @@ -207,7 +211,9 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_as_strided( * @param self Destination tensor (must not be null) * @param src Source tensor to copy from (must not be null) * @param non_blocking If true, the copy may be asynchronous (currently ignored) - * @return AOTITorchError error code (Error::Ok on success) + * @return Error::Ok on success, Error::NotSupported if either tensor is a CUDA + * tensor and the current stream is capturing a CUDA graph, which cannot record + * the synchronous copy */ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_copy_(SlimTensor* self, SlimTensor* src, int32_t non_blocking); @@ -231,7 +237,8 @@ aoti_torch_item_bool(SlimTensor* tensor, bool* ret_value); * * @param tensor Single-element tensor (must not be null) * @param ret_value Output parameter for the extracted value - * @return AOTITorchError error code (Error::Ok on success) + * @return Error::Ok on success, Error::NotSupported for a CUDA tensor while the + * current stream is capturing a CUDA graph, which cannot record the copy */ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_item_uint8(SlimTensor* tensor, uint8_t* ret_value); diff --git a/backends/cuda/runtime/shims/tests/test_aoti_torch_copy_.cpp b/backends/cuda/runtime/shims/tests/test_aoti_torch_copy_.cpp index 392a3e29244..df67f002184 100644 --- a/backends/cuda/runtime/shims/tests/test_aoti_torch_copy_.cpp +++ b/backends/cuda/runtime/shims/tests/test_aoti_torch_copy_.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -449,6 +450,83 @@ TEST_F(AOTITorchCopySlimTest, CudaToCpu) { EXPECT_EQ(aoti_torch_delete_tensor_object(dst), Error::Ok); } +TEST_F(AOTITorchCopySlimTest, RefusedInsideGraphCapture) { + if (!isCudaAvailable()) { + GTEST_SKIP() << "CUDA not available"; + } + + std::vector sizes = {2, 3}; + Tensor* device = createTestTensor( + sizes, + {}, + static_cast(slim_c10::ScalarType::Float), + static_cast(slim_c10::DeviceType::CUDA), + 0); + Tensor* other_device = createTestTensor( + sizes, + {}, + static_cast(slim_c10::ScalarType::Float), + static_cast(slim_c10::DeviceType::CUDA), + 0); + Tensor* host = createTestTensor(sizes); + std::vector empty_sizes = {0, 3}; + Tensor* empty_device = createTestTensor( + empty_sizes, + {}, + static_cast(slim_c10::ScalarType::Float), + static_cast(slim_c10::DeviceType::CUDA), + 0); + Tensor* other_empty_device = createTestTensor( + empty_sizes, + {}, + static_cast(slim_c10::ScalarType::Float), + static_cast(slim_c10::DeviceType::CUDA), + 0); + ASSERT_NE(device, nullptr); + ASSERT_NE(other_device, nullptr); + ASSERT_NE(host, nullptr); + ASSERT_NE(empty_device, nullptr); + ASSERT_NE(other_empty_device, nullptr); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + ASSERT_EQ(setCurrentCUDAStream(stream, 0), Error::Ok); + ASSERT_EQ( + cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed), + cudaSuccess); + // A capture cannot record these synchronous copies, so each one refuses + // instead of failing inside the copy. + EXPECT_EQ(aoti_torch_copy_(host, device, 0), Error::NotSupported); + EXPECT_EQ(aoti_torch_copy_(device, host, 0), Error::NotSupported); + EXPECT_EQ(aoti_torch_copy_(other_device, device, 0), Error::NotSupported); + Tensor* clone = nullptr; + EXPECT_EQ(aoti_torch_clone(device, &clone), Error::NotSupported); + EXPECT_EQ( + aoti_torch_clone_preserve_strides(device, &clone), Error::NotSupported); + // A copy in host memory does not touch the stream, and an empty one copies + // nothing. + EXPECT_EQ(aoti_torch_clone(host, &clone), Error::Ok); + if (clone != nullptr) { + EXPECT_EQ(aoti_torch_delete_tensor_object(clone), Error::Ok); + } + EXPECT_EQ(aoti_torch_copy_(empty_device, other_empty_device, 0), Error::Ok); + cudaGraph_t graph = nullptr; + // Nothing reached the stream, so the capture is still valid. + EXPECT_EQ(cudaStreamEndCapture(stream, &graph), cudaSuccess); + if (graph != nullptr) { + EXPECT_EQ(cudaGraphDestroy(graph), cudaSuccess); + } + EXPECT_EQ(aoti_torch_copy_(host, device, 0), Error::Ok); + + clearCurrentCUDAStream(0); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); + EXPECT_EQ(aoti_torch_delete_tensor_object(device), Error::Ok); + EXPECT_EQ(aoti_torch_delete_tensor_object(other_device), Error::Ok); + EXPECT_EQ(aoti_torch_delete_tensor_object(host), Error::Ok); + EXPECT_EQ(aoti_torch_delete_tensor_object(empty_device), Error::Ok); + EXPECT_EQ(aoti_torch_delete_tensor_object(other_empty_device), Error::Ok); +} + // ============================================================================ // Non-blocking Tests // ============================================================================ diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 7d1a0e91fea..1e9489479bc 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -556,13 +556,21 @@ def forward(self, q, k, v, cu_seq): (big, big, v, cu_seq, cu_seq, 16, 16), ] table = CudaBackend.get_decomposition_table() - # On ROCm, varlen_attn runs AOTriton, which lines an explicit window up - # with the start of each sequence. The decomposition follows Flash - # attention on CUDA, which lines it up with the end, so compare on CUDA. - if torch.version.hip is None: - for args in cases: - for actual, expected in zip(table[op](*args)[:2], op(*args)[:2]): - torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + for args in cases: + # On ROCm, varlen_attn runs AOTriton, which lines an explicit window + # up with the start of each sequence. The decomposition follows + # Flash attention on CUDA, which lines it up with the end. The two + # agree when the query and key sequences line up, as in + # self-attention. + window = args[9] if len(args) > 9 else None + if ( + torch.version.hip is not None + and window is not None + and not torch.equal(args[3], args[4]) + ): + continue + for actual, expected in zip(table[op](*args)[:2], op(*args)[:2]): + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) # Flash attention versions differ on this window, since a query sequence # can be longer than max_k + 1. with self.assertRaisesRegex(NotImplementedError, "window"): @@ -1000,6 +1008,79 @@ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]: for output, expected in zip(outputs, model(x), strict=True): torch.testing.assert_close(output.cpu(), expected) + @unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", + ) + def test_complex_rotary_embedding(self): + """Inductor generates no code for complex numbers, so the backend + computes a rotary embedding written with them on real numbers.""" + from executorch.runtime import Runtime + + class ComplexRope(torch.nn.Module): + def __init__(self, precision): + super().__init__() + self.precision = precision + angles = torch.outer(torch.arange(4.0), torch.rand(12)).to(precision) + self.register_buffer( + "freqs", torch.polar(torch.ones_like(angles), angles) + ) + + def forward(self, x, angles): + # A grid of positions takes each axis's part of the frequencies. + rows, cols = self.freqs.split([8, 4], dim=1) + grid = torch.cat( + [ + rows[:2].view(2, 1, -1).expand(2, 2, -1), + cols[:2].view(1, 2, -1).expand(2, 2, -1), + ], + dim=-1, + ).view(4, 1, -1) + spin = torch.polar(angles.abs(), angles) + pairs = torch.view_as_complex( + x.to(self.precision).unflatten(-1, (-1, 2)) + ) + rotated = torch.view_as_real(pairs * grid * spin.unsqueeze(1) * 0.5) + # A view_as_real can also be an output of the program, and two + # of one complex tensor become one output in two slots. + return ( + rotated.flatten(-2).type_as(x), + torch.view_as_real(spin), + torch.view_as_real(spin), + ) + + angles = torch.randn(4, 12) + # Rotary embeddings are often computed in float64, so in complex128. + for dtype, precision in ( + (torch.float32, torch.float32), + (torch.bfloat16, torch.float32), + (torch.float32, torch.float64), + ): + with self.subTest(dtype=dtype, precision=precision): + model = ComplexRope(precision) + x = torch.randn(4, 3, 24, dtype=dtype) + program = self._export_to_cuda_with_lower( + model, (x, angles) + ).to_executorch() + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + program.write_to_file(f) + program.write_tensor_data_to_file(directory) + (data,) = [ + os.path.join(directory, name) + for name in os.listdir(directory) + if name.endswith(".ptd") + ] + method = ( + Runtime.get() + .load_program(path, data_path=data) + .load_method("forward") + ) + outputs = method.execute([x, angles]) + for output, expected in zip(outputs, model(x, angles), strict=True): + torch.testing.assert_close(output.cpu(), expected) + def test_device_info_propagated_to_cuda_delegate_outputs(self): """ Verify that, for a CUDA-delegated graph, every memory-planned tensor's diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index a53156cd3f0..176acd089e5 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -316,6 +316,17 @@ def forward(self, x, mask): return self.calls.clone(), x[mask] +class _ConstantRow(nn.Module): + """AOTI returns a part of a constant through aoti_torch_clone.""" + + def __init__(self): + super().__init__() + self.register_buffer("table", torch.randn(3, 4)) + + def forward(self, x): + return self.table[1], x * 2 + + class _HostFlag(nn.Module): """A torch.cond on a flag the program keeps in host memory and updates on every call.""" @@ -531,6 +542,20 @@ def test_mask_retry_updates_state_once(self): self.assertIsInstance(out, torch.Tensor) torch.testing.assert_close(out, torch.tensor([i + 1.0])) + def test_constant_copy_runs_without_a_graph(self): + # A capture cannot record the copy that returns a part of a constant, + # so the capture call has to run without a graph. + torch.manual_seed(0) + module = _ConstantRow().eval() + inputs = [torch.randn(4) for _ in range(6)] + with _RuntimeLog() as log: + outputs = self._run_with_cuda_graph(module, (inputs[0],), inputs) + self.assertEqual(log.text.count("capture failed for 'forward'"), 1) + for i, out in enumerate(outputs): + with self.subTest(call=i): + self.assertIsInstance(out, torch.Tensor) + torch.testing.assert_close(out, module.table[1]) + def test_output_shape_changes_at_capture(self): torch.manual_seed(0) # The capture call has a larger batch than the warmup calls before it.