InputGen: introduce meta argument engine

Reviewed By: SS-JIA

Differential Revision: D52047182

fbshipit-source-id: c9e2320bdb2dad23f8b0b1defca9ffdf11c7ff99

Author: manuelcandales

inputgen/argument/engine.py

@@ -5,12 +5,15 @@
 # LICENSE file in the root directory of this source tree.
 
 import random
-from typing import Any, List, Optional
+from typing import Any, List, Optional, Tuple, Union
 
+import torch
 from inputgen.argument.type import ArgType
 from inputgen.attribute.engine import AttributeEngine
 from inputgen.attribute.model import Attribute
-from inputgen.specs.model import Constraint
+from inputgen.attribute.solve import AttributeSolver
+from inputgen.specs.model import Constraint, ConstraintSuffix
+from inputgen.variable.type import ScalarDtype
 
 
 class StructuralEngine:
@@ -27,6 +30,11 @@ def __init__(
         self.valid = valid
         self.hierarchy = StructuralEngine.hierarchy(argtype)
 
+        self.gen_list_mode = set()
+        for constraint in constraints:
+            if constraint.suffix == ConstraintSuffix.GEN:
+                self.gen_list_mode.add(constraint.attribute)
+
     @staticmethod
     def hierarchy(argtype) -> List[Attribute]:
         """Return the structural hierarchy for a given argument type"""
@@ -47,6 +55,11 @@ def gen_structure_with_depth_and_length(
             return
 
         attr = self.hierarchy[-(depth + 1)]
+
+        if attr in self.gen_list_mode:
+            yield from self.gen_structure_with_depth(depth, focus, length)
+            return
+
         focus_ixs = range(length) if focus == attr else (random.choice(range(length)),)
         for focus_ix in focus_ixs:
             values = [()]
@@ -93,3 +106,167 @@ def gen_structure_with_depth(
     def gen(self, focus: Attribute):
         depth = len(self.hierarchy) - 1
         yield from self.gen_structure_with_depth(depth, focus)
+
+
+class MetaArg:
+    def __init__(
+        self,
+        argtype: ArgType,
+        *,
+        optional: bool = False,
+        dtype: Optional[
+            Union[torch.dtype, List[Optional[torch.dtype]], ScalarDtype]
+        ] = None,
+        structure: Optional[Tuple] = None,
+        value: Optional[Any] = None,
+    ):
+        self.argtype = argtype
+        self.optional = optional
+        self.dtype = dtype
+        self.structure = structure
+        self.value = value
+
+        if not self.argtype.is_optional() and self.optional:
+            raise ValueError("Only optional argtypes can have optional instances")
+
+        if self.argtype.is_tensor_list():
+            if len(self.structure) != len(self.dtype):
+                raise ValueError(
+                    "Structure and dtype must be same length when tensor list"
+                )
+            if self.argtype == ArgType.TensorList and any(
+                d is None for d in self.dtype
+            ):
+                raise ValueError("Only TensorOptList can have None in list of dtypes")
+
+        if not self.optional and Attribute.DTYPE not in Attribute.hierarchy(
+            self.argtype
+        ):
+            if argtype.is_list():
+                self.value = list(self.structure)
+            else:
+                self.value = self.structure
+
+    def __str__(self):
+        if self.optional:
+            strval = "None"
+        elif self.argtype.is_tensor_list():
+            strval = (
+                "["
+                + ", ".join(
+                    [
+                        f"{self.dtype[i]} {self.structure[i]}"
+                        for i in range(len(self.dtype))
+                    ]
+                )
+                + "]"
+            )
+        elif self.argtype.is_tensor():
+            strval = f"{self.dtype} {self.structure}"
+        else:
+            strval = str(self.value)
+        return f"{self.argtype} {strval}"
+
+    def length(self):
+        if self.argtype.is_list():
+            return len(self.structure)
+        else:
+            return None
+
+    def rank(self, ix=None):
+        if self.argtype.is_tensor():
+            return len(self.structure)
+        elif self.argtype.is_tensor_list():
+            if ix is None:
+                return (len(s) for s in self.structure)
+            else:
+                return len(self.structure[ix])
+        else:
+            return None
+
+
+class MetaArgEngine:
+    def __init__(
+        self,
+        argtype: ArgType,
+        constraints: List[Constraint],
+        deps: List[Any],
+        valid: bool,
+    ):
+        self.argtype = argtype
+        self.constraints = constraints
+        self.deps = deps
+        self.valid = valid
+
+    def gen_structures(self, focus):
+        if self.argtype.is_scalar():
+            yield None
+        else:
+            yield from StructuralEngine(
+                self.argtype, self.constraints, self.deps, self.valid
+            ).gen(focus)
+
+    def gen_dtypes(self, focus):
+        if Attribute.DTYPE not in Attribute.hierarchy(self.argtype):
+            return {None}
+        engine = AttributeEngine(
+            Attribute.DTYPE, self.constraints, self.valid, self.argtype
+        )
+        if self.argtype.is_scalar() and focus == Attribute.VALUE:
+            # if focused on a scalar value, must generate all dtypes too
+            focus = Attribute.DTYPE
+        return engine.gen(focus, self.deps)
+
+    def gen_optional(self):
+        engine = AttributeEngine(
+            Attribute.OPTIONAL, self.constraints, self.valid, self.argtype
+        )
+        return True in engine.gen(Attribute.OPTIONAL, self.deps)
+
+    def gen_scalars(self, scalar_dtype, focus):
+        engine = AttributeEngine(
+            Attribute.VALUE, self.constraints, self.valid, self.argtype, scalar_dtype
+        )
+        return engine.gen(focus, self.deps, scalar_dtype)
+
+    def gen_value_spaces(self, focus, dtype, struct):
+        if not self.argtype.is_tensor() and not self.argtype.is_tensor_list():
+            return [None]
+        solver = AttributeSolver(Attribute.VALUE, self.argtype)
+        variables = list(
+            solver.solve(self.constraints, focus, self.valid, self.deps, dtype, struct)
+        )
+        if focus == Attribute.VALUE:
+            return [v.space for v in variables]
+        else:
+            return [random.choice(variables).space]
+
+    def gen(self, focus):
+        # TODO(mcandales): Enable Tensor List generation
+
+        if focus in [None, Attribute.OPTIONAL]:
+            if self.argtype.is_optional() and self.gen_optional():
+                yield MetaArg(self.argtype, optional=True)
+            if focus == Attribute.OPTIONAL:
+                return
+
+        if self.argtype.is_scalar():
+            scalar_dtypes = self.gen_dtypes(focus)
+            for scalar_dtype in scalar_dtypes:
+                for value in self.gen_scalars(scalar_dtype, focus):
+                    yield MetaArg(self.argtype, dtype=scalar_dtype, value=value)
+        else:
+            if focus == Attribute.DTYPE:
+                for dtype in self.gen_dtypes(focus):
+                    for struct in self.gen_structures(focus):
+                        for space in self.gen_value_spaces(focus, dtype, struct):
+                            yield MetaArg(
+                                self.argtype, dtype=dtype, structure=struct, value=space
+                            )
+            else:
+                for struct in self.gen_structures(focus):
+                    for dtype in self.gen_dtypes(focus):
+                        for space in self.gen_value_spaces(focus, dtype, struct):
+                            yield MetaArg(
+                                self.argtype, dtype=dtype, structure=struct, value=space
+                            )

test/inputgen/test_meta_arg_engine.py

@@ -0,0 +1,68 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+from inputgen.argument.engine import MetaArgEngine
+from inputgen.argument.type import ArgType
+from inputgen.attribute.model import Attribute
+from inputgen.specs.model import ConstraintProducer as cp
+from inputgen.variable.type import SUPPORTED_TENSOR_DTYPES
+
+
+class TestMetaArgEngine(unittest.TestCase):
+    def test_tensor(self):
+        constraints = [
+            cp.Rank.Le(lambda deps: deps[0] + 2),
+            cp.Size.NotIn(lambda deps, length, ix: [1, 3]),
+            cp.Size.Le(lambda deps, length, ix: 5),
+            cp.Value.Ne(lambda deps, dtype, struct: 0),
+        ]
+        deps = [2]
+
+        engine = MetaArgEngine(ArgType.Tensor, constraints, deps, True)
+        ms = list(engine.gen(Attribute.DTYPE))
+        self.assertEqual(len(ms), len(SUPPORTED_TENSOR_DTYPES))
+        self.assertEqual({m.dtype for m in ms}, set(SUPPORTED_TENSOR_DTYPES))
+        self.assertTrue(all(0 <= m.rank() <= 4 for m in ms))
+        for m in ms:
+            self.assertTrue(
+                all(0 <= size <= 5 and size not in [1, 3] for size in m.structure)
+            )
+        for m in ms:
+            self.assertEqual(str(m.value), "[-inf, 0.0) (0.0, inf]")
+
+        ms = list(engine.gen(Attribute.RANK))
+        self.assertEqual(len(ms), 4)
+        ranks = {len(m.structure) for m in ms}
+        self.assertTrue(0 in ranks)
+        self.assertTrue(4 in ranks)
+        self.assertTrue(all(0 <= r <= 4 for r in ranks))
+
+    def test_dim_list(self):
+        constraints = [
+            cp.Length.Le(lambda deps: deps[0] + deps[1]),
+            cp.Value.Gen(
+                lambda deps, length: ({(deps[0],) * length}, {(deps[1],) * length})
+            ),
+        ]
+        deps = [2, 3]
+
+        engine = MetaArgEngine(ArgType.DimList, constraints, deps, True)
+        ms = list(engine.gen(Attribute.VALUE))
+        self.assertEqual(len(ms), 1)
+        self.assertTrue(1 <= len(ms[0].value) <= 5)
+        self.assertTrue(all(v == 2 for v in ms[0].value))
+
+        engine = MetaArgEngine(ArgType.DimList, constraints, deps, False)
+        ms = list(engine.gen(Attribute.VALUE))
+        self.assertEqual(len(ms), 1)
+        self.assertTrue(1 <= len(ms[0].value) <= 5)
+        self.assertTrue(all(v == 3 for v in ms[0].value))
+
+
+if __name__ == "__main__":
+    unittest.main()