| 75 | #ifdef USE_DNNL |
| 76 | |
| 77 | Tensor CpuBatchNormForwardInference(const BatchNormHandle& bnh, const Tensor& x, |
| 78 | const Tensor& bnScale, const Tensor& bnBias, |
| 79 | Tensor& running_mean, Tensor& running_var) { |
| 80 | CHECK_EQ(x.device()->lang(), kCpp); |
| 81 | Tensor y; |
| 82 | y.ResetLike(x); |
| 83 | |
| 84 | Tensor w = get_bn_weight_from(bnScale, bnBias); |
| 85 | |
| 86 | y.device()->Exec( |
| 87 | [y, w, x, &running_mean, &running_var, &bnh](Context* ctx) mutable { |
| 88 | auto eng = ctx->dnnl_engine; |
| 89 | using namespace dnnl; |
| 90 | |
| 91 | auto x_mem = memory(bnh.x_md, eng, x.block()->mutable_data()); |
| 92 | auto y_mem = memory(bnh.x_md, eng, y.block()->mutable_data()); |
| 93 | // indicates using scale&bias and running mean&var |
| 94 | auto flags_ = normalization_flags::use_scale_shift | |
| 95 | normalization_flags::use_global_stats; |
| 96 | |
| 97 | auto bn_fwd_d = batch_normalization_forward::desc( |
| 98 | prop_kind::forward_inference, bnh.x_md, bnh.epsilon, flags_); |
| 99 | auto bn_fwd_pd = |
| 100 | batch_normalization_forward::primitive_desc(bn_fwd_d, eng); |
| 101 | auto m_mem = memory(bn_fwd_pd.mean_desc(), eng, |
| 102 | running_mean.block()->mutable_data()); |
| 103 | auto v_mem = memory(bn_fwd_pd.variance_desc(), eng, |
| 104 | running_var.block()->mutable_data()); |
| 105 | auto w_mem = |
| 106 | memory(bn_fwd_pd.weights_desc(), eng, w.block()->mutable_data()); |
| 107 | |
| 108 | // execution |
| 109 | batch_normalization_forward(bn_fwd_pd).execute( |
| 110 | ctx->dnnl_stream, {{DNNL_ARG_SRC, x_mem}, |
| 111 | {DNNL_ARG_DST, y_mem}, |
| 112 | {DNNL_ARG_SCALE_SHIFT, w_mem}, |
| 113 | {DNNL_ARG_MEAN, m_mem}, |
| 114 | {DNNL_ARG_VARIANCE, v_mem}}); |
| 115 | ctx->dnnl_stream.wait(); |
| 116 | }, |
| 117 | {x.block(), w.block(), running_mean.block(), running_var.block()}, |
| 118 | {y.block(), running_mean.block(), running_var.block()}, "CpuBatchNormForwardInference"); |
| 119 | |
| 120 | return y; |
| 121 | } |
| 122 | |
| 123 | const std::vector<Tensor> CpuBatchNormForwardTraining( |
| 124 | const BatchNormHandle& bnh, const Tensor& x, const Tensor& bnScale, |