| 94 | } |
| 95 | |
| 96 | Tensor CpuConvForward(const Tensor &x, Tensor &W, Tensor &b, |
| 97 | const ConvHandle &ch) { |
| 98 | CHECK_EQ(x.device()->lang(), kCpp); |
| 99 | |
| 100 | CHECK(x.shape(1) == ch.channels && x.shape(2) == ch.height && |
| 101 | x.shape(3) == ch.width) |
| 102 | << "input sample shape should not change"; |
| 103 | |
| 104 | CHECK(W.shape(0) == ch.num_filters && W.shape(1) == ch.channels && |
| 105 | W.shape(2) == ch.kernel_h && W.shape(3) == ch.kernel_w) |
| 106 | << "weights shape should not change"; |
| 107 | |
| 108 | #ifdef USE_DNNL |
| 109 | DataType dtype = x.data_type(); |
| 110 | auto dev = x.device(); |
| 111 | |
| 112 | Shape shape{ch.batchsize, ch.num_filters, ch.conv_height, ch.conv_width}; |
| 113 | Tensor output(shape, dev, dtype); |
| 114 | |
| 115 | output.device()->Exec( |
| 116 | [output, x, &W, &b, &ch](Context *ctx) mutable { |
| 117 | using namespace dnnl; |
| 118 | using tag = memory::format_tag; |
| 119 | auto eng = ctx->dnnl_engine; |
| 120 | auto s = ctx->dnnl_stream; |
| 121 | auto dtype = dnnl::memory::data_type::f32; |
| 122 | |
| 123 | // dnnl design pattern |
| 124 | // xxx_user_xxx_memory(and its format tag) is defined by user, which may |
| 125 | // need to be reordered |
| 126 | auto conv_user_src_memory = memory({{ch.x_dims}, dtype, tag::nchw}, eng, |
| 127 | x.block()->mutable_data()); |
| 128 | auto conv_user_weights_memory = memory({{ch.w_dims}, dtype, tag::goihw}, |
| 129 | eng, W.block()->mutable_data()); |
| 130 | auto conv_user_bias_memory = memory({{ch.b_dims}, dtype, tag::x}, eng, |
| 131 | b.block()->mutable_data()); |
| 132 | |
| 133 | // xxx_xxx_memory_md is created for creating conv_desc, and format tag |
| 134 | // is defined as any |
| 135 | auto conv_src_md = memory::desc({ch.x_dims}, dtype, tag::any); |
| 136 | auto conv_bias_md = memory::desc({ch.b_dims}, dtype, tag::any); |
| 137 | auto conv_weights_md = memory::desc({ch.w_dims}, dtype, tag::any); |
| 138 | auto conv_dst_md = memory::desc({ch.o_dims}, dtype, |
| 139 | tag::nchw); // could not set to any |
| 140 | |
| 141 | auto conv_desc = convolution_forward::desc( |
| 142 | prop_kind::forward, algorithm::convolution_direct, conv_src_md, |
| 143 | conv_weights_md, conv_bias_md, conv_dst_md, ch.s_dims, ch.p_dims, |
| 144 | ch.p_dims); |
| 145 | auto conv_pd = convolution_forward::primitive_desc(conv_desc, eng); |
| 146 | |
| 147 | // auto conv_pd = *ch.conv_pd; // 1ms to 70 ms slower |
| 148 | |
| 149 | // memory placeholder for reorder |
| 150 | auto conv_src_memory = conv_user_src_memory; |
| 151 | auto conv_weights_memory = conv_user_weights_memory; |
| 152 | |
| 153 | // output memory |