| 116 | |
| 117 | |
| 118 | class Conv2DProblemSize: |
| 119 | def __init__( |
| 120 | self, n: int, h: int, w: int, c: int, |
| 121 | k: int, r: int, s: int, c_: int, |
| 122 | pad_h: int, pad_w: int, stride_h: int, stride_w: int, |
| 123 | dilation_h: int, dilation_w: int, mode: ConvMode=ConvMode.CrossCorrelation, |
| 124 | split_k_slices: int=1, groups: int=1): |
| 125 | |
| 126 | self.N = n |
| 127 | self.H = h |
| 128 | self.W = w |
| 129 | self.C = c |
| 130 | self.K = k |
| 131 | self.R = r |
| 132 | self.S = s |
| 133 | self.pad_h = pad_h |
| 134 | self.pad_w = pad_w |
| 135 | self.stride_h = stride_h |
| 136 | self.stride_w = stride_w |
| 137 | self.dilation_h = dilation_h |
| 138 | self.dilation_w = dilation_w |
| 139 | self.mode = int(mode) |
| 140 | self.split_k_slices = split_k_slices |
| 141 | self.groups = groups |
| 142 | self.P = ((h + pad_h * 2 - r * dilation_h) // stride_h) + 1 |
| 143 | self.Q = ((w + pad_w * 2 - s * dilation_w) // stride_w) + 1 |
| 144 | |
| 145 | @property |
| 146 | def ctype(self) -> Conv2DProblemSize_: |
| 147 | return Conv2DProblemSize_(self) |
| 148 | |
| 149 | def implicit_gemm_size(self, kind: ConvKind): |
| 150 | if kind == ConvKind.Fprop: |
| 151 | return GemmCoord( |
| 152 | self.N * self.P * self.Q, |
| 153 | self.K, |
| 154 | self.R * self.S * self.C // self.groups |
| 155 | ) |
| 156 | elif kind == ConvKind.Dgrad: |
| 157 | return GemmCoord( |
| 158 | self.N * self.H * self.W, |
| 159 | self.C, |
| 160 | self.R * self.S * self.K |
| 161 | ) |
| 162 | elif kind == ConvKind.Wgrad: |
| 163 | return GemmCoord( |
| 164 | self.K, |
| 165 | self.R * self.S * self.C, |
| 166 | self.N * self.P * self.Q |
| 167 | ) |
| 168 | |
| 169 | @staticmethod |
| 170 | def from_sizes(input_size, weight_size): |
| 171 | K, R, S, _ = weight_size |
| 172 | pad_h = R // 2 |
| 173 | pad_w = S // 2 |
| 174 | stride_h = 1 |
| 175 | stride_w = 1 |
no outgoing calls