| 183 | } |
| 184 | |
| 185 | void BatchNormForward::get_output_var_shape( |
| 186 | const TensorShapeArray& inp_shape, TensorShapeArray& out_shape) const { |
| 187 | mgb_assert( |
| 188 | inp_shape[0].ndim == 4 && inp_shape[0].ndim == 4 && inp_shape[1].ndim == 4, |
| 189 | "expect input, scale and bias to be 4 dim tensor, but " |
| 190 | "got input dim: %zu, scale dim: %zu, bias dim: %zu", |
| 191 | inp_shape[0].ndim, inp_shape[1].ndim, inp_shape[2].ndim); |
| 192 | |
| 193 | size_t channel_idx; |
| 194 | if (param().param_dim == Param::ParamDim::DIM_111C) { |
| 195 | channel_idx = 3; |
| 196 | } else { |
| 197 | channel_idx = 1; |
| 198 | } |
| 199 | size_t inp_c = inp_shape[0][channel_idx], scale_c = inp_shape[1][channel_idx], |
| 200 | bias_c = inp_shape[2][channel_idx]; |
| 201 | mgb_assert( |
| 202 | inp_c == scale_c && inp_c == bias_c, |
| 203 | "inconsistent channel size, input channel: %zu, scale channel: %zu, bias " |
| 204 | "channel: %zu", |
| 205 | inp_c, scale_c, bias_c); |
| 206 | |
| 207 | out_shape[5] = inp_shape[0]; |
| 208 | for (size_t i = 0; i < 4; ++i) { |
| 209 | out_shape[i] = inp_shape[1]; |
| 210 | } |
| 211 | if (!need_stats()) { |
| 212 | out_shape[0] = out_shape[1] = {0}; |
| 213 | } |
| 214 | if (inp_shape[0].is_empty()) { |
| 215 | out_shape[4] = {0}; |
| 216 | } else { |
| 217 | out_shape[4] = { |
| 218 | megdnn_opr()->get_reserve_in_bytes({inp_shape[0], input(0)->dtype()})}; |
| 219 | } |
| 220 | } |
| 221 | |
| 222 | size_t BatchNormForward::get_workspace_size_bytes( |
| 223 | const TensorShapeArray& input_shapes, |
nothing calls this directly
no test coverage detected