| 251 | |
| 252 | template<typename Ti, typename Tw, typename To> |
| 253 | void meanFirst(Param out, Param in, Param inWeight) { |
| 254 | uint threads_x = nextpow2(std::max(32u, (uint)in.info.dims[0])); |
| 255 | threads_x = std::min(threads_x, THREADS_PER_GROUP); |
| 256 | uint threads_y = THREADS_PER_GROUP / threads_x; |
| 257 | |
| 258 | uint groups_x = divup(in.info.dims[0], threads_x * REPEAT); |
| 259 | uint groups_y = divup(in.info.dims[1], threads_y); |
| 260 | |
| 261 | Param tmpOut = out; |
| 262 | Param noWeight; |
| 263 | noWeight.info.offset = 0; |
| 264 | for (int k = 0; k < 4; ++k) { |
| 265 | noWeight.info.dims[k] = 0; |
| 266 | noWeight.info.strides[k] = 0; |
| 267 | } |
| 268 | // Does not matter what the value is it will not be used. Just needs to be |
| 269 | // valid. |
| 270 | noWeight.data = inWeight.data; |
| 271 | |
| 272 | Param tmpWeight = noWeight; |
| 273 | |
| 274 | if (groups_x > 1) { |
| 275 | tmpOut.data = bufferAlloc(groups_x * in.info.dims[1] * in.info.dims[2] * |
| 276 | in.info.dims[3] * sizeof(To)); |
| 277 | |
| 278 | tmpWeight.data = |
| 279 | bufferAlloc(groups_x * in.info.dims[1] * in.info.dims[2] * |
| 280 | in.info.dims[3] * sizeof(Tw)); |
| 281 | |
| 282 | tmpOut.info.dims[0] = groups_x; |
| 283 | for (int k = 1; k < 4; k++) tmpOut.info.strides[k] *= groups_x; |
| 284 | tmpWeight.info = tmpOut.info; |
| 285 | } |
| 286 | |
| 287 | meanFirstLauncher<Ti, Tw, To>(tmpOut, tmpWeight, in, inWeight, threads_x, |
| 288 | groups_x, groups_y); |
| 289 | |
| 290 | if (groups_x > 1) { |
| 291 | // No Weight is needed when writing out the output. |
| 292 | meanFirstLauncher<Ti, Tw, To>(out, noWeight, tmpOut, tmpWeight, |
| 293 | threads_x, 1, groups_y); |
| 294 | |
| 295 | bufferFree(tmpOut.data); |
| 296 | bufferFree(tmpWeight.data); |
| 297 | } |
| 298 | } |
| 299 | |
| 300 | template<typename Ti, typename Tw, typename To> |
| 301 | void meanWeighted(Param out, Param in, Param inWeight, int dim) { |
nothing calls this directly
no test coverage detected