InputGen: introduce argument attributes

Reviewed By: SS-JIA

Differential Revision: D52047184

fbshipit-source-id: 4a7e81a7965f791ffe6c793235f59c671730f8df

Author: manuelcandales

inputgen/attribute/__init__.py

@@ -0,0 +1,133 @@
+# 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 enum import Enum
+from typing import List, Optional, Tuple
+
+import torch
+from inputgen.argument.type import ArgType
+from inputgen.variable.type import ScalarDtype
+
+
+class Attribute(str, Enum):
+    OPTIONAL = "optional"
+    LENGTH = "len"
+    DTYPE = "dtype"
+    RANK = "rank"
+    SIZE = "size"
+    VALUE = "value"
+
+    @staticmethod
+    def hierarchy(argtype: ArgType) -> List["Attribute"]:
+        if argtype.is_tensor_list():
+            if argtype == ArgType.TensorOptList:
+                return [
+                    Attribute.LENGTH,
+                    Attribute.OPTIONAL,
+                    Attribute.DTYPE,
+                    Attribute.RANK,
+                    Attribute.SIZE,
+                    Attribute.VALUE,
+                ]
+            else:
+                return [
+                    Attribute.LENGTH,
+                    Attribute.DTYPE,
+                    Attribute.RANK,
+                    Attribute.SIZE,
+                    Attribute.VALUE,
+                ]
+        opt = [Attribute.OPTIONAL] if argtype.is_optional() else []
+        if argtype.is_tensor():
+            return opt + [
+                Attribute.DTYPE,
+                Attribute.RANK,
+                Attribute.SIZE,
+                Attribute.VALUE,
+            ]
+        elif argtype.is_scalar():
+            return opt + [Attribute.DTYPE, Attribute.VALUE]
+        elif argtype.is_list():
+            return opt + [Attribute.LENGTH, Attribute.VALUE]
+        else:
+            return opt + [Attribute.VALUE]
+
+    def get_vtype(
+        self,
+        argtype: Optional[ArgType] = None,
+        scalar_dtype: Optional[ScalarDtype] = None,
+    ) -> type:
+        if self == Attribute.OPTIONAL:
+            return bool
+        if self == Attribute.DTYPE:
+            if argtype is None:
+                raise ValueError(f"Attribute {self} requires an argtype")
+            if argtype.is_scalar():
+                assert isinstance(ScalarDtype, type)
+                return ScalarDtype
+            else:
+                assert isinstance(torch.dtype, type)
+                return torch.dtype
+        if self in [Attribute.LENGTH, Attribute.RANK, Attribute.SIZE]:
+            return int
+        if self == Attribute.VALUE:
+            if argtype is None:
+                raise ValueError(f"Attribute {self} requires an argtype")
+            if argtype.has_integer_value():
+                return int
+            if argtype.is_bool():
+                return bool
+            if argtype.is_float():
+                return float
+            if argtype.is_string():
+                return str
+            if argtype.is_memory_format():
+                return str
+            if argtype.is_scalar():
+                if scalar_dtype is None:
+                    raise ValueError(
+                        "Attribute value for argtype scalar requires a scalar_dtype"
+                    )
+                assert isinstance(scalar_dtype, ScalarDtype)
+                assert isinstance(scalar_dtype.value, type)
+                return scalar_dtype.value
+            if argtype.is_scalar_type():
+                assert isinstance(torch.dtype, type)
+                return torch.dtype
+        return float
+
+    def get_custom_limits(
+        self, argtype: Optional[ArgType] = None
+    ) -> Optional[Tuple[int, int]]:
+        RANK_MAX = 6
+        SIZE_MAX = 8
+        TL_LEN_MAX = 6
+        LIST_LEN_MAX = 8
+        VALUE_LENGTH_MIN = -9
+        VALUE_LENGTH_MAX = 9
+        VALUE_MIN = -20
+        VALUE_MAX = 20
+
+        if self == Attribute.LENGTH:
+            if argtype is None:
+                raise ValueError(f"Attribute {self} requires an argtype")
+            if argtype.is_tensor_list():
+                return (0, TL_LEN_MAX)
+            if argtype.is_shape():
+                return (0, RANK_MAX)
+            return (0, LIST_LEN_MAX)
+        elif self == Attribute.RANK:
+            return (0, RANK_MAX)
+        elif self == Attribute.SIZE:
+            return (0, SIZE_MAX)
+        elif self == Attribute.VALUE:
+            if argtype is None:
+                raise ValueError(f"Attribute {self} requires an argtype")
+            if argtype.is_shape():
+                return (-SIZE_MAX, SIZE_MAX)
+            if argtype.is_length() or argtype.is_length_list():
+                return (VALUE_LENGTH_MIN, VALUE_LENGTH_MAX)
+        return None

inputgen/attribute/model.py

@@ -0,0 +1,78 @@
+# 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.argument.type import ArgType
+from inputgen.attribute.model import Attribute
+from inputgen.variable.type import ScalarDtype
+
+
+class TestAttribute(unittest.TestCase):
+    def test_hierarchy(self):
+        self.assertEqual(
+            Attribute.hierarchy(ArgType.ScalarOpt),
+            [Attribute.OPTIONAL, Attribute.DTYPE, Attribute.VALUE],
+        )
+        self.assertEqual(
+            Attribute.hierarchy(ArgType.TensorOpt),
+            [
+                Attribute.OPTIONAL,
+                Attribute.DTYPE,
+                Attribute.RANK,
+                Attribute.SIZE,
+                Attribute.VALUE,
+            ],
+        )
+
+    def test_vtype(self):
+        attr = Attribute.OPTIONAL
+        self.assertEqual(attr.get_vtype(), bool)
+
+        attr = Attribute.LENGTH
+        self.assertEqual(attr.get_vtype(), int)
+
+        attr = Attribute.RANK
+        self.assertEqual(attr.get_vtype(), int)
+
+        attr = Attribute.SIZE
+        self.assertEqual(attr.get_vtype(), int)
+
+        attr = Attribute.DTYPE
+        self.assertEqual(attr.get_vtype(ArgType.Tensor), torch.dtype)
+
+        attr = Attribute.DTYPE
+        self.assertEqual(attr.get_vtype(ArgType.Scalar), ScalarDtype)
+
+        attr = Attribute.VALUE
+        self.assertEqual(attr.get_vtype(ArgType.Dim), int)
+
+        attr = Attribute.VALUE
+        self.assertEqual(attr.get_vtype(ArgType.Float), float)
+
+        attr = Attribute.VALUE
+        self.assertEqual(attr.get_vtype(ArgType.String), str)
+
+        attr = Attribute.VALUE
+        self.assertEqual(attr.get_vtype(ArgType.ScalarType), torch.dtype)
+
+        attr = Attribute.VALUE
+        self.assertEqual(attr.get_vtype(ArgType.Scalar, ScalarDtype.bool), bool)
+
+        attr = Attribute.VALUE
+        self.assertEqual(attr.get_vtype(ArgType.Scalar, ScalarDtype.int), int)
+
+        attr = Attribute.VALUE
+        self.assertEqual(attr.get_vtype(ArgType.Scalar, ScalarDtype.float), float)
+
+    def test_custom_limits(self):
+        attr = Attribute.OPTIONAL
+        self.assertEqual(attr.get_custom_limits(), None)
+
+
+if __name__ == "__main__":
+    unittest.main()