|
| 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 | +# pyre-strict |
| 8 | + |
| 9 | +import copy |
| 10 | +import functools |
| 11 | +import traceback |
| 12 | +import unittest |
| 13 | +from typing import Any, Callable, Sequence |
| 14 | + |
| 15 | +import torch |
| 16 | +from executorch.backends.test.harness.tester import Tester as TesterBase |
| 17 | +from executorch.backends.xnnpack.test.tester.tester import Tester as XnnpackTester |
| 18 | +from facto.inputgen.argtuple.gen import ArgumentTupleGenerator |
| 19 | +from facto.inputgen.specs.model import ConstraintProducer as cp, Spec |
| 20 | +from facto.inputgen.utils.random_manager import random_manager |
| 21 | +from facto.specdb.db import SpecDictDB |
| 22 | +from torch._ops import OpOverload |
| 23 | + |
| 24 | +from .facto_specs import ExtraSpecDB |
| 25 | + |
| 26 | +CombinedSpecDB = SpecDictDB | ExtraSpecDB |
| 27 | + |
| 28 | +COMMON_TENSOR_CONSTRAINTS = [ |
| 29 | + cp.Rank.Ge(lambda deps: 1), |
| 30 | + cp.Rank.Le(lambda deps: 4), |
| 31 | + cp.Size.Ge(lambda deps, r, d: 1), |
| 32 | + cp.Size.Le(lambda deps, r, d: 2**9), |
| 33 | +] |
| 34 | + |
| 35 | +COMMON_SCALAR_CONSTRAINS = [ |
| 36 | + cp.Value.Ge(lambda deps, dtype: -1000), |
| 37 | + cp.Value.Le(lambda deps, dtype: 1000), |
| 38 | +] |
| 39 | + |
| 40 | +# Operator args are treated as runtime graph inputs if the argument name is |
| 41 | +# in this list. |
| 42 | +RUNTIME_INPUT_NAMES = { |
| 43 | + "self", |
| 44 | + "tensor", |
| 45 | + "other", |
| 46 | +} |
| 47 | + |
| 48 | + |
| 49 | +def _patch_spec(spec: Spec) -> Spec: |
| 50 | + spec = copy.deepcopy(spec) |
| 51 | + for inspec in spec.inspec: |
| 52 | + if inspec.type.is_tensor(): |
| 53 | + inspec.constraints.extend(COMMON_TENSOR_CONSTRAINTS) |
| 54 | + elif inspec.type.is_scalar(): |
| 55 | + inspec.constraints.extend(COMMON_SCALAR_CONSTRAINS) |
| 56 | + return spec |
| 57 | + |
| 58 | + |
| 59 | +class OpModel(torch.nn.Module): |
| 60 | + """ |
| 61 | + Wraps a single torch operator in an nn.Module. |
| 62 | + """ |
| 63 | + |
| 64 | + def __init__( |
| 65 | + self, |
| 66 | + op: OpOverload, |
| 67 | + runtime_input_count: int, |
| 68 | + fixed_args: Sequence[Any], |
| 69 | + fixed_kwargs: dict[str, Any], |
| 70 | + ): |
| 71 | + super().__init__() |
| 72 | + self.op = op |
| 73 | + self.runtime_input_count = runtime_input_count |
| 74 | + self.fixed_kwargs = fixed_kwargs |
| 75 | + |
| 76 | + # Register parameters for fixed tensors. Some things will choke on |
| 77 | + # constant tensor weights, for example. |
| 78 | + new_args = [] |
| 79 | + for i, arg in enumerate(fixed_args): |
| 80 | + if isinstance(arg, torch.Tensor): |
| 81 | + param = torch.nn.Parameter(arg, requires_grad=False) |
| 82 | + param_name = f"arg_{i}_param" |
| 83 | + setattr(self, param_name, param) |
| 84 | + self.register_parameter(param_name, param) |
| 85 | + new_args.append(param) |
| 86 | + else: |
| 87 | + new_args.append(arg) |
| 88 | + self.fixed_args = tuple(new_args) |
| 89 | + |
| 90 | + def forward(self, *args, **kwargs): |
| 91 | + return self.op(*(args + self.fixed_args), **(kwargs | self.fixed_kwargs)) |
| 92 | + |
| 93 | + |
| 94 | +class ConvModel(OpModel): |
| 95 | + def forward(self, *args, **kwargs): |
| 96 | + weight, bias, stride, padding, dilation, transposed, output_padding, groups = ( |
| 97 | + self.fixed_args |
| 98 | + ) |
| 99 | + |
| 100 | + if not transposed: |
| 101 | + if len(weight.shape) == 3: |
| 102 | + op = torch.nn.functional.conv1d |
| 103 | + elif len(weight.shape) == 4: |
| 104 | + op = torch.nn.functional.conv2d |
| 105 | + elif len(weight.shape) == 5: |
| 106 | + op = torch.nn.functional.conv3d |
| 107 | + |
| 108 | + return op(args[0], weight, bias, stride, padding, dilation, groups) |
| 109 | + else: |
| 110 | + if len(weight.shape) == 3: |
| 111 | + op = torch.nn.functional.conv_transpose1d |
| 112 | + elif len(weight.shape) == 4: |
| 113 | + op = torch.nn.functional.conv_transpose2d |
| 114 | + elif len(weight.shape) == 5: |
| 115 | + op = torch.nn.functional.conv_transpose3d |
| 116 | + |
| 117 | + return op( |
| 118 | + args[0], weight, bias, stride, padding, output_padding, groups, dilation |
| 119 | + ) |
| 120 | + |
| 121 | + |
| 122 | +def get_module_for_op(op: OpOverload): |
| 123 | + if op == torch.ops.aten.convolution.default: |
| 124 | + return ConvModel |
| 125 | + else: |
| 126 | + return OpModel |
| 127 | + |
| 128 | + |
| 129 | +class FactoTestsBase(unittest.TestCase): |
| 130 | + def __init__(self, tester_factory: Callable[[], TesterBase], *args, **kwargs): |
| 131 | + super().__init__(*args, **kwargs) |
| 132 | + self._tester_factory = tester_factory |
| 133 | + |
| 134 | + @staticmethod |
| 135 | + def _generate_test(op_name: str) -> None: |
| 136 | + # Find the torch op with the given name. |
| 137 | + sections = op_name.split(".") |
| 138 | + torch_op = functools.reduce(getattr, sections, torch.ops.aten) |
| 139 | + |
| 140 | + test_name = "test_" + op_name.replace(".", "_") |
| 141 | + |
| 142 | + def test_body(self): |
| 143 | + self._test_op(torch_op) |
| 144 | + |
| 145 | + setattr(FactoTestsBase, test_name, test_body) |
| 146 | + |
| 147 | + @staticmethod |
| 148 | + def get_runtime_input_count(spec: Spec): |
| 149 | + # Determine which inputs are fixed at tracing time (weights, for example), |
| 150 | + # vs inputs to the runtime graph. We currently assume that the runtime graph |
| 151 | + # inputs start at the beginning of the arg list and are contiguous. |
| 152 | + # |
| 153 | + # Args are consider to be runtime inputs if they are positional and are named |
| 154 | + # one of RUNTIME_INPUT_NAMES. If none match, we assume only the first arg is a |
| 155 | + # runtime input. |
| 156 | + runtime_input_count = 0 |
| 157 | + for inspec in spec.inspec: |
| 158 | + is_runtime_input = ( |
| 159 | + inspec.type.is_tensor() and inspec.name.lower() in RUNTIME_INPUT_NAMES |
| 160 | + ) |
| 161 | + if is_runtime_input: |
| 162 | + runtime_input_count += 1 |
| 163 | + else: |
| 164 | + break |
| 165 | + |
| 166 | + return max(1, runtime_input_count) |
| 167 | + |
| 168 | + def setUp(self): |
| 169 | + torch.set_printoptions(threshold=3) |
| 170 | + |
| 171 | + def _test_op(self, op: OpOverload) -> None: # noqa |
| 172 | + random_manager.seed(0) |
| 173 | + |
| 174 | + # Strip namespace |
| 175 | + op_name = op.name().split("::")[-1] |
| 176 | + |
| 177 | + # Default to .default overload |
| 178 | + if "." not in op_name: |
| 179 | + op_name += ".default" |
| 180 | + |
| 181 | + # Find and patch op spec |
| 182 | + if op_name not in CombinedSpecDB: |
| 183 | + raise ValueError(f"Operator {op_name} not found in SpecDictDB.") |
| 184 | + spec = _patch_spec(CombinedSpecDB[op_name]) |
| 185 | + |
| 186 | + runtime_input_count = FactoTestsBase.get_runtime_input_count(spec) |
| 187 | + |
| 188 | + print(f"Op: {op_name}, {runtime_input_count} runtime inputs") |
| 189 | + |
| 190 | + # Run test cases |
| 191 | + success_count_delegated = 0 |
| 192 | + success_count_undelegated = 0 |
| 193 | + fail_count = 0 |
| 194 | + |
| 195 | + i = 0 |
| 196 | + for posargs, inkwargs, _ in ArgumentTupleGenerator(spec).gen(): |
| 197 | + i += 1 |
| 198 | + |
| 199 | + try: |
| 200 | + if isinstance(posargs[0], torch.Tensor): |
| 201 | + # Temporary for getting around XNN crashes (https://github.com/pytorch/executorch/issues/10960). |
| 202 | + # TODO Re-enable when resolved. |
| 203 | + if posargs[0].dtype in {torch.int8, torch.uint8}: |
| 204 | + print("Skipping (u)int8 case.") |
| 205 | + continue |
| 206 | + |
| 207 | + module_cls = get_module_for_op(op) |
| 208 | + model = module_cls( |
| 209 | + op, runtime_input_count, posargs[runtime_input_count:], inkwargs |
| 210 | + ) |
| 211 | + |
| 212 | + # Sanity check to make sure it runs in eager. This can present nicer error |
| 213 | + # messages sometimes compared to tracing. |
| 214 | + try: |
| 215 | + model(*posargs[:runtime_input_count]) |
| 216 | + except Exception as e: |
| 217 | + print(f"Eager execution failed: {e}") |
| 218 | + continue |
| 219 | + |
| 220 | + tester = self._tester_factory( |
| 221 | + model, tuple(posargs[:runtime_input_count]) |
| 222 | + ) |
| 223 | + |
| 224 | + # Dynamo will also fail to handle some patterns that are valid in eager. |
| 225 | + try: |
| 226 | + tester.export() |
| 227 | + except Exception: |
| 228 | + print("Export failed.") |
| 229 | + continue |
| 230 | + |
| 231 | + tester.to_edge_transform_and_lower() |
| 232 | + |
| 233 | + is_delegated = any( |
| 234 | + n.target == torch._higher_order_ops.executorch_call_delegate |
| 235 | + for n in tester.stages[tester.cur].graph_module.graph.nodes |
| 236 | + if n.op == "call_function" |
| 237 | + ) |
| 238 | + |
| 239 | + # Only run the runtime test if the op was delegated. |
| 240 | + if is_delegated: |
| 241 | + ( |
| 242 | + tester.to_executorch() |
| 243 | + .serialize() |
| 244 | + .run_method_and_compare_outputs() |
| 245 | + ) |
| 246 | + |
| 247 | + if is_delegated: |
| 248 | + success_count_delegated += 1 |
| 249 | + else: |
| 250 | + success_count_undelegated += 1 |
| 251 | + except Exception: |
| 252 | + fail_count += 1 |
| 253 | + print("Args:") |
| 254 | + for arg in posargs: |
| 255 | + if isinstance(arg, torch.Tensor): |
| 256 | + print(f" {arg.dtype} {arg.shape}") |
| 257 | + else: |
| 258 | + print(f" {arg}") |
| 259 | + |
| 260 | + traceback.print_exc() |
| 261 | + |
| 262 | + print( |
| 263 | + f"{success_count_delegated + success_count_undelegated} PASS, {fail_count} FAIL" |
| 264 | + ) |
| 265 | + print( |
| 266 | + f" {success_count_delegated} DELEGATED, {success_count_undelegated} UNDELEGATED" |
| 267 | + ) |
| 268 | + |
| 269 | + |
| 270 | +# Programatically generate tests for each operator. |
| 271 | +for op_name in CombinedSpecDB.keys(): |
| 272 | + FactoTestsBase._generate_test(op_name) |
| 273 | + |
| 274 | + |
| 275 | +# TODO Figure out where to put these |
| 276 | +class FactoTestsXNNPACK(FactoTestsBase): |
| 277 | + def __init__(self, *args, **kwargs): |
| 278 | + super().__init__(XnnpackTester, *args, **kwargs) |
| 279 | + |
| 280 | + |
| 281 | +try: |
| 282 | + from executorch.backends.apple.coreml.test.tester import CoreMLTester |
| 283 | + |
| 284 | + class FactoTestsCoreML(FactoTestsBase): |
| 285 | + def __init__(self, *args, **kwargs): |
| 286 | + super().__init__(CoreMLTester, *args, **kwargs) |
| 287 | + |
| 288 | +except: |
| 289 | + print("Skipping Core ML facto tests as Core ML AOT is not available.") |
0 commit comments