pytorch

1
import itertools
2
from typing import List, Sequence, Union
3

4
from torchgen.api import cpp
5

6
from torchgen.api.types import ArgName, Binding, CType, NamedCType
7
from torchgen.model import (
8
    Argument,
9
    FunctionSchema,
10
    Return,
11
    SelfArgument,
12
    TensorOptionsArguments,
13
    Type,
14
)
15
from torchgen.utils import assert_never, concatMap
16

17
# This file describes the translation of JIT schema to the dispatcher
18
# API, the *unboxed* calling convention by which invocations through
19
# the dispatcher are made.  Historically, the dispatcher API matched
20
# the C++ API, but with the establishment of the boxed API, we've
21
# made changes to the dispatcher API to so that the unboxed API
22
# better aligns with the boxed API.  The dispatcher API hooks heavily
23
# into our template based boxing/unboxing machinery, so changes
24
# to this convention will usually need template updates too.
25
#
26
# Prominent characteristics of the dispatcher API:
27
#
28
#   - dtype, layout, device and pin_memory are represented as separate
29
#     arguments.
30
#
31

32

33
def name(func: FunctionSchema) -> str:
34
    return cpp.name(func)
35

36

37
def argumenttype_type(
38
    t: Type,
39
    *,
40
    mutable: bool,
41
    binds: ArgName,
42
    remove_non_owning_ref_types: bool = False,
43
    symint: bool = True,
44
) -> NamedCType:
45
    # This is a faux amis.  If it makes sense in the future to add
46
    # more special cases here, or invert things so cpp.argument_type
47
    # calls this, or just completely inline the function, please do
48
    # it.
49
    return cpp.argumenttype_type(
50
        t,
51
        mutable=mutable,
52
        binds=binds,
53
        symint=symint,
54
        remove_non_owning_ref_types=remove_non_owning_ref_types,
55
    )
56

57

58
def argument_type(
59
    a: Argument,
60
    *,
61
    binds: ArgName,
62
    remove_non_owning_ref_types: bool = False,
63
    symint: bool = True,
64
) -> NamedCType:
65
    return argumenttype_type(
66
        a.type,
67
        mutable=a.is_write,
68
        binds=binds,
69
        remove_non_owning_ref_types=remove_non_owning_ref_types,
70
        symint=symint,
71
    )
72

73

74
def returns_type(rs: Sequence[Return], *, symint: bool = True) -> CType:
75
    # At present, there is no difference. But there could be!
76
    return cpp.returns_type(rs, symint=symint)
77

78

79
def jit_arguments(func: FunctionSchema) -> List[Argument]:
80
    def to_argument(
81
        a: Union[Argument, TensorOptionsArguments, SelfArgument]
82
    ) -> List[Argument]:
83
        if isinstance(a, Argument):
84
            return [a]
85
        elif isinstance(a, SelfArgument):
86
            return [a.argument]
87
        elif isinstance(a, TensorOptionsArguments):
88
            return [a.dtype, a.layout, a.device, a.pin_memory]
89
        else:
90
            assert_never(a)
91

92
    return list(
93
        concatMap(
94
            to_argument,
95
            itertools.chain(
96
                func.arguments.positional, func.arguments.kwarg_only, func.arguments.out
97
            ),
98
        )
99
    )
100

101

102
def argument(
103
    a: Argument, *, remove_non_owning_ref_types: bool = False, symint: bool = True
104
) -> Binding:
105
    return Binding(
106
        nctype=argument_type(
107
            a,
108
            binds=a.name,
109
            remove_non_owning_ref_types=remove_non_owning_ref_types,
110
            symint=symint,
111
        ),
112
        name=a.name,
113
        argument=a,
114
    )
115

116

117
def arguments(func: FunctionSchema, *, symint: bool = True) -> List[Binding]:
118
    return [argument(a, symint=symint) for a in jit_arguments(func)]
119