mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
164bee1d09
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/55046 Updating `returns` in the codegen to return a CType instead of a raw string. This has benefit of putting all stringifying logic through CType, which is useful in the followup PR when I add namespaces. I also added new CTypes for other templated C++ types: array, vector and tuple. Mostly because it makes the namespacing logic in the next PR significantly easier. It also seems more natural to me that `BaseCType` shouldn't represent specializations of templated types. There's a little bit of weirdness, types that are currently *only* used for returns, i.e. `TupleCType`. Returns aren't named, so I opted not to give it one- so we can add it in later if we discover that we need it. Test Plan: Imported from OSS Reviewed By: bhosmer Differential Revision: D27708348 Pulled By: bdhirsh fbshipit-source-id: 230b210c3e53be1bd362105fbea8451055dc59a8
310 lines
12 KiB
Python
310 lines
12 KiB
Python
from tools.codegen.model import (Argument, Arguments, BaseTy, BaseType,
|
|
FunctionSchema, ListType, NativeFunction,
|
|
OptionalType, Return, SelfArgument,
|
|
TensorOptionsArguments, Type, assert_never)
|
|
from tools.codegen.api.types import (ArgName, BaseCType, Binding,
|
|
ConstRefCType, CType, MutRefCType,
|
|
OptionalCType, SpecialArgName,
|
|
TupleCType, ArrayCType, ListCType, VectorCType, ArrayRefCType)
|
|
from typing import Optional, Sequence, Union, List, Set
|
|
|
|
# This file describes the translation of JIT schema to the public C++
|
|
# API, which is what people use when they call functions like at::add.
|
|
#
|
|
# Prominent characteristics of the C++ API:
|
|
#
|
|
# - dtype, layout, device and pin_memory are collected into
|
|
# a single C++ type TensorOptions (the native functions API
|
|
# also has this, but tensor options is really most relevant
|
|
# for the C++ API; it makes calling kwarg factory functions
|
|
# pleasant)
|
|
#
|
|
# - defaulting lives here (in fact, the dispatcher is completely
|
|
# oblivious of defaults!)
|
|
#
|
|
# BTW: policy on name collisions: we try not to have types with
|
|
# collisions, but functions are fair game to collide
|
|
|
|
def name(func: FunctionSchema, *, faithful_name_for_out_overloads: bool = False) -> str:
|
|
name = str(func.name.name)
|
|
if func.is_out_fn():
|
|
if faithful_name_for_out_overloads:
|
|
name += '_outf'
|
|
else:
|
|
name += '_out'
|
|
|
|
return name
|
|
|
|
# Translation of "value types" in JIT schema to C++ API type. Value
|
|
# types look the same no matter if they are argument types or return
|
|
# types. Returns None if the type in question is not a value type.
|
|
def valuetype_type(t: Type, *, binds: ArgName) -> Optional[CType]:
|
|
if isinstance(t, BaseType):
|
|
if t.name == BaseTy.Tensor or t.name == BaseTy.Scalar:
|
|
return None
|
|
elif t.name == BaseTy.int:
|
|
return BaseCType('int64_t', binds)
|
|
elif t.name == BaseTy.float:
|
|
return BaseCType('double', binds)
|
|
elif t.name == BaseTy.str:
|
|
return BaseCType('std::string', binds)
|
|
elif t.name in [BaseTy.bool, BaseTy.QScheme, BaseTy.Scalar,
|
|
BaseTy.ScalarType, BaseTy.Generator, BaseTy.Storage,
|
|
BaseTy.Layout, BaseTy.Device, BaseTy.MemoryFormat,
|
|
BaseTy.Dimname, BaseTy.Stream, BaseTy.ConstQuantizerPtr]:
|
|
# These C++ names line up with their schema names
|
|
return BaseCType(t.name.name, binds)
|
|
else:
|
|
raise AssertionError(f"unsupported type: {t}")
|
|
elif isinstance(t, OptionalType):
|
|
elem = valuetype_type(t.elem, binds=binds)
|
|
if elem is None:
|
|
return None
|
|
return OptionalCType(elem)
|
|
elif isinstance(t, ListType):
|
|
if str(t.elem) == 'bool':
|
|
assert t.size is not None
|
|
return ArrayCType(BaseCType("bool", binds), t.size)
|
|
else:
|
|
return None
|
|
else:
|
|
raise AssertionError(f"unrecognized type {repr(t)}")
|
|
|
|
# Translation of types occuring in JIT arguments to a C++ argument type.
|
|
def argumenttype_type(t: Type, *, mutable: bool, binds: ArgName) -> CType:
|
|
# If it's a value type, do the value type translation
|
|
r = valuetype_type(t, binds=binds)
|
|
if r is not None:
|
|
return r
|
|
|
|
if isinstance(t, BaseType):
|
|
if t.name == BaseTy.Tensor:
|
|
if mutable:
|
|
return MutRefCType(BaseCType('Tensor', binds))
|
|
else:
|
|
return ConstRefCType(BaseCType('Tensor', binds))
|
|
elif t.name == BaseTy.Scalar:
|
|
return ConstRefCType(BaseCType('Scalar', binds))
|
|
else:
|
|
raise AssertionError(f"base type should have been value type {t}")
|
|
elif isinstance(t, OptionalType):
|
|
if str(t.elem) == 'Tensor':
|
|
if mutable:
|
|
return MutRefCType(BaseCType('Tensor', binds)) # TODO: fix this discrepancy
|
|
else:
|
|
return ConstRefCType(OptionalCType(BaseCType('Tensor', binds)))
|
|
elif str(t.elem) == 'Scalar':
|
|
return ConstRefCType(OptionalCType(BaseCType('Scalar', binds)))
|
|
elem = argumenttype_type(t.elem, mutable=mutable, binds=binds)
|
|
return OptionalCType(elem)
|
|
elif isinstance(t, ListType):
|
|
# TODO: remove these special cases, ArrayRef fallthrough works fine
|
|
if str(t.elem) == 'int':
|
|
return BaseCType("IntArrayRef", binds)
|
|
elif str(t.elem) == 'Tensor':
|
|
return BaseCType("TensorList", binds)
|
|
elif str(t.elem) == 'Scalar':
|
|
return ArrayRefCType(BaseCType("Scalar", binds))
|
|
elif str(t.elem) == 'Dimname':
|
|
return BaseCType("DimnameList", binds)
|
|
elif str(t.elem) == 'Tensor?':
|
|
return ConstRefCType(ListCType(OptionalCType(BaseCType("Tensor", binds))))
|
|
elem = argumenttype_type(t.elem, mutable=mutable, binds=binds)
|
|
return ArrayRefCType(elem)
|
|
else:
|
|
raise AssertionError(f"unrecognized type {repr(t)}")
|
|
|
|
# Translate a JIT argument into its C++ type
|
|
def argument_type(a: Argument, *, binds: ArgName) -> CType:
|
|
return argumenttype_type(a.type, mutable=a.is_write, binds=binds)
|
|
|
|
# Translation of a (non-multi) return type from JIT to C++
|
|
# NB: if need translations on return types, make this return CType too. Need to
|
|
# take care; ArgName is misnomer now, and inputs are permitted to conflict with outputs
|
|
# so need to make sure you don't have trouble
|
|
def returntype_type(t: Type, *, mutable: bool) -> CType:
|
|
# placeholder is ignored
|
|
r = valuetype_type(t, binds="__placeholder__")
|
|
if r is not None:
|
|
return r
|
|
|
|
if isinstance(t, BaseType):
|
|
if t.name == BaseTy.Tensor:
|
|
if mutable:
|
|
return MutRefCType(BaseCType('Tensor', "__placeholder__"))
|
|
else:
|
|
return BaseCType('Tensor', "__placeholder__")
|
|
elif t.name == BaseTy.Scalar:
|
|
return BaseCType('Scalar', "__placeholder__")
|
|
elif isinstance(t, ListType):
|
|
elem = returntype_type(t.elem, mutable=mutable)
|
|
assert t.size is None, f"fixed size list returns not supported: {t}"
|
|
return VectorCType(elem)
|
|
|
|
raise AssertionError(f"unrecognized return type {t}")
|
|
|
|
# Translation of a single return to its C++ type
|
|
def return_type(r: Return) -> CType:
|
|
return returntype_type(r.type, mutable=r.is_write)
|
|
|
|
# Translation of a full (possibly multi) return from JIT to its C++ type
|
|
def returns_type(rs: Sequence[Return]) -> CType:
|
|
if len(rs) == 0:
|
|
return BaseCType('void', "__placeholder__")
|
|
elif len(rs) == 1:
|
|
return return_type(rs[0])
|
|
else:
|
|
return TupleCType([return_type(r) for r in rs])
|
|
|
|
def return_names(f: NativeFunction) -> Sequence[str]:
|
|
returns: List[str] = []
|
|
for i, r in enumerate(f.func.returns):
|
|
# If we have an inplace function, the return argument is
|
|
# implicitly named self.
|
|
# TODO: Consider incorporating this into the data model
|
|
if f.func.name.name.inplace:
|
|
assert i == 0, "illegal inplace function with multiple returns"
|
|
name = 'self'
|
|
# If we are out function, the name is the name of the
|
|
# corresponding output function (r.name will get recorded
|
|
# in field_name later.)
|
|
elif f.func.is_out_fn():
|
|
name = f.func.arguments.out[i].name
|
|
# If the return argument is explicitly named...
|
|
elif r.name:
|
|
name_conflict = any(r.name == a.name for a in f.func.schema_order_arguments())
|
|
if name_conflict and not f.func.is_out_fn():
|
|
name = f'{r.name}_return'
|
|
else:
|
|
name = r.name
|
|
# If there is no explicit name, we just name the output result,
|
|
# unless it's a multi-return, in which case it's result0,
|
|
# result1, etc (zero-indexed)
|
|
else:
|
|
name = 'result' if len(f.func.returns) == 1 else f'result{i}'
|
|
returns.append(name)
|
|
return returns
|
|
|
|
JIT_TO_CPP_DEFAULT = {
|
|
'False': 'false',
|
|
'True': 'true',
|
|
'None': 'c10::nullopt', # UGH this one is type directed
|
|
'Mean': 'at::Reduction::Mean',
|
|
'[]': '{}',
|
|
'contiguous_format': 'MemoryFormat::Contiguous',
|
|
'long': 'at::kLong',
|
|
}
|
|
|
|
# Convert a JIT default into C++ expression representing the default
|
|
def default_expr(d: str, t: Type) -> str:
|
|
if d == 'None' and str(t) == 'Tensor?':
|
|
return '{}'
|
|
if isinstance(t, BaseType) and t.name is BaseTy.str:
|
|
# Schema allows single quotes but C++ needs double
|
|
if len(d) >= 2 and d[0] == "'" and d[-1] == "'":
|
|
s = ''
|
|
i = 1
|
|
while i + 1 < len(d):
|
|
if d[i] != '\\':
|
|
if d[i] == '"':
|
|
s += '\\"'
|
|
else:
|
|
s += d[i]
|
|
i += 1
|
|
else:
|
|
if d[i + 1] == "'":
|
|
s += "'"
|
|
else:
|
|
s += d[i:i + 2]
|
|
i += 2
|
|
|
|
return f'"{s}"'
|
|
|
|
if isinstance(t, OptionalType):
|
|
if d == 'None':
|
|
return 'c10::nullopt'
|
|
|
|
return default_expr(d, t.elem)
|
|
|
|
if isinstance(t, ListType):
|
|
if (d.startswith('[') and d.endswith(']')):
|
|
return '{' + d[1:-1] + '}'
|
|
elif t.size is None:
|
|
# NOTE: Sized lists can have scalar defaults
|
|
raise ValueError(f"Expected a list default '[...]' but found: '{d}'")
|
|
|
|
return JIT_TO_CPP_DEFAULT.get(d, d)
|
|
|
|
# Convert an argument into its C++ API form
|
|
|
|
def argument(
|
|
a: Union[Argument, TensorOptionsArguments, SelfArgument],
|
|
*, cpp_no_default_args: Set[str], method: bool, faithful: bool,
|
|
has_tensor_options: bool
|
|
) -> List[Binding]:
|
|
def sub_argument(a: Union[Argument, TensorOptionsArguments, SelfArgument]) -> List[Binding]:
|
|
return argument(
|
|
a, cpp_no_default_args=cpp_no_default_args, method=method, faithful=faithful,
|
|
has_tensor_options=has_tensor_options)
|
|
|
|
if isinstance(a, Argument):
|
|
binds: ArgName
|
|
if a.name == "memory_format" and has_tensor_options:
|
|
binds = SpecialArgName.possibly_redundant_memory_format
|
|
else:
|
|
binds = a.name
|
|
default: Optional[str] = None
|
|
if a.name not in cpp_no_default_args and a.default is not None:
|
|
default = default_expr(a.default, a.type)
|
|
return [Binding(
|
|
ctype=argument_type(a, binds=binds),
|
|
name=a.name,
|
|
default=default,
|
|
argument=a,
|
|
)]
|
|
elif isinstance(a, TensorOptionsArguments):
|
|
if faithful:
|
|
return sub_argument(a.dtype) + sub_argument(a.layout) + \
|
|
sub_argument(a.device) + sub_argument(a.pin_memory)
|
|
else:
|
|
default = None
|
|
# Enforced by NativeFunction.__post_init__
|
|
assert 'options' not in cpp_no_default_args
|
|
if all(x.default == "None" for x in a.all()):
|
|
default = '{}'
|
|
elif a.dtype.default == "long":
|
|
default = 'at::kLong' # TODO: this is wrong
|
|
return [Binding(
|
|
ctype=BaseCType('TensorOptions', 'options'),
|
|
name='options',
|
|
default=default,
|
|
argument=a,
|
|
)]
|
|
elif isinstance(a, SelfArgument):
|
|
if method:
|
|
# Caller is responsible for installing implicit this in context!
|
|
return []
|
|
else:
|
|
return sub_argument(a.argument)
|
|
else:
|
|
assert_never(a)
|
|
|
|
def arguments(
|
|
arguments: Arguments,
|
|
*, faithful: bool, method: bool, cpp_no_default_args: Set[str]
|
|
) -> List[Binding]:
|
|
args: List[Union[Argument, TensorOptionsArguments, SelfArgument]] = []
|
|
if faithful:
|
|
args.extend(arguments.non_out)
|
|
args.extend(arguments.out)
|
|
else:
|
|
args.extend(arguments.out)
|
|
args.extend(arguments.non_out)
|
|
return [
|
|
r.no_default() if faithful else r for a in args
|
|
for r in argument(
|
|
a, faithful=faithful, method=method,
|
|
has_tensor_options=arguments.tensor_options is not None,
|
|
cpp_no_default_args=cpp_no_default_args)
|
|
]
|