Introduce tensor conditions (empty, transposed, permuted, strided) (#15)

Summary:
This PR introduces support for `TensorConfig` in FACTO's tensor generation pipeline, enabling fine-grained control over tensor properties like empty tensors, transposed layouts, permuted dimensions, and strided memory patterns. The implementation features a clean separation between transformation generation and tensor creation through a new `TensorTransformationGenerator` class.

Pull Request resolved: #15

Reviewed By: digantdesai

Differential Revision: D79295653

Pulled By: manuelcandales

fbshipit-source-id: 64e0f252f6ed2beb572f9369cb22f59b2e7894dc

Author: manuelcandales

examples/config_example.py

@@ -0,0 +1,90 @@
+# 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 torch
+from facto.inputgen.argtuple.gen import ArgumentTupleGenerator
+from facto.inputgen.utils.config import TensorConfig
+from facto.specdb.db import SpecDictDB
+
+
+def qualify_tensor(tensor):
+    order = tensor.dim_order()
+    dims = sum(i != order[i] for i in range(len(order)))
+    empty = tensor.numel() == 0
+    transposed = dims == 2
+    permuted = dims > 2
+    strided = len(tensor.storage()) - tensor.storage_offset() != tensor.numel()
+    return empty, transposed, permuted, strided
+
+
+def qualify_tensor_string(tensor):
+    empty, transposed, permuted, strided = qualify_tensor(tensor)
+    s = "E" if empty else ""
+    s += "P" if permuted else "T" if transposed else ""
+    s += "S" if strided else ""
+    return s
+
+
+def pretty_print_add_args(posargs, inkwargs, outargs):
+    return "".join(
+        [
+            "Tensor{",
+            qualify_tensor_string(posargs[0]),
+            "}(",
+            str(list(posargs[0].shape)),
+            ", ",
+            str(posargs[0].dtype)[6:],
+            ") + Tensor{",
+            qualify_tensor_string(posargs[1]),
+            "}(",
+            str(list(posargs[1].shape)),
+            ", dtype=",
+            str(posargs[0].dtype)[6:],
+            ") alpha = ",
+            str(inkwargs["alpha"]),
+        ]
+    )
+
+
+def generate_inputs():
+    spec = SpecDictDB["add.Tensor"]
+
+    config = TensorConfig(
+        empty=False,
+        transposed=False,
+        permuted=True,
+        strided=True,
+    ).set_probability(0.7)
+
+    generator = ArgumentTupleGenerator(spec, config=config)
+    for ix, tup in enumerate(generator.gen()):
+        posargs, inkwargs, outargs = tup
+        # Pretty printing the inputs and outputs
+        print(f"Tuple #{ix}: {pretty_print_add_args(posargs, inkwargs, outargs)}")
+        yield posargs, inkwargs, outargs
+
+
+def test_add_op():
+    op = torch.ops.aten.add.Tensor
+    for posargs, inkwargs, outargs in generate_inputs():
+        try:
+            op(*posargs, **inkwargs, **outargs)
+        except Exception:
+            return False
+    return True
+
+
+def main():
+    print("Testing add.Tensor with the following input tuples:")
+    success = test_add_op()
+    if success:
+        print("Success!")
+    else:
+        print("Failure!")
+
+
+if __name__ == "__main__":
+    main()

facto/inputgen/argtuple/gen.py

@@ -5,17 +5,55 @@
 # LICENSE file in the root directory of this source tree.
 
 from collections import OrderedDict
-from typing import Any, Generator, List, Tuple
+from copy import deepcopy
+from typing import Any, Generator, List, Optional, Tuple
 
 from facto.inputgen.argtuple.engine import MetaArgTupleEngine
 from facto.inputgen.argument.engine import MetaArg
 from facto.inputgen.argument.gen import ArgumentGenerator
-from facto.inputgen.specs.model import Spec
+from facto.inputgen.specs.model import ConstraintProducer as cp, Spec
+from facto.inputgen.utils.config import Condition, TensorConfig
 
 
 class ArgumentTupleGenerator:
-    def __init__(self, spec: Spec):
+    def __init__(self, spec: Spec, config: Optional[TensorConfig] = None):
         self.spec = spec
+        self.config = config
+        self._modified_spec = self._apply_config_constraints(spec, config)
+
+    def _apply_config_constraints(
+        self, spec: Spec, config: Optional[TensorConfig]
+    ) -> Spec:
+        """Apply TensorConfig constraints to the spec by modifying argument constraints."""
+
+        if config is None:
+            return spec
+
+        # Create a copy of the spec with modified constraints
+        modified_inspec = []
+        for arg in spec.inspec:
+            modified_arg = self._apply_constraints_to_arg(arg, config)
+            modified_inspec.append(modified_arg)
+
+        modified_outspec = []
+        for arg in spec.outspec:
+            modified_arg = self._apply_constraints_to_arg(arg, config)
+            modified_outspec.append(modified_arg)
+
+        return Spec(spec.op, modified_inspec, modified_outspec)
+
+    def _apply_constraints_to_arg(self, arg, config: TensorConfig):
+        """Apply config constraints to a single argument."""
+        # Create a copy of the argument with potentially modified constraints
+        modified_arg = deepcopy(arg)
+
+        # Add size constraints for tensor arguments when empty tensors are not allowed
+        if not config.is_allowed(Condition.ALLOW_EMPTY):
+            if arg.type.is_tensor() or arg.type.is_tensor_list():
+                size_constraint = cp.Size.Ge(lambda deps, r, d: 1)
+                modified_arg.constraints = modified_arg.constraints + [size_constraint]
+
+        return modified_arg
 
     def gen_tuple(
         self, meta_tuple: Tuple[MetaArg], *, out: bool = False
@@ -25,15 +63,15 @@ def gen_tuple(
         outargs = OrderedDict()
         for ix, arg in enumerate(self.spec.inspec):
             m = meta_tuple[ix]
-            val = ArgumentGenerator(m).gen()
+            val = ArgumentGenerator(m, config=self.config).gen()
             if arg.kw:
                 inkwargs[arg.name] = val
             else:
                 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()
+                val = ArgumentGenerator(m, config=self.config).gen()
                 outargs[arg.name] = val
         return posargs, inkwargs, outargs
 
@@ -42,6 +80,6 @@ def gen(
     ) -> Generator[
         Tuple[List[Any], OrderedDict[str, Any], OrderedDict[str, Any]], Any, Any
     ]:
-        engine = MetaArgTupleEngine(self.spec, out=out)
+        engine = MetaArgTupleEngine(self._modified_spec, out=out)
         for meta_tuple in engine.gen(valid=valid):
             yield self.gen_tuple(meta_tuple, out=out)