Use sample input for edge ops (#291)

Summary:

Use sample inputs to generate dtype constraints for edge ops.

This is depending on diff stack D49088856. The purpose of the whole stack is to increase edge ops dtype constraints coverage.

This diff uses the sample input generator DtypeRunner to generate all possible dtype arguments and get back valid ones for each operator.

Reviewed By: manuelcandales

Differential Revision: D49182359

Author: larryliu0820

exir/dialects/edge/_ops.py

@@ -87,15 +87,15 @@ def __init__(
             alias: AllowedDtypeSet(set(types)) for alias, types in type_alias.items()
         }
         self.type_constraint: List[Dict[str, str]] = type_constraint
-        # type_constraint's non return entries should be same as all tensor args.
+        # type_constraint's non return entries should include all tensor-like arguments.
         for t_constraint in self.type_constraint:
             type_constraint_names = set(t_constraint)
             all_tensor_arg_names = set(
                 self.essential_tensor_io_names + self.optional_tensor_io_names
             )
-            if type_constraint_names != all_tensor_arg_names:
+            if not all_tensor_arg_names.issubset(type_constraint_names):
                 raise RuntimeError(
-                    "Each input entry of type_constraint must be tensor-like,"
+                    "Input entries of type_constraint must contain all tensor-like arguments, "
                     + f"but get {type_constraint_names} and {all_tensor_arg_names}"
                 )
 

exir/dialects/edge/arg/model.py

@@ -362,8 +362,3 @@ def to_out(self, *, name: Optional[str] = None) -> OutArg:
             value=self.value,
             deps=self.deps,
         )
-
-
-def get_callable(name):
-    main, suffix = name.split(".")
-    return getattr(getattr(torch.ops.aten, main), suffix)

exir/dialects/edge/dtype/TARGETS

@@ -7,9 +7,11 @@ python_library(
     srcs = [
         "runner.py",
         "supported.py",
+        "utils.py",
     ],
     deps = [
         "//caffe2:torch",
         "//executorch/exir/dialects/edge/arg:lib",
+        "//executorch/exir/dialects/edge/op:lib",
     ],
 )

exir/dialects/edge/dtype/runner.py

@@ -10,14 +10,10 @@
 
 import torch
 import torch.testing._internal.common_dtype as common_dtype
-from executorch.exir.dialects.edge.arg.model import (
-    ArgMode,
-    BaseArg,
-    BaseKwarg,
-    GenMode,
-    get_callable,
-)
+from executorch.exir.dialects.edge.arg.model import ArgMode, BaseArg, BaseKwarg, GenMode
 from executorch.exir.dialects.edge.arg.type import ArgType
+from executorch.exir.dialects.edge.dtype.utils import extract_return_dtype
+from executorch.exir.dialects.edge.op.api import get_callable
 
 
 class DtypeRunner:
@@ -92,12 +88,13 @@ def _get_type_tuples(
         types = DtypeRunner._get_types(inputs)
 
         def mapping(t):
+            type_dtypes = []
             if t.is_optional():
-                return [None]
-            elif t.is_scalar():
-                return self.scalar_dtypes
+                type_dtypes = [None]
+            if t.is_scalar():
+                return type_dtypes + self.scalar_dtypes
             elif t.is_scalar_type() or t.is_tensor() or t.is_tensor_list():
-                return self.tensor_dtypes
+                return type_dtypes + self.tensor_dtypes
             else:
                 raise ValueError("Type {t.name} does not have dtype")
 
@@ -142,19 +139,29 @@ def run_dtypes(
         args, kwargs = DtypeRunner._get_args_kwargs(inputs, dtypes, argmode)
         op = get_callable(name)
         try:
-            op(*args, **kwargs)
-            return (True, name, dtypes, args, kwargs)
+            res = op(*args, **kwargs)
+            ret_dtypes = ()
+            if "returns" in inputs:
+                ret_dtypes = tuple(extract_return_dtype(res, inputs["returns"]))
+            return (True, name, dtypes + ret_dtypes, args, kwargs)
+        except AssertionError as e:
+            raise RuntimeError(
+                f"opname: {name}, inputs: {inputs}, dtypes: {dtypes}, argmode {argmode}"
+            ) from e
         except Exception as e:
             if argmode == ArgMode.ONES:
                 return (False, name, dtypes, args, kwargs)
             ones_args, ones_kwargs = DtypeRunner._get_args_kwargs(
                 inputs, dtypes, ArgMode.ONES
             )
             try:
-                op(*ones_args, **ones_kwargs)
+                res = op(*args, **kwargs)
+                ret_dtypes = ()
+                if "returns" in inputs:
+                    ret_dtypes = tuple(extract_return_dtype(res, inputs["returns"]))
                 print(e)
                 print(name, dtypes, args, kwargs)
-                return (True, name, dtypes, ones_args, ones_kwargs)
+                return (True, name, dtypes + ret_dtypes, ones_args, ones_kwargs)
             except Exception:
                 return (False, name, dtypes, ones_args, ones_kwargs)
 

exir/dialects/edge/dtype/utils.py

@@ -0,0 +1,36 @@
+# 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 collections
+from typing import Any, List
+
+import torch
+from executorch.exir.dialects.edge.arg.model import BaseArg
+
+from executorch.exir.dialects.edge.arg.type import ArgType
+
+
+def extract_return_dtype(
+    returns: Any, sample_returns: List[BaseArg]
+) -> List[torch.dtype]:
+    """Extract the dtype from a return value."""
+    if not isinstance(returns, collections.abc.Sequence):
+        returns = [returns]
+    result = []
+    for ret, sample in zip(returns, sample_returns):
+        if sample.type == ArgType.TensorList or sample.type == ArgType.TensorOptList:
+            # Assuming all tensors in tensor list has the same dtype, and we only add 1 dtype to result.
+            assert (
+                ret is not None
+            ), f"Expecting non-None return value for {sample} but got None"
+            result.append(ret.dtype)
+            break
+        elif sample.type == ArgType.Tensor or sample.type == ArgType.TensorOpt:
+            assert (
+                ret is not None
+            ), f"Expecting non-None return value for {sample} but got None"
+            result.append(ret.dtype)
+    return result