Rename positional and kwarg_only to have flat prefix (#49042)

ezyang · facebook-github-bot · commit 267641a245ce · 2020-12-10T18:20:43.000-08:00
Summary: Pull Request resolved: #49042 I want the names positional and kwarg_only to give the unflat representation (e.g., preserving TensorOptionsArguments in the returned Union). So I regret my original naming choice when I moved grouping to model. This renames them to have flat_ prefix and also adds a flat_non_out argument for cases where you just want to look at non-out arguments. Signed-off-by: Edward Z. Yang <ezyang@fb.com> Test Plan: Imported from OSS Reviewed By: H-Huang Differential Revision: D25455884 Pulled By: ezyang fbshipit-source-id: f923f8881267a3e3e8e9521519412f7cc25034fc
diff --git a/tools/autograd/gen_annotated_fn_args.py b/tools/autograd/gen_annotated_fn_args.py
@@ -52,7 +52,7 @@ def gen_annotated(native_yaml_path: str, out: str, autograd_dir: str) -> None:
 @with_native_function
 def gen_annotated_args(f: NativeFunction) -> str:
     out_args: List[Dict[str, Any]] = []
-    for arg in f.func.arguments.positional:
+    for arg in f.func.arguments.flat_positional:
         if arg.default is not None:
             continue
         out_arg: Dict[str, Any] = {}
diff --git a/tools/autograd/gen_trace_type.py b/tools/autograd/gen_trace_type.py
@@ -150,8 +150,7 @@ def dispatch_trace_input(arg: Union[Argument, TensorOptionsArguments]) -> Sequen
         # Factories are a bit special because their out-of-place overloads
         # take an extra TensorOptions argument, which is missing in the _out function
         has_tensor_return = any(r.type.is_tensor_like() for r in f.func.returns)
-        has_tensor_input_arg = any(a.type.is_tensor_like()
-                                   for a in itertools.chain(f.func.arguments.positional, f.func.arguments.kwarg_only))
+        has_tensor_input_arg = any(a.type.is_tensor_like() for a in f.func.arguments.flat_non_out)
         is_factory_method = f.category_override == 'factory' or (has_tensor_return and not has_tensor_input_arg)
 
         # HACK: preserve old codegen behavior - the old codegen set the `is_factory_method`
diff --git a/tools/codegen/api/dispatcher.py b/tools/codegen/api/dispatcher.py
@@ -68,7 +68,11 @@ def name(func: FunctionSchema) -> str:
 
 def arguments(func: FunctionSchema) -> Tuple[DispatcherArgument, ...]:
     if local.use_c10_dispatcher().dispatcher_uses_new_style():
-        return tuple(map(argument, itertools.chain(func.arguments.positional, func.arguments.kwarg_only, func.arguments.out)))
+        return tuple(map(argument, itertools.chain(
+            func.arguments.flat_positional,
+            func.arguments.flat_kwarg_only,
+            func.arguments.out
+        )))
     else:
         return tuple(
             DispatcherArgument(type=la.type, name=la.name, argument=la.argument)
diff --git a/tools/codegen/api/meta.py b/tools/codegen/api/meta.py
@@ -4,7 +4,6 @@
 import tools.codegen.api.dispatcher as dispatcher
 
 from typing import Sequence
-import itertools
 
 # Follows dispatcher calling convention, but:
 #   - Mutable arguments not allowed.  Meta functions are always
@@ -29,4 +28,4 @@ def argument(a: Argument) -> MetaArgument:
 
 def arguments(func: FunctionSchema) -> Sequence[MetaArgument]:
     assert not func.arguments.out
-    return list(map(argument, itertools.chain(func.arguments.positional, func.arguments.kwarg_only)))
+    return list(map(argument, func.arguments.flat_non_out))
diff --git a/tools/codegen/api/python.py b/tools/codegen/api/python.py
@@ -1,4 +1,3 @@
-import itertools
 from dataclasses import dataclass
 from typing import Optional, Union, Sequence, Set, List, Tuple, Dict
 
@@ -734,8 +733,8 @@ def signature(f: NativeFunction, *, method: bool = False, pyi: bool = False) ->
     cpp_args = cpp.group_arguments(f.func, method=method, faithful=True)
     args = tuple(a for a in cpp_args if isinstance(a, Argument))
 
-    input_arg_set = set(a.name for a in f.func.arguments.positional)
-    kwarg_only_set = set(a.name for a in f.func.arguments.kwarg_only)
+    input_arg_set = set(a.name for a in f.func.arguments.flat_positional)
+    kwarg_only_set = set(a.name for a in f.func.arguments.flat_kwarg_only)
     out_arg_set = set(a.name for a in f.func.arguments.out)
 
     input_args = tuple(map(argument, filter(lambda a: a.name in input_arg_set, args)))
@@ -750,8 +749,7 @@ def signature(f: NativeFunction, *, method: bool = False, pyi: bool = False) ->
     # to the original versions in the yaml, this recreation is a potential
     # source of drift between eager and JIT. Pull this logic out to a shared place.
 
-    has_tensor_input_arg = any(a.type.is_tensor_like()
-                               for a in itertools.chain(f.func.arguments.positional, f.func.arguments.kwarg_only))
+    has_tensor_input_arg = any(a.type.is_tensor_like() for a in f.func.arguments.flat_non_out)
     if any(a.name == 'requires_grad' for a in f.func.schema_order_arguments()):
         raise ValueError('argument named requires_grad is reserved, should not explicitly add it in the schema')
 
diff --git a/tools/codegen/gen.py b/tools/codegen/gen.py
@@ -486,11 +486,15 @@ def gen_unstructured(self, f: NativeFunction) -> Optional[str]:
 
             cuda_guard = ""
             if is_generic_dispatch_key(self.dispatch_key) or is_cuda_dispatch_key(self.dispatch_key):
-                self_args = (a for a in f.func.arguments.positional if a.name == "self")
+                self_arg = [f.func.arguments.self_arg.argument] if f.func.arguments.self_arg is not None else []
 
                 # There is precedence for which argument we use to do
                 # device guard.  This describes the precedence order.
-                candidate_args = itertools.chain(self_args, f.func.arguments.out, f.func.arguments.positional)
+                candidate_args = itertools.chain(
+                    self_arg,
+                    f.func.arguments.out,
+                    f.func.arguments.flat_positional
+                )
 
                 # Only tensor like arguments are eligible
                 device_of = next((f'{a.name}' for a in candidate_args if a.type.is_tensor_like()), None)
@@ -619,8 +623,7 @@ def __call__(self, f: NativeFunction) -> Optional[str]:
             return None
 
         assert not f.func.is_out_fn()
-        assert len(f.func.arguments.positional) > 0
-        assert sum(a.name == 'self' for a in f.func.arguments.positional) == 1
+        assert f.func.arguments.self_arg is not None
 
         name = cpp.name(f.func)
 
@@ -992,7 +995,7 @@ def compute_declaration_yaml(f: NativeFunction) -> object:
 
     # These sets are used to conveniently test if an argument is a
     # kwarg-only or out argument
-    kwarg_only_set = set(a.name for a in f.func.arguments.kwarg_only)
+    kwarg_only_set = set(a.name for a in f.func.arguments.flat_kwarg_only)
     out_arg_set = set(a.name for a in f.func.arguments.out)
 
     sig_group = CppSignatureGroup.from_schema(f.func, method=False)
diff --git a/tools/codegen/model.py b/tools/codegen/model.py
@@ -398,7 +398,11 @@ class FunctionSchema:
     returns: Tuple['Return', ...]
 
     def schema_order_arguments(self) -> Iterator['Argument']:
-        return itertools.chain(self.arguments.positional, self.arguments.kwarg_only, self.arguments.out)
+        return itertools.chain(
+            self.arguments.flat_positional,
+            self.arguments.flat_kwarg_only,
+            self.arguments.out
+        )
 
     @staticmethod
     def parse(func: str) -> 'FunctionSchema':
@@ -428,7 +432,7 @@ def __post_init__(self) -> None:
         # This means that all mutable returns should be aliased to a keyword argument
         # (except for "self", which we explicitly don't treat as an out argument because of its use in methods)
         # See Note [is_out_fn]
-        out_and_self = list(self.arguments.out) + [arg for arg in self.arguments.positional if arg.name == "self"]
+        out_and_self = list(self.arguments.out) + [arg for arg in self.arguments.flat_positional if arg.name == "self"]
         mutable_returns = [ret for ret in self.returns if ret.annotation is not None and ret.annotation.is_write]
         for ret in mutable_returns:
             assert any([ret.annotation == arg.annotation for arg in out_and_self]), \
@@ -899,7 +903,14 @@ class Arguments:
     out: Tuple[Argument, ...]  # these are also kwarg-only
 
     @property
-    def positional(self) -> Sequence[Argument]:
+    def flat_non_out(self) -> Sequence[Argument]:
+        ret: List[Argument] = []
+        ret.extend(self.flat_positional)
+        ret.extend(self.flat_kwarg_only)
+        return ret
+
+    @property
+    def flat_positional(self) -> Sequence[Argument]:
         ret: List[Argument] = []
         ret.extend(self.pre_self_positional)
         if self.self_arg is not None:
@@ -909,7 +920,7 @@ def positional(self) -> Sequence[Argument]:
 
     # NB: doesn't contain out arguments
     @property
-    def kwarg_only(self) -> Sequence[Argument]:
+    def flat_kwarg_only(self) -> Sequence[Argument]:
         ret: List[Argument] = []
         ret.extend(self.pre_tensor_options_kwarg_only)
         if self.tensor_options is not None:
@@ -1056,10 +1067,10 @@ def pred(name: str, ty: Type) -> Callable[[Argument], bool]:
 
     def __str__(self) -> str:
         all_arguments: List[str] = []
-        all_arguments.extend(map(str, self.positional))
-        if self.kwarg_only or self.out:
+        all_arguments.extend(map(str, self.flat_positional))
+        if self.flat_kwarg_only or self.out:
             all_arguments.append('*')
-        all_arguments.extend(map(str, self.kwarg_only))
+        all_arguments.extend(map(str, self.flat_kwarg_only))
         all_arguments.extend(map(str, self.out))
         return ', '.join(all_arguments)
 
