RepDL/repdl/ops.py at main · microsoft/RepDL · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

from typing import Union

import torch

from .backend import cpu
from .backend import cuda


def conv2d(
    input: torch.Tensor,
    weight: torch.Tensor,
    stride: tuple[int, int],
    padding: tuple[int, int],
    groups: int,
) -> torch.Tensor:
    if input.is_cuda:
        return cuda.conv2d(
            input, weight, stride[0], stride[1], padding[0], padding[1], groups
        )
    else:
        return cpu.conv2d(
            input, weight, stride[0], stride[1], padding[0], padding[1], groups
        )


def conv2d_grad_input(
    input_shape: tuple[int, ...],
    weight: torch.Tensor,
    grad_output: torch.Tensor,
    stride: tuple[int, int],
    padding: tuple[int, int],
) -> torch.Tensor:
    if grad_output.is_cuda:
        return cuda.conv2d_grad_input(
            grad_output,
            weight,
            input_shape[2],
            input_shape[3],
            stride[0],
            stride[1],
            padding[0],
            padding[1],
        )
    else:
        return cpu.conv2d_grad_input(
            grad_output,
            weight,
            input_shape[2],
            input_shape[3],
            stride[0],
            stride[1],
            padding[0],
            padding[1],
        )


def conv2d_grad_kernel(
    input: torch.Tensor,
    weight_shape: tuple[int, ...],
    grad_output: torch.Tensor,
    stride: tuple[int, int],
    padding: tuple[int, int],
) -> torch.Tensor:
    if input.is_cuda:
        return cuda.conv2d_grad_kernel(
            grad_output,
            input,
            weight_shape[2],
            weight_shape[3],
            stride[0],
            stride[1],
            padding[0],
            padding[1],
        )
    else:
        return cpu.conv2d_grad_kernel(
            grad_output,
            input,
            weight_shape[2],
            weight_shape[3],
            stride[0],
            stride[1],
            padding[0],
            padding[1],
        )


def cross_entropy(input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    if input.is_cuda:
        return cuda.cross_entropy(input, target)
    else:
        return cpu.cross_entropy(input, target)


def div(input: torch.Tensor, other: Union[torch.Tensor, float]) -> torch.Tensor:
    # See pytorch/aten/src/ATen/native/cuda/BinaryMulDivKernel.cu
    if isinstance(other, torch.Tensor) or input.device == "cpu":
        return input.div(other)
    else:
        return input.div(torch.tensor(other, device=input.device))


def mm(
    input: torch.Tensor, mat2: torch.Tensor, transA: bool = False, transB: bool = False
) -> torch.Tensor:
    if input.is_cuda:
        return cuda.mm(input, mat2, transA, transB)
    else:
        return cpu.mm(input, mat2, transA, transB)


def softmax(input: torch.Tensor) -> torch.Tensor:
    if input.is_cuda:
        return cuda.softmax(input)
    else:
        return cpu.softmax(input)


def sqrt(input: torch.Tensor) -> torch.Tensor:
    return input.to(dtype=torch.float64).sqrt().to(dtype=torch.float32)


def sum1d(input: torch.Tensor, average: bool = False) -> torch.Tensor:
    if input.is_cuda:
        return cuda.sum1d(input, average)
    else:
        return cpu.sum1d(input, average)


def sum2d_dim0(
    input: torch.Tensor, keepdim: bool = False, average: bool = False
) -> torch.Tensor:
    if input.is_cuda:
        return cuda.sum2d_dim0(input, keepdim, average)
    else:
        return cpu.sum2d_dim0(input, keepdim, average)


def sum2d_dim1(
    input: torch.Tensor, keepdim: bool = False, average: bool = False
) -> torch.Tensor:
    if input.is_cuda:
        return cuda.sum2d_dim1(input, keepdim, average)
    else:
        return cpu.sum2d_dim1(input, keepdim, average)


def sum4d_dim023(
    input: torch.Tensor, keepdim: bool = False, average: bool = False
) -> torch.Tensor:
    if input.is_cuda:
        return cuda.sum4d_dim023(input, keepdim, average)
    else:
        return cpu.sum4d_dim023(input, keepdim, average)