mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +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
67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
from tools.codegen.model import (Argument, FunctionSchema, Return,
|
|
SelfArgument, TensorOptionsArguments, Type,
|
|
assert_never)
|
|
|
|
from tools.codegen.api.types import ArgName, Binding, CType
|
|
from tools.codegen.api import cpp
|
|
|
|
import itertools
|
|
from typing import Sequence, List, Union
|
|
|
|
# This file describes the translation of JIT schema to the dispatcher
|
|
# API, the *unboxed* calling convention by which invocations through
|
|
# the dispatcher are made. Historically, the dispatcher API matched
|
|
# the C++ API, but with the establishment of the boxed API, we've
|
|
# made changes to the dispatcher API to so that the unboxed API
|
|
# better aligns with the boxed API. The dispatcher API hooks heavily
|
|
# into our template based boxing/unboxing machinery, so changes
|
|
# to this convention will usually need template updates too.
|
|
#
|
|
# Prominent characteristics of the dispatcher API:
|
|
#
|
|
# - dtype, layout, device and pin_memory are represented as separate
|
|
# arguments.
|
|
#
|
|
|
|
def name(func: FunctionSchema) -> str:
|
|
return cpp.name(func)
|
|
|
|
def argumenttype_type(t: Type, *, mutable: bool, binds: ArgName) -> CType:
|
|
# This is a faux amis. If it makes sense in the future to add
|
|
# more special cases here, or invert things so cpp.argument_type
|
|
# calls this, or just completely inline the function, please do
|
|
# it.
|
|
return cpp.argumenttype_type(t, mutable=mutable, binds=binds)
|
|
|
|
def argument_type(a: Argument, *, binds: ArgName) -> CType:
|
|
return argumenttype_type(a.type, mutable=a.is_write, binds=binds)
|
|
|
|
def returns_type(rs: Sequence[Return]) -> CType:
|
|
# At present, there is no difference. But there could be!
|
|
return cpp.returns_type(rs)
|
|
|
|
def argument(
|
|
a: Union[Argument, TensorOptionsArguments, SelfArgument]
|
|
) -> List[Binding]:
|
|
if isinstance(a, Argument):
|
|
return [Binding(
|
|
ctype=argument_type(a, binds=a.name),
|
|
name=a.name,
|
|
argument=a,
|
|
)]
|
|
elif isinstance(a, SelfArgument):
|
|
return argument(a.argument)
|
|
elif isinstance(a, TensorOptionsArguments):
|
|
return argument(a.dtype) + argument(a.layout) + argument(a.device) + argument(a.pin_memory)
|
|
else:
|
|
assert_never(a)
|
|
|
|
def arguments(func: FunctionSchema) -> List[Binding]:
|
|
return [
|
|
r for a in itertools.chain(
|
|
func.arguments.positional,
|
|
func.arguments.kwarg_only,
|
|
func.arguments.out
|
|
) for r in argument(a)
|
|
]
|