| 267 | |
| 268 | template<typename Ti, typename Tw, typename To, int dim> |
| 269 | void mean_dim(Param<To> out, Param<Ti> in, Param<Tw> iwt) { |
| 270 | uint threads_y = std::min(THREADS_Y, nextpow2(in.info.dims[dim])); |
| 271 | uint threads_x = THREADS_X; |
| 272 | |
| 273 | dim_t blocks_dim[] = {divup(in.info.dims[0], threads_x), in.info.dims[1], |
| 274 | in.info.dims[2], in.info.dims[3]}; |
| 275 | |
| 276 | blocks_dim[dim] = divup(in.info.dims[dim], threads_y * REPEAT); |
| 277 | |
| 278 | Array<To> tmpOut = createEmptyArray<To>(dim4()); |
| 279 | Array<Tw> tmpWt = createEmptyArray<Tw>(dim4()); |
| 280 | |
| 281 | if (blocks_dim[dim] > 1) { |
| 282 | dim4 dims(4, out.info.dims); |
| 283 | dims[dim] = blocks_dim[dim]; |
| 284 | tmpOut = createEmptyArray<To>(dims); |
| 285 | tmpWt = createEmptyArray<Tw>(dims); |
| 286 | } else { |
| 287 | tmpOut = createParamArray(out, false); |
| 288 | } |
| 289 | |
| 290 | mean_dim_launcher<Ti, Tw, To, dim>(tmpOut, tmpWt, in, iwt, threads_y, |
| 291 | blocks_dim); |
| 292 | |
| 293 | if (blocks_dim[dim] > 1) { |
| 294 | blocks_dim[dim] = 1; |
| 295 | |
| 296 | Array<Tw> owt = createEmptyArray<Tw>(dim4()); |
| 297 | mean_dim_launcher<To, Tw, To, dim>(out, owt, tmpOut, tmpWt, threads_y, |
| 298 | blocks_dim); |
| 299 | } |
| 300 | } |
| 301 | |
| 302 | // Calculate mean along the first dimension. If wt is an empty Param, use |
| 303 | // weight as 1 and treat it as count. If owt is empty Param, do not write |
nothing calls this directly
no test coverage detected