InputGen: Introduce specs

Reviewed By: SS-JIA

Differential Revision: D52308291

fbshipit-source-id: 4d4cb8d54ea6585e8e7a6403a700a1f243f3e17c

Author: manuelcandales

inputgen/specs/model.py

@@ -6,8 +6,9 @@
 
 from dataclasses import dataclass
 from enum import Enum
-from typing import Callable
+from typing import Callable, List, Optional
 
+from inputgen.argument.type import ArgType
 from inputgen.attribute.model import Attribute
 
 
@@ -63,3 +64,86 @@ class ConstraintProducer:
     Rank = ConstraintAttributeSuffixes(Attribute.RANK)
     Size = ConstraintAttributeSuffixes(Attribute.SIZE)
     Value = ConstraintAttributeSuffixes(Attribute.VALUE)
+
+
+class BaseArg:
+    def __init__(
+        self,
+        argtype: ArgType,
+        name: str,
+        deps: Optional[List[int]] = None,
+        constraints: Optional[List[Constraint]] = None,
+    ):
+        self.name: str = name
+        self.type: ArgType = argtype
+        self.deps = () if deps is None else tuple(deps)
+        self.constraints = [] if constraints is None else constraints
+        self._kw: bool = False
+        self._out: bool = False
+        self._ret: bool = False
+
+    @property
+    def kw(self):
+        return self._kw
+
+    @kw.setter
+    def kw(self, v):
+        if not isinstance(v, bool):
+            raise ValueError("kw property should be boolean")
+        self._kw = v
+
+    @property
+    def out(self):
+        return self._out
+
+    @out.setter
+    def out(self, v):
+        if not isinstance(v, bool):
+            raise ValueError("out property should be boolean")
+        self._out = v
+
+    @property
+    def ret(self):
+        return self._ret
+
+    @ret.setter
+    def ret(self, v):
+        if not isinstance(v, bool):
+            raise ValueError("ret property should be boolean")
+        self._ret = v
+
+
+class InArg(BaseArg):
+    def __init__(self, *args, **kwargs):
+        BaseArg.__init__(self, *args, **kwargs)
+
+
+class InPosArg(InArg):
+    def __init__(self, *args, **kwargs):
+        BaseArg.__init__(self, *args, **kwargs)
+
+
+class InKwArg(InArg):
+    def __init__(self, *args, **kwargs):
+        BaseArg.__init__(self, *args, **kwargs)
+        self._kw = True
+
+
+class OutArg(BaseArg):
+    def __init__(self, argtype: ArgType, name: str = "out", *args, **kwargs):
+        BaseArg.__init__(self, argtype, name, *args, **kwargs)
+        self._kw = True
+        self._out = True
+
+
+class Return(BaseArg):
+    def __init__(self, argtype: ArgType, name: str = "__ret", *args, **kwargs):
+        BaseArg.__init__(self, argtype, name, *args, **kwargs)
+        self._ret = True
+
+
+class Spec:
+    def __init__(self, op: str, inspec: List[InArg], outspec: List[OutArg]):
+        self.op = op
+        self.inspec = inspec
+        self.outspec = outspec

test/inputgen/test_specs.py

@@ -0,0 +1,48 @@
+# 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.type import ArgType
+from inputgen.specs.model import InKwArg, InPosArg, OutArg, Return
+
+
+class TestArgSpecs(unittest.TestCase):
+    def test_inpos(self):
+        arg = InPosArg(ArgType.Tensor, name="self")
+        self.assertEqual(arg.name, "self")
+        self.assertEqual(arg.type, ArgType.Tensor)
+        self.assertFalse(arg.kw)
+        self.assertFalse(arg.out)
+        self.assertFalse(arg.ret)
+
+    def test_inkw(self):
+        arg = InKwArg(ArgType.Scalar, name="alpha")
+        self.assertEqual(arg.name, "alpha")
+        self.assertEqual(arg.type, ArgType.Scalar)
+        self.assertTrue(arg.kw)
+        self.assertFalse(arg.out)
+        self.assertFalse(arg.ret)
+
+    def test_out(self):
+        arg = OutArg(ArgType.TensorList)
+        self.assertEqual(arg.name, "out")
+        self.assertEqual(arg.type, ArgType.TensorList)
+        self.assertTrue(arg.kw)
+        self.assertTrue(arg.out)
+        self.assertFalse(arg.ret)
+
+    def test_ret(self):
+        arg = Return(ArgType.Tensor)
+        self.assertEqual(arg.name, "__ret")
+        self.assertEqual(arg.type, ArgType.Tensor)
+        self.assertFalse(arg.kw)
+        self.assertFalse(arg.out)
+        self.assertTrue(arg.ret)
+
+
+if __name__ == "__main__":
+    unittest.main()