| 93 | #ifdef USE_DNNL |
| 94 | |
| 95 | Tensor CpuPoolingForward(const PoolingHandle &ph, const Tensor &x) { |
| 96 | CHECK_EQ(x.device()->lang(), kCpp); |
| 97 | Tensor y({(unsigned long)ph.batchsize, (unsigned long)ph.channels, |
| 98 | (unsigned long)ph.pooled_height, (unsigned long)ph.pooled_width}, |
| 99 | x.device(), x.data_type()); |
| 100 | |
| 101 | y.device()->Exec( |
| 102 | [y, x, &ph](Context *ctx) mutable { |
| 103 | auto eng = ctx->dnnl_engine; |
| 104 | using namespace dnnl; |
| 105 | |
| 106 | memory x_mem(ph.x_md, eng, x.block()->mutable_data()); |
| 107 | memory y_mem(ph.y_md, eng, y.block()->mutable_data()); |
| 108 | |
| 109 | pooling_forward(ph.pool_fwd_pd) |
| 110 | .execute(ctx->dnnl_stream, {{DNNL_ARG_SRC, x_mem}, |
| 111 | {DNNL_ARG_DST, y_mem}, |
| 112 | {DNNL_ARG_WORKSPACE, ph.ws_mem}}); |
| 113 | ctx->dnnl_stream.wait(); |
| 114 | }, |
| 115 | {x.block()}, {y.block()}, "CpuPoolingForward"); |
| 116 | |
| 117 | return y; |
| 118 | } |
| 119 | |
| 120 | Tensor CpuPoolingBackward(const PoolingHandle &ph, const Tensor &grad, |
| 121 | const Tensor &x, const Tensor &y) { |