InputGen: Add support for out tensor arguments

Summary:
Enables out tensor generation, for out variants.
In the case of out tensors, only dtype information is relevant, since shape is determined by inputs.

Reviewed By: SS-JIA

Differential Revision: D52919719

fbshipit-source-id: 20ad8c59bbdf1e0d30243f672ed09fa4fe2e94f1

Author: manuelcandales

inputgen/argtuple/engine.py

@@ -39,8 +39,11 @@ def inverse_permutation(permutation):
 class MetaArgTupleEngine:
     def __init__(self, spec: Spec, out: bool = False):
         if out:
-            raise NotImplementedError("out=True is not supported yet")
-        self.args = spec.inspec
+            for arg in spec.outspec:
+                arg.deps = list(range(len(spec.inspec)))
+            self.args = spec.inspec + spec.outspec
+        else:
+            self.args = spec.inspec
         self.order = self._sort_dependencies()
         self.order_inverse_perm = inverse_permutation(self.order)
 
@@ -77,7 +80,9 @@ def gen_meta_tuples(self, valid: bool, focus_ix: int):
             for focus in focuses:
                 for meta_tuple in tuples:
                     deps = self._get_deps(meta_tuple, arg.deps)
-                    engine = MetaArgEngine(arg.type, arg.constraints, deps, valid)
+                    engine = MetaArgEngine(
+                        arg.out, arg.type, arg.constraints, deps, valid
+                    )
                     for meta_arg in engine.gen(focus):
                         new_tuples.append(meta_tuple + (meta_arg,))
             tuples = new_tuples
@@ -98,7 +103,7 @@ def gen_invalid_from_valid(self, valid_tuple):
             # Generating invalid argument {ix} {arg.type}
             deps = tuple(valid_value_tuple[i] for i in arg.deps)
             for focus in Attribute.hierarchy(arg.type):
-                engine = MetaArgEngine(arg.type, arg.constraints, deps, False)
+                engine = MetaArgEngine(arg.out, arg.type, arg.constraints, deps, False)
                 for meta_arg in engine.gen(focus):
                     invalid_tuple = (
                         valid_tuple[:ix] + (meta_arg,) + valid_tuple[ix + 1 :]

inputgen/argtuple/gen.py

@@ -19,21 +19,22 @@ def __init__(self, spec: Spec):
     def gen_tuple(
         self, meta_tuple: Tuple[MetaArg], *, out: bool = False
     ) -> Tuple[List[Any], OrderedDict[str, Any]]:
-        args = []
-        kwargs = OrderedDict()
+        posargs = []
+        inkwargs = OrderedDict()
+        outargs = OrderedDict()
         for ix, arg in enumerate(self.spec.inspec):
             m = meta_tuple[ix]
             val = ArgumentGenerator(m).gen()
             if arg.kw:
-                kwargs[arg.name] = val
+                inkwargs[arg.name] = val
             else:
-                args.append(val)
+                posargs.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
+                outargs[arg.name] = val
+        return posargs, inkwargs, outargs
 
     def gen(
         self, *, valid: bool = True, out: bool = False

inputgen/argument/engine.py

@@ -188,11 +188,13 @@ def rank(self, ix=None):
 class MetaArgEngine:
     def __init__(
         self,
+        out: bool,
         argtype: ArgType,
         constraints: List[Constraint],
         deps: List[Any],
         valid: bool,
     ):
+        self.out = out
         self.argtype = argtype
         self.constraints = constraints
         self.deps = deps
@@ -244,6 +246,21 @@ def gen_value_spaces(self, focus, dtype, struct):
     def gen(self, focus):
         # TODO(mcandales): Enable Tensor List generation
 
+        if self.out:
+            if self.argtype.is_tensor():
+                if focus in [None, Attribute.DTYPE]:
+                    struct = (0,)
+                    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
+                            )
+                return
+            elif self.argtype.is_tensor_list():
+                raise NotImplementedError("Tensor List output not implemented yet")
+            else:
+                raise ValueError("Output argtype must be tensor or tensor list")
+
         if focus in [None, Attribute.OPTIONAL]:
             if self.argtype.is_optional() and self.gen_optional():
                 yield MetaArg(self.argtype, optional=True)

test/inputgen/test_argtuple_generator.py

@@ -39,11 +39,12 @@ def test_gen(self):
             outspec=[],
         )
 
-        for args, kwargs in ArgumentTupleGenerator(spec).gen():
-            self.assertEqual(len(args), 2)
-            self.assertEqual(kwargs, {})
-            t = args[0]
-            dim = args[1]
+        for posargs, inkwargs, outargs in ArgumentTupleGenerator(spec).gen():
+            self.assertEqual(len(posargs), 2)
+            self.assertEqual(inkwargs, {})
+            self.assertEqual(outargs, {})
+            t = posargs[0]
+            dim = posargs[1]
             self.assertTrue(isinstance(t, torch.Tensor))
             self.assertTrue(isinstance(dim, int))
             if t.dim() == 0: