conv_transpose_mkldnn_op.cc 11.5 KB
Newer Older
J
Jacek Czaja 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License. */

15
#include "boost/optional.hpp"
J
Jacek Czaja 已提交
16 17 18
#include "paddle/fluid/framework/data_layout_transform.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/memory/malloc.h"
19
#include "paddle/fluid/operators/conv_op.h"
J
Jacek Czaja 已提交
20 21 22 23 24 25 26 27 28 29 30 31
#include "paddle/fluid/platform/mkldnn_reuse.h"

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
using framework::DataLayout;

template <typename T>
class ConvTransposeMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
 public:
  void Compute(const paddle::framework::ExecutionContext& ctx) const override {
32 33 34
    PADDLE_ENFORCE_EQ(platform::is_cpu_place(ctx.GetPlace()), true,
                      paddle::platform::errors::PreconditionNotMet(
                          "Operator DNNL ConvTranspose must use CPUPlace"));
J
Jacek Czaja 已提交
35
    const bool is_test = ctx.Attr<bool>("is_test");
F
FDInSky 已提交
36 37 38 39
    PADDLE_ENFORCE_EQ(is_test, true,
                      platform::errors::InvalidArgument(
                          "ConvTransposeMKLDNN works only for inference. "
                          "Set is_test = True. but got is_test=False ."));
J
Jacek Czaja 已提交
40 41 42 43 44 45 46 47 48 49

    auto& dev_ctx =
        ctx.template device_context<paddle::platform::MKLDNNDeviceContext>();
    const auto& mkldnn_engine = dev_ctx.GetEngine();

    auto* input = ctx.Input<Tensor>("Input");
    auto* filter = ctx.Input<Tensor>("Filter");
    auto* bias = ctx.HasInput("Bias") ? ctx.Input<Tensor>("Bias") : nullptr;
    auto* output = ctx.Output<Tensor>("Output");

F
FDInSky 已提交
50 51 52 53
    PADDLE_ENFORCE_EQ(
        input->layout(), DataLayout::kMKLDNN,
        platform::errors::InvalidArgument(
            "Got wrong layout = %d for Input tensor.", input->layout()));
A
Adam 已提交
54
    PADDLE_ENFORCE_NE(input->format(), MKLDNNMemoryFormat::undef,
F
FDInSky 已提交
55 56 57 58 59 60 61 62
                      platform::errors::InvalidArgument(
                          "Got wrong format for Input tensor."));

    PADDLE_ENFORCE_EQ(
        filter->layout(), DataLayout::kMKLDNN,
        platform::errors::InvalidArgument(
            "The filter tensor's laytout should be %d, but got %d.",
            DataLayout::kMKLDNN, filter->layout()));
A
Adam 已提交
63
    PADDLE_ENFORCE_NE(filter->format(), MKLDNNMemoryFormat::undef,
F
FDInSky 已提交
64 65 66 67 68 69 70 71 72 73 74 75 76
                      platform::errors::InvalidArgument(
                          "Got wrong formats for Filter tensor."));

    PADDLE_ENFORCE_EQ(
        input->dims().size(), 4,
        platform::errors::InvalidArgument(
            "Input must be with 4 dimensions, i.e. NCHW. but got dimension =%d",
            input->dims().size()));
    PADDLE_ENFORCE_EQ(
        filter->dims().size(), 4,
        platform::errors::InvalidArgument("Filter must be with 4 dimensions, "
                                          "i.e. OIHW, but got dimension =%d",
                                          filter->dims().size()));
J
Jacek Czaja 已提交
77 78

    if (bias) {
F
FDInSky 已提交
79 80 81 82 83
      PADDLE_ENFORCE_EQ(
          bias->layout(), DataLayout::kMKLDNN,
          platform::errors::InvalidArgument(
              "The bias tensor's laytout should be %d, but got %d.",
              DataLayout::kMKLDNN, bias->layout()));
A
Adam 已提交
84
      PADDLE_ENFORCE_NE(bias->format(), MKLDNNMemoryFormat::undef,
F
FDInSky 已提交
85 86 87 88 89 90 91 92
                        platform::errors::InvalidArgument(
                            "Got wrong format for Bias tensor."));

      PADDLE_ENFORCE_EQ(
          bias->dims().size(), 1,
          platform::errors::InvalidArgument("Bias must only have 1 dimension, "
                                            "i.e. X, but got dimension = %d .",
                                            bias->dims().size()));
J
Jacek Czaja 已提交
93 94
    }

A
Adam 已提交
95 96 97 98 99 100 101 102 103
    std::vector<int> strides_temp = ctx.Attr<std::vector<int>>("strides");
    std::vector<int64_t> strides(begin(strides_temp), end(strides_temp));

    std::vector<int> paddings_temp = ctx.Attr<std::vector<int>>("paddings");
    std::vector<int64_t> paddings(begin(paddings_temp), end(paddings_temp));

    std::vector<int> dilations_temp = ctx.Attr<std::vector<int>>("dilations");
    std::vector<int64_t> dilations(begin(dilations_temp), end(dilations_temp));

J
Jacek Czaja 已提交
104
    int groups = ctx.Attr<int>("groups");
105 106
    std::string padding_algorithm = ctx.Attr<std::string>("padding_algorithm");

107 108 109 110 111
    PADDLE_ENFORCE_EQ(
        strides.size(), 2,
        platform::errors::Unimplemented(
            "Now we only support 2d oneDNN convolution transpose op"));

112 113 114 115 116 117
    auto input_dims = input->dims();
    auto data_dims = framework::slice_ddim(input_dims, 2, input_dims.size());
    auto filter_dims = filter->dims();
    auto filter_data_dims =
        framework::slice_ddim(filter_dims, 2, filter_dims.size());

A
Adam 已提交
118
    auto ksize = framework::vectorize(filter_data_dims);
119 120 121

    UpdatePaddingAndDilation(&paddings, &dilations, padding_algorithm,
                             data_dims, strides, ksize);
J
Jacek Czaja 已提交
122

123 124
    std::transform(dilations.begin(), dilations.end(), dilations.begin(),
                   [](int64_t i) { return i - 1; });
J
Jacek Czaja 已提交
125 126 127 128

    const T* input_data = input->data<T>();
    const T* filter_data = filter->data<T>();

A
Adam 已提交
129 130 131
    auto src_tz = paddle::framework::vectorize<int64_t>(input->dims());
    auto iohw_weights_tz =
        paddle::framework::vectorize<int64_t>(filter->dims());
132 133
    auto weights_tz = iohw_weights_tz;

J
Jacek Czaja 已提交
134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171
    // IOHW -> OIHW
    weights_tz[0] = iohw_weights_tz[1];
    weights_tz[1] = iohw_weights_tz[0];

    // Custom Reorder from IOHW to OIHW
    auto iohw2oihw_reorder =
        [&iohw_weights_tz](const T* filter_data) -> std::shared_ptr<T> {
      int o = iohw_weights_tz[1];
      int c = iohw_weights_tz[0];
      int h = iohw_weights_tz[2];
      int w = iohw_weights_tz[3];
      std::shared_ptr<T> reordered_filter_data(new T[o * c * h * w](),
                                               std::default_delete<T[]>());
      for (int i = 0; i < c; ++i) {
        for (int j = 0; j < o; ++j) {
          int in_offset = j * h * w + i * o * h * w;
          int out_offset = j * c * h * w + i * h * w;
          std::memcpy(&(reordered_filter_data.get())[out_offset],
                      &filter_data[in_offset], h * w * sizeof(T));
        }
      }

      return reordered_filter_data;
    };

    int g = std::max(groups, 1);
    if (g > 1) {
      int o = weights_tz[0];
      int i = weights_tz[1];
      int h = weights_tz[2];
      int w = weights_tz[3];
      weights_tz.resize(5);
      weights_tz[0] = g;
      weights_tz[1] = o / g;
      weights_tz[2] = i;
      weights_tz[3] = h;
      weights_tz[4] = w;
    }
A
Adam 已提交
172
    auto dst_tz = paddle::framework::vectorize<int64_t>(output->dims());
J
Jacek Czaja 已提交
173 174

    // Get unique name for storing MKLDNN primitives
175
    const std::string key =
176
        platform::CreateKey(dev_ctx, src_tz, ctx.OutputName("Output"));
J
Jacek Czaja 已提交
177 178 179 180 181

    std::vector<mkldnn::primitive> pipeline;

    auto user_src_md = platform::MKLDNNMemDesc(
        {src_tz}, platform::MKLDNNGetDataType<T>(), input->format());
182 183 184
    auto user_weights_md = platform::MKLDNNMemDesc(
        {weights_tz}, platform::MKLDNNGetDataType<T>(),
        (g == 1) ? MKLDNNMemoryFormat::oihw : MKLDNNMemoryFormat::goihw);
J
Jacek Czaja 已提交
185 186 187 188 189

    /* create memory descriptor for convolution without specified format
     * ('any') which lets a primitive (convolution in this case) choose
     * the memory format preferred for best performance
     */
190
    auto chosen_memory_format = MKLDNNMemoryFormat::any;
191 192 193
    std::string fuse_activation = ctx.Attr<std::string>("fuse_activation");
    float fuse_alpha = ctx.Attr<float>("fuse_alpha");
    float fuse_beta = ctx.Attr<float>("fuse_beta");
J
Jacek Czaja 已提交
194 195 196 197 198

    auto src_md = platform::MKLDNNMemDesc(
        src_tz, platform::MKLDNNGetDataType<T>(), chosen_memory_format);
    auto weights_md = platform::MKLDNNMemDesc(
        weights_tz, platform::MKLDNNGetDataType<T>(), chosen_memory_format);
A
Adam 已提交
199
    std::vector<int64_t> bias_tz;
J
Jacek Czaja 已提交
200 201 202
    auto dst_md = platform::MKLDNNMemDesc(
        dst_tz, platform::MKLDNNGetDataType<T>(), chosen_memory_format);

203
    platform::ConvTransposeMKLDNNHandler handler(dev_ctx, mkldnn_engine, key);
J
Jacek Czaja 已提交
204 205 206 207 208 209 210
    // create a deconv(conv transpose) primitive descriptor and save it for
    // usage in backward
    std::shared_ptr<mkldnn::deconvolution_forward::primitive_desc>
        conv_transpose_pd;
    auto fwd_prop_kind = is_test ? mkldnn::prop_kind::forward_inference
                                 : mkldnn::prop_kind::forward_training;
    if (bias) {
A
Adam 已提交
211
      bias_tz = paddle::framework::vectorize<int64_t>(bias->dims());
J
Jacek Czaja 已提交
212
      auto bias_md = platform::MKLDNNMemDesc(
213
          bias_tz, platform::MKLDNNGetDataType<T>(), MKLDNNMemoryFormat::x);
214
      conv_transpose_pd = handler.AcquireConvolutionPrimitiveDescriptor(
215 216 217
          src_md, weights_md, bias_md, dst_md, strides, dilations, paddings,
          mkldnn_engine, fuse_activation, fuse_alpha, fuse_beta, false,
          fwd_prop_kind);
J
Jacek Czaja 已提交
218
    } else {
219
      conv_transpose_pd = handler.AcquireConvolutionPrimitiveDescriptor(
220
          src_md, weights_md, boost::none, dst_md, strides, dilations, paddings,
221 222
          mkldnn_engine, fuse_activation, fuse_alpha, fuse_beta, false,
          fwd_prop_kind);
J
Jacek Czaja 已提交
223 224 225 226 227 228 229 230 231 232 233 234 235 236 237
    }

    // create mkldnn memory from input tensors (data/weights)
    auto user_src_memory_p = handler.AcquireSrcMemory(
        user_src_md, platform::to_void_cast<T>(input_data));
    auto user_weights_memory_p = handler.AcquireWeightsMemory(
        user_weights_md, platform::to_void_cast<T>(filter_data),
        is_test ? iohw2oihw_reorder : platform::user_function());

    // create reorder primitive if the input format is not the preferred one
    auto src_memory_p =
        handler.AcquireSrcMemoryFromPrimitive(user_src_memory_p, pipeline);
    auto weights_memory_p = handler.AcquireWeightsMemoryFromPrimitive(
        user_weights_memory_p, pipeline, is_test);

238 239
    auto output_data =
        output->mutable_data<T>(ctx.GetPlace(), handler.GetDstMemorySize());
A
Adam 已提交
240
    auto dst_memory_p = handler.AcquireDstMemoryFromPrimitive(
J
Jacek Czaja 已提交
241 242
        platform::to_void_cast<T>(output_data));

A
Adam 已提交
243 244
    auto conv_p = handler.AcquireConvolution();

245
    auto& astream = platform::MKLDNNDeviceContext::tls().get_stream();
J
Jacek Czaja 已提交
246 247
    if (bias) {
      const T* bias_data = bias->data<T>();
248 249
      auto user_bias_md = platform::MKLDNNMemDesc(
          {bias_tz}, platform::MKLDNNGetDataType<T>(), MKLDNNMemoryFormat::x);
J
Jacek Czaja 已提交
250 251 252 253 254
      auto user_bias_memory_p = handler.AcquireBiasMemory(
          user_bias_md, platform::to_void_cast<T>(bias_data));

      auto bias_memory_p =
          handler.AcquireBiasMemoryFromPrimitive(user_bias_memory_p, pipeline);
A
Adam 已提交
255 256 257 258 259

      conv_p->execute(astream, {{MKLDNN_ARG_SRC, *src_memory_p},
                                {MKLDNN_ARG_WEIGHTS, *weights_memory_p},
                                {MKLDNN_ARG_BIAS, *bias_memory_p},
                                {MKLDNN_ARG_DST, *dst_memory_p}});
J
Jacek Czaja 已提交
260
    } else {
A
Adam 已提交
261 262 263
      conv_p->execute(astream, {{MKLDNN_ARG_SRC, *src_memory_p},
                                {MKLDNN_ARG_WEIGHTS, *weights_memory_p},
                                {MKLDNN_ARG_DST, *dst_memory_p}});
J
Jacek Czaja 已提交
264
    }
A
Adam 已提交
265
    astream.wait();
J
Jacek Czaja 已提交
266

267 268
    output->set_layout(DataLayout::kMKLDNN);
    output->set_format(platform::GetMKLDNNFormat(*dst_memory_p));
J
Jacek Czaja 已提交
269 270 271 272 273 274 275 276 277 278
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

REGISTER_OP_KERNEL(conv2d_transpose, MKLDNN, ::paddle::platform::CPUPlace,
                   ops::ConvTransposeMKLDNNOpKernel<float>);