| 274 | |
| 275 | template <typename Pooler, typename IdxGetter, typename ctype = typename Pooler::ctype> |
| 276 | void pooling_forward_impl( |
| 277 | const ctype* __restrict src, ctype* __restrict dst, DType src_dtype, size_t N, |
| 278 | size_t C, size_t IH, size_t IW, size_t OH, size_t OW, size_t PH, size_t PW, |
| 279 | size_t SH, size_t SW, size_t FH, size_t FW) { |
| 280 | rep(n, N) rep(c, C) rep(oh, OH) rep(ow, OW) { |
| 281 | Pooler pooler(FH * FW, src_dtype); |
| 282 | pooler.init(); |
| 283 | rep(fh, FH) rep(fw, FW) { |
| 284 | size_t ih = -PH + oh * SH + fh; |
| 285 | size_t iw = -PW + ow * SW + fw; |
| 286 | if (ih < IH && iw < IW) { |
| 287 | size_t idx = IdxGetter::get_idx(n, c, ih, iw, N, C, IH, IW); |
| 288 | pooler.feed(src[idx]); |
| 289 | } |
| 290 | } |
| 291 | size_t idx = IdxGetter::get_idx(n, c, oh, ow, N, C, OH, OW); |
| 292 | dst[idx] = pooler.get_ans(); |
| 293 | } |
| 294 | } |
| 295 | |
| 296 | template <typename ctype, typename IdxGetter> |
| 297 | void pooling_backward_avg_impl( |