InputGen: introduce argument tuple generator

Reviewed By: SS-JIA

Differential Revision: D52308290

fbshipit-source-id: c4e50f7f13e574e0c32f797d9d4f5e4e298751be

Author: manuelcandales

inputgen/argtuple/gen.py

@@ -0,0 +1,43 @@
+# 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.
+
+from typing import Any, List, OrderedDict, Tuple
+
+from inputgen.argtuple.engine import MetaArgTupleEngine
+from inputgen.argument.engine import MetaArg
+from inputgen.argument.gen import ArgumentGenerator
+from inputgen.specs.model import Spec
+
+
+class ArgumentTupleGenerator:
+    def __init__(self, spec: Spec):
+        self.spec = spec
+
+    def gen_tuple(
+        self, meta_tuple: Tuple[MetaArg], *, out: bool = False
+    ) -> Tuple[List[Any], OrderedDict[str, Any]]:
+        args = []
+        kwargs = OrderedDict()
+        for ix, arg in enumerate(self.spec.inspec):
+            m = meta_tuple[ix]
+            val = ArgumentGenerator(m).gen()
+            if arg.kw:
+                kwargs[arg.name] = val
+            else:
+                args.append(val)
+        if out:
+            for ix, arg in enumerate(self.spec.outspec):
+                m = meta_tuple[ix + len(self.spec.inspec)]
+                val = ArgumentGenerator(m).gen()
+                kwargs[arg.name] = val
+        return args, kwargs
+
+    def gen(
+        self, *, valid: bool = True, out: bool = False
+    ) -> Tuple[List[Any], OrderedDict[str, Any]]:
+        engine = MetaArgTupleEngine(self.spec, out=out)
+        for meta_tuple in engine.gen(valid=valid):
+            yield self.gen_tuple(meta_tuple, out=out)

test/inputgen/test_argtuple_generator.py

@@ -0,0 +1,57 @@
+# 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
+
+import torch
+from inputgen.argtuple.gen import ArgumentTupleGenerator
+from inputgen.argument.type import ArgType
+from inputgen.specs.model import ConstraintProducer as cp, InPosArg, Spec
+
+
+class TestArgumentTupleGenerator(unittest.TestCase):
+    def test_gen(self):
+        spec = Spec(
+            op="test_size",  # (Tensor self, int dim) -> int
+            inspec=[
+                InPosArg(ArgType.Tensor, name="self"),
+                InPosArg(
+                    ArgType.Dim,
+                    name="dim",
+                    deps=[0],
+                    constraints=[
+                        cp.Value.Ge(
+                            lambda deps: -deps[0].dim() if deps[0].dim() > 0 else None
+                        ),
+                        cp.Value.Ge(lambda deps: -1 if deps[0].dim() == 0 else None),
+                        cp.Value.Le(
+                            lambda deps: deps[0].dim() - 1
+                            if deps[0].dim() > 0
+                            else None
+                        ),
+                        cp.Value.Le(lambda deps: 0 if deps[0].dim() == 0 else None),
+                    ],
+                ),
+            ],
+            outspec=[],
+        )
+
+        for args, kwargs in ArgumentTupleGenerator(spec).gen():
+            self.assertEqual(len(args), 2)
+            self.assertEqual(kwargs, {})
+            t = args[0]
+            dim = args[1]
+            self.assertTrue(isinstance(t, torch.Tensor))
+            self.assertTrue(isinstance(dim, int))
+            if t.dim() == 0:
+                self.assertTrue(dim in [-1, 0])
+            else:
+                self.assertTrue(dim >= -t.dim())
+                self.assertTrue(dim < t.dim())
+
+
+if __name__ == "__main__":
+    unittest.main()