InputGen: introduce structural engine

Reviewed By: SS-JIA

Differential Revision: D52047181

fbshipit-source-id: 925bbcd95f969dd1373f9df8e3de192f245d55b3

Author: manuelcandales

inputgen/argument/engine.py

@@ -0,0 +1,95 @@
+# 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 random
+from typing import Any, List, Optional
+
+from inputgen.argument.type import ArgType
+from inputgen.attribute.engine import AttributeEngine
+from inputgen.attribute.model import Attribute
+from inputgen.specs.model import Constraint
+
+
+class StructuralEngine:
+    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
+        self.hierarchy = StructuralEngine.hierarchy(argtype)
+
+    @staticmethod
+    def hierarchy(argtype) -> List[Attribute]:
+        """Return the structural hierarchy for a given argument type"""
+        if argtype.is_tensor_list():
+            return [Attribute.LENGTH, Attribute.RANK, Attribute.SIZE]
+        elif argtype.is_tensor():
+            return [Attribute.RANK, Attribute.SIZE]
+        elif argtype.is_list():
+            return [Attribute.LENGTH, Attribute.VALUE]
+        else:
+            return [Attribute.VALUE]
+
+    def gen_structure_with_depth_and_length(
+        self, depth: int, length: int, focus: Attribute
+    ):
+        if length == 0:
+            yield ()
+            return
+
+        attr = self.hierarchy[-(depth + 1)]
+        focus_ixs = range(length) if focus == attr else (random.choice(range(length)),)
+        for focus_ix in focus_ixs:
+            values = [()]
+            for ix in range(length):
+                if ix == focus_ix:
+                    elements = self.gen_structure_with_depth(depth, focus, length, ix)
+                else:
+                    elements = self.gen_structure_with_depth(depth, None, length, ix)
+                new_values = []
+                for elem in elements:
+                    new_values += [t + (elem,) for t in values]
+                values = new_values
+            yield from values
+
+    def gen_structure_with_depth(
+        self,
+        depth: int,
+        focus: Attribute,
+        length: Optional[int] = None,
+        ix: Optional[int] = None,
+    ):
+        attr = self.hierarchy[-(depth + 1)]
+
+        if ix is not None:
+            args = (self.deps, length, ix)
+        elif length is not None:
+            args = (
+                self.deps,
+                length,
+            )
+        else:
+            args = (self.deps,)
+
+        values = AttributeEngine(attr, self.constraints, self.valid, self.argtype).gen(
+            focus, *args
+        )
+
+        for v in values:
+            if depth == 0:
+                yield v
+            else:
+                yield from self.gen_structure_with_depth_and_length(depth - 1, v, focus)
+
+    def gen(self, focus: Attribute):
+        depth = len(self.hierarchy) - 1
+        yield from self.gen_structure_with_depth(depth, focus)

test/inputgen/test_structural_engine.py

@@ -0,0 +1,32 @@
+# 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 StructuralEngine
+from inputgen.argument.type import ArgType
+from inputgen.attribute.model import Attribute
+from inputgen.specs.model import ConstraintProducer as cp
+
+
+class TestStructuralEngine(unittest.TestCase):
+    def test_engine(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: 0),
+        ]
+        deps = [2]
+
+        engine = StructuralEngine(ArgType.Tensor, constraints, deps, True)
+        for s in engine.gen(Attribute.VALUE):
+            self.assertTrue(1 <= len(s) <= 4)
+            self.assertTrue(all(v in [2, 4, 5] for v in s))
+
+
+if __name__ == "__main__":
+    unittest.main()