|
| 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 pytest |
| 8 | +import torch |
| 9 | +from executorch.backends.cortex_m.edge_compile_config import ( |
| 10 | + cortex_m_edge_compile_config, |
| 11 | +) |
| 12 | +from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager |
| 13 | +from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer |
| 14 | +from executorch.backends.cortex_m.test.tester import CortexMTester |
| 15 | +from executorch.backends.test.harness.stages import StageType |
| 16 | +from executorch.exir import to_edge_transform_and_lower |
| 17 | +from executorch.exir.dialects._ops import ops as exir_ops |
| 18 | +from torch.export import Dim, export |
| 19 | +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e |
| 20 | + |
| 21 | + |
| 22 | +class ExpandBetweenQuantizedOps(torch.nn.Module): |
| 23 | + def forward(self, x, y, z): |
| 24 | + return (x * y).expand(4, 8) + z |
| 25 | + |
| 26 | + |
| 27 | +class ExpandAsBetweenQuantizedOps(torch.nn.Module): |
| 28 | + def forward(self, x, y, z): |
| 29 | + return (x * y).expand_as(z) + z |
| 30 | + |
| 31 | + |
| 32 | +class ExpandBetweenFloatOps(torch.nn.Module): |
| 33 | + def forward(self, x, z): |
| 34 | + return torch.cumsum(torch.cumsum(x, dim=1).expand(4, 8), dim=0) |
| 35 | + |
| 36 | + |
| 37 | +class ExpandAsBetweenFloatOps(torch.nn.Module): |
| 38 | + def forward(self, x, z): |
| 39 | + return torch.cumsum(torch.cumsum(x, dim=1).expand_as(z), dim=0) |
| 40 | + |
| 41 | + |
| 42 | +class ExpandAsWithQuantizedShapeOperand(torch.nn.Module): |
| 43 | + def forward(self, x, y, z): |
| 44 | + other = z + z |
| 45 | + return (x * y).expand_as(other) + other |
| 46 | + |
| 47 | + |
| 48 | +def quantize(model, inputs, dynamic_shapes=None): |
| 49 | + captured = export(model, inputs, dynamic_shapes=dynamic_shapes).module() |
| 50 | + prepared = prepare_pt2e(captured, CortexMQuantizer()) |
| 51 | + prepared(*inputs) |
| 52 | + return convert_pt2e(prepared) |
| 53 | + |
| 54 | + |
| 55 | +def is_qdq(node) -> bool: |
| 56 | + return "quantize_per_tensor" in str(node.target) |
| 57 | + |
| 58 | + |
| 59 | +def single_expand(gm): |
| 60 | + (expand,) = gm.graph.find_nodes( |
| 61 | + op="call_function", target=torch.ops.aten.expand.default |
| 62 | + ) |
| 63 | + return expand |
| 64 | + |
| 65 | + |
| 66 | +@pytest.mark.parametrize( |
| 67 | + "model,inputs", |
| 68 | + [ |
| 69 | + ( |
| 70 | + ExpandBetweenQuantizedOps(), |
| 71 | + (torch.randn(1, 8), torch.randn(1, 8), torch.randn(4, 8)), |
| 72 | + ), |
| 73 | + ( |
| 74 | + ExpandBetweenQuantizedOps(), |
| 75 | + (torch.randn(8), torch.randn(8), torch.randn(4, 8)), |
| 76 | + ), |
| 77 | + ( |
| 78 | + ExpandAsBetweenQuantizedOps(), |
| 79 | + (torch.randn(1, 8), torch.randn(1, 8), torch.randn(4, 8)), |
| 80 | + ), |
| 81 | + ( |
| 82 | + ExpandAsBetweenQuantizedOps(), |
| 83 | + (torch.randn(8), torch.randn(8), torch.randn(4, 8)), |
| 84 | + ), |
| 85 | + ], |
| 86 | + ids=["expand", "expand_rank_change", "expand_as", "expand_as_rank_change"], |
| 87 | +) |
| 88 | +def test_expand_with_quantized_neighbors_runs_on_int8(model, inputs): |
| 89 | + tester = CortexMTester(model.eval(), inputs) |
| 90 | + tester.quantize().export().to_edge().run_passes() |
| 91 | + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() |
| 92 | + |
| 93 | + (expand,) = program.graph.find_nodes( |
| 94 | + op="call_function", target=exir_ops.edge.aten.expand_copy.default |
| 95 | + ) |
| 96 | + assert expand.args[0].meta["val"].dtype == torch.int8 |
| 97 | + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) |
| 98 | + |
| 99 | + |
| 100 | +@pytest.mark.parametrize( |
| 101 | + "model", [ExpandBetweenFloatOps(), ExpandAsBetweenFloatOps()], ids=["expand", "expand_as"] |
| 102 | +) |
| 103 | +def test_expand_between_float_ops_stays_float(model): |
| 104 | + expand = single_expand(quantize(model, (torch.randn(1, 8), torch.randn(4, 8)))) |
| 105 | + |
| 106 | + assert not is_qdq(expand.args[0]) |
| 107 | + assert not any(is_qdq(user) for user in expand.users) |
| 108 | + |
| 109 | + |
| 110 | +def test_expand_as_does_not_share_qparams_with_its_shape_operand(): |
| 111 | + inputs = (torch.randn(1, 8), torch.randn(1, 8), 10 * torch.randn(4, 8)) |
| 112 | + expand = single_expand(quantize(ExpandAsWithQuantizedShapeOperand(), inputs)) |
| 113 | + |
| 114 | + dequantize = expand.args[0] |
| 115 | + assert is_qdq(dequantize) |
| 116 | + assert dequantize.args[0].args[0].target == torch.ops.aten.mul.Tensor |
| 117 | + (quantize_after_expand,) = expand.users |
| 118 | + (add,) = [ |
| 119 | + user |
| 120 | + for dq in quantize_after_expand.users |
| 121 | + for user in dq.users |
| 122 | + if user.target == torch.ops.aten.add.Tensor |
| 123 | + ] |
| 124 | + other_dequantize = add.args[1] |
| 125 | + assert is_qdq(other_dequantize) |
| 126 | + assert dequantize.args[1] != other_dequantize.args[1] |
| 127 | + assert all(isinstance(size, int) for size in expand.args[1]) |
| 128 | + |
| 129 | + |
| 130 | +def test_expand_as_with_dynamic_shape_reads_the_size_at_runtime(): |
| 131 | + inputs = (torch.randn(1, 8), torch.randn(1, 8), torch.randn(4, 8)) |
| 132 | + dynamic_shapes = ({}, {}, {0: Dim("B")}) |
| 133 | + quantized = quantize(ExpandAsBetweenQuantizedOps(), inputs, dynamic_shapes) |
| 134 | + expand = single_expand(quantized) |
| 135 | + |
| 136 | + assert any( |
| 137 | + isinstance(size, torch.fx.Node) and size.target == torch.ops.aten.sym_size.int |
| 138 | + for size in expand.args[1] |
| 139 | + ) |
| 140 | + assert is_qdq(expand.args[0]) |
| 141 | + assert any(is_qdq(user) for user in expand.users) |
| 142 | + |
| 143 | + edge = to_edge_transform_and_lower( |
| 144 | + export(quantized, inputs, dynamic_shapes=dynamic_shapes), |
| 145 | + compile_config=cortex_m_edge_compile_config(), |
| 146 | + transform_passes=CortexMPassManager(), |
| 147 | + ) |
| 148 | + (expand_copy,) = edge.exported_program().graph.find_nodes( |
| 149 | + op="call_function", target=exir_ops.edge.aten.expand_copy.default |
| 150 | + ) |
| 151 | + assert expand_copy.args[0].meta["val"].dtype == torch.int8 |
0 commit comments