| # this code should be common among cwrap and ATen preprocessing | |
| # for now, I have put it in one place but right now is copied out of cwrap | |
| import copy | |
| from typing import Any, Dict, Iterable, List, Union | |
| Arg = Dict[str, Any] | |
| def parse_arguments(args: List[Union[str, Arg]]) -> List[Arg]: | |
| new_args = [] | |
| for arg in args: | |
| # Simple arg declaration of form "<type> <name>" | |
| if isinstance(arg, str): | |
| t, _, name = arg.partition(" ") | |
| new_args.append({"type": t, "name": name}) | |
| elif isinstance(arg, dict): | |
| if "arg" in arg: | |
| arg["type"], _, arg["name"] = arg["arg"].partition(" ") | |
| del arg["arg"] | |
| new_args.append(arg) | |
| else: | |
| raise AssertionError() | |
| return new_args | |
| Declaration = Dict[str, Any] | |
| def set_declaration_defaults(declaration: Declaration) -> None: | |
| if "schema_string" not in declaration: | |
| # This happens for legacy TH bindings like | |
| # _thnn_conv_depthwise2d_backward | |
| declaration["schema_string"] = "" | |
| declaration.setdefault("arguments", []) | |
| declaration.setdefault("return", "void") | |
| if "cname" not in declaration: | |
| declaration["cname"] = declaration["name"] | |
| if "backends" not in declaration: | |
| declaration["backends"] = ["CPU", "CUDA"] | |
| assert "api_name" not in declaration | |
| declaration["api_name"] = declaration["name"] | |
| # NB: keep this in sync with gen_autograd.py | |
| if declaration.get("overload_name"): | |
| declaration["type_wrapper_name"] = "{}_{}".format( | |
| declaration["name"], declaration["overload_name"] | |
| ) | |
| else: | |
| declaration["type_wrapper_name"] = declaration["name"] | |
| # TODO: Uggggh, parsing the schema string here, really??? | |
| declaration["operator_name_with_overload"] = declaration["schema_string"].split( | |
| "(" | |
| )[0] | |
| if declaration["schema_string"]: | |
| declaration["unqual_schema_string"] = declaration["schema_string"].split("::")[ | |
| 1 | |
| ] | |
| declaration["unqual_operator_name_with_overload"] = declaration[ | |
| "operator_name_with_overload" | |
| ].split("::")[1] | |
| else: | |
| declaration["unqual_schema_string"] = "" | |
| declaration["unqual_operator_name_with_overload"] = "" | |
| # Simulate multiple dispatch, even if it's not necessary | |
| if "options" not in declaration: | |
| declaration["options"] = [ | |
| { | |
| "arguments": copy.deepcopy(declaration["arguments"]), | |
| "schema_order_arguments": copy.deepcopy( | |
| declaration["schema_order_arguments"] | |
| ), | |
| } | |
| ] | |
| del declaration["arguments"] | |
| del declaration["schema_order_arguments"] | |
| # Parse arguments (some of them can be strings) | |
| for option in declaration["options"]: | |
| option["arguments"] = parse_arguments(option["arguments"]) | |
| option["schema_order_arguments"] = parse_arguments( | |
| option["schema_order_arguments"] | |
| ) | |
| # Propagate defaults from declaration to options | |
| for option in declaration["options"]: | |
| for k, v in declaration.items(): | |
| # TODO(zach): why does cwrap not propagate 'name'? I need it | |
| # propagaged for ATen | |
| if k != "options": | |
| option.setdefault(k, v) | |
| # TODO(zach): added option to remove keyword handling for C++ which cannot | |
| # support it. | |
| Option = Dict[str, Any] | |
| def filter_unique_options( | |
| options: Iterable[Option], | |
| allow_kwarg: bool, | |
| type_to_signature: Dict[str, str], | |
| remove_self: bool, | |
| ) -> List[Option]: | |
| def exclude_arg(arg: Arg) -> bool: | |
| return arg["type"] == "CONSTANT" # type: ignore[no-any-return] | |
| def exclude_arg_with_self_check(arg: Arg) -> bool: | |
| return exclude_arg(arg) or (remove_self and arg["name"] == "self") | |
| def signature(option: Option, num_kwarg_only: int) -> str: | |
| if num_kwarg_only == 0: | |
| kwarg_only_count = None | |
| else: | |
| kwarg_only_count = -num_kwarg_only | |
| arg_signature = "#".join( | |
| type_to_signature.get(arg["type"], arg["type"]) | |
| for arg in option["arguments"][:kwarg_only_count] | |
| if not exclude_arg_with_self_check(arg) | |
| ) | |
| if kwarg_only_count is None: | |
| return arg_signature | |
| kwarg_only_signature = "#".join( | |
| arg["name"] + "#" + arg["type"] | |
| for arg in option["arguments"][kwarg_only_count:] | |
| if not exclude_arg(arg) | |
| ) | |
| return arg_signature + "#-#" + kwarg_only_signature | |
| seen_signatures = set() | |
| unique = [] | |
| for option in options: | |
| # if only check num_kwarg_only == 0 if allow_kwarg == False | |
| limit = len(option["arguments"]) if allow_kwarg else 0 | |
| for num_kwarg_only in range(0, limit + 1): | |
| sig = signature(option, num_kwarg_only) | |
| if sig not in seen_signatures: | |
| if num_kwarg_only > 0: | |
| for arg in option["arguments"][-num_kwarg_only:]: | |
| arg["kwarg_only"] = True | |
| unique.append(option) | |
| seen_signatures.add(sig) | |
| break | |
| return unique | |
| def sort_by_number_of_args(declaration: Declaration, reverse: bool = True) -> None: | |
| def num_args(option: Option) -> int: | |
| return len(option["arguments"]) | |
| declaration["options"].sort(key=num_args, reverse=reverse) | |
| class Function(object): | |
| def __init__(self, name: str) -> None: | |
| self.name = name | |
| self.arguments: List["Argument"] = [] | |
| def add_argument(self, arg: "Argument") -> None: | |
| assert isinstance(arg, Argument) | |
| self.arguments.append(arg) | |
| def __repr__(self) -> str: | |
| return self.name + "(" + ", ".join(a.__repr__() for a in self.arguments) + ")" | |
| class Argument(object): | |
| def __init__(self, _type: str, name: str, is_optional: bool): | |
| self.type = _type | |
| self.name = name | |
| self.is_optional = is_optional | |
| def __repr__(self) -> str: | |
| return self.type + " " + self.name | |
| def parse_header(path: str) -> List[Function]: | |
| with open(path, "r") as f: | |
| lines: Iterable[Any] = f.read().split("\n") | |
| # Remove empty lines and prebackend directives | |
| lines = filter(lambda l: l and not l.startswith("#"), lines) | |
| # Remove line comments | |
| lines = (l.partition("//") for l in lines) | |
| # Select line and comment part | |
| lines = ((l[0].strip(), l[2].strip()) for l in lines) | |
| # Remove trailing special signs | |
| lines = ((l[0].rstrip(");").rstrip(","), l[1]) for l in lines) | |
| # Split arguments | |
| lines = ((l[0].split(","), l[1]) for l in lines) | |
| # Flatten lines | |
| new_lines = [] | |
| for l, c in lines: | |
| for split in l: | |
| new_lines.append((split, c)) | |
| lines = new_lines | |
| del new_lines | |
| # Remove unnecessary whitespace | |
| lines = ((l[0].strip(), l[1]) for l in lines) | |
| # Remove empty lines | |
| lines = filter(lambda l: l[0], lines) | |
| generic_functions = [] | |
| for l, c in lines: | |
| if l.startswith("TH_API void THNN_"): | |
| fn_name = l[len("TH_API void THNN_") :] | |
| if fn_name[0] == "(" and fn_name[-2] == ")": | |
| fn_name = fn_name[1:-2] | |
| else: | |
| fn_name = fn_name[:-1] | |
| generic_functions.append(Function(fn_name)) | |
| elif l.startswith("TORCH_CUDA_CPP_API void THNN_"): | |
| fn_name = l[len("TORCH_CUDA_CPP_API void THNN_") :] | |
| if fn_name[0] == "(" and fn_name[-2] == ")": | |
| fn_name = fn_name[1:-2] | |
| else: | |
| fn_name = fn_name[:-1] | |
| generic_functions.append(Function(fn_name)) | |
| elif l.startswith("TORCH_CUDA_CU_API void THNN_"): | |
| fn_name = l[len("TORCH_CUDA_CU_API void THNN_") :] | |
| if fn_name[0] == "(" and fn_name[-2] == ")": | |
| fn_name = fn_name[1:-2] | |
| else: | |
| fn_name = fn_name[:-1] | |
| generic_functions.append(Function(fn_name)) | |
| elif l: | |
| t, name = l.split() | |
| if "*" in name: | |
| t = t + "*" | |
| name = name[1:] | |
| generic_functions[-1].add_argument(Argument(t, name, "[OPTIONAL]" in c)) | |
| return generic_functions |