transfer_layout_kernel.cc 6.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* Copyright (c) 2022 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 "paddle/phi/kernels/transfer_layout_kernel.h"
16

17 18 19
#include <sstream>
#include <string>

20 21
#include "paddle/phi/backends/all_context.h"
#include "paddle/phi/core/kernel_registry.h"
22
#include "paddle/phi/core/visit_type.h"
23
#include "paddle/phi/kernels/funcs/data_layout_transform.h"
24
#include "paddle/phi/kernels/funcs/math_function.h"
25
#ifdef PADDLE_WITH_MKLDNN
26
#include "paddle/phi/backends/onednn/onednn_helper.h"
27
#endif
28
namespace phi {
29 30 31 32 33

std::vector<int> GetAxis(const DataLayout& from, const DataLayout& to) {
  PADDLE_ENFORCE_NE(
      from,
      to,
34
      phi::errors::InvalidArgument(
35 36 37 38 39 40
          "Layout transform should transform between different layout."));
  if (from == DataLayout::NCHW && to == DataLayout::NHWC) {
    return {0, 2, 3, 1};
  } else if (from == DataLayout::NHWC && to == DataLayout::NCHW) {
    return {0, 3, 1, 2};
  } else {
41
    PADDLE_THROW(phi::errors::InvalidArgument("Unsupported layout transform."));
42 43 44 45 46 47 48 49
  }
}

template <typename T, typename Context>
void CastDataLayout(const Context& dev_ctx,
                    const DenseTensor& x,
                    const std::vector<int>& axis,
                    DenseTensor* out) {
50
  funcs::Transpose<Context, T, 4> trans4;
51 52 53 54
  trans4(dev_ctx, x, out, axis);
}

template <typename Context>
55 56 57 58
void TransferLayoutGeneral(const Context& dev_ctx,
                           const DenseTensor& x,
                           DataLayout dst_layout,
                           DenseTensor* out) {
59 60 61 62 63 64 65 66 67 68
  auto src_dim = x.dims();

  auto axis = GetAxis(x.layout(), dst_layout);

  std::vector<int64_t> dst_dim;
  dst_dim.resize(axis.size());
  for (size_t i = 0; i < axis.size(); i++) {
    dst_dim[i] = src_dim[axis[i]];
  }

69 70
  out->Resize(phi::make_ddim(dst_dim));
  dev_ctx.Alloc(out, x.dtype());
71 72 73 74 75 76

  PD_VISIT_ALL_TYPES(x.dtype(), "CastDataLayout", ([&] {
                       CastDataLayout<data_t, Context>(dev_ctx, x, axis, out);
                     }));
}

77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111
#ifdef PADDLE_WITH_MKLDNN
template <typename Context>
void TransferLayoutMKLDNN(const Context& dev_ctx,
                          const DenseTensor& x,
                          DataLayout src_layout,
                          DataLayout dst_layout,
                          DenseTensor* out) {
  auto print_tensor_meta = [](const DenseTensor& x) {
    std::ostringstream oss;

    oss << "[";
    oss << "layout:" << x.layout() << " ,";
    oss << "dims:" << x.dims() << " ,";
    if (x.IsInitialized()) oss << "place:" << x.place();
    oss << "]";

    return oss.str();
  };
  VLOG(10) << " x: " << print_tensor_meta(x);
  VLOG(10) << " out: " << print_tensor_meta(*out) << " " << out;

  // NOTE(zhiqiu): to handle the special case in ApplyDataTransform() in
  // data_transfer.cc
  if (!x.IsInitialized() && src_layout == DataLayout::MKLDNN &&
      dst_layout == DataLayout::NHWC) {
    VLOG(4) << src_layout << "->" << dst_layout << " " << x.layout();
    out->Resize(x.dims());
    out->set_layout(dst_layout);
    funcs::MatchShapeToLayout(out, src_layout, dst_layout);
    return;
  }

  if (src_layout != DataLayout::MKLDNN && dst_layout == DataLayout::MKLDNN) {
    // Case1 - transform from Non-MKLDNN OPKernel to MKLDNN OPKernel
    // Just set layout/format. No real transform occur
112 113
    auto out_format = funcs::OneDNNFormatForSize(
        x.dims().size(), funcs::ToOneDNNFormat(src_layout));
114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129

    out->ShareDataWith(x);
    // For NHWC data we need reshape of tensors as MKL-DNN
    // is expecting NHWC dims description order
    if (src_layout == DataLayout::NHWC) {
      VLOG(4) << "NHWC";
      funcs::MatchShapeToLayout(out, src_layout, dst_layout);
      OneDNNContext::tls().set_cur_paddle_data_layout(src_layout);
    }

    out->set_layout(DataLayout::MKLDNN);
    out->set_format(out_format);
  } else if (src_layout == DataLayout::MKLDNN &&
             dst_layout != DataLayout::MKLDNN) {
    // Case2 - transfrom from MKLDNN OPKernel to Non-MKLDNN OPKernel
    // Do transform via MKLDNN lib
130
    funcs::innerTransDataLayoutFromOneDNN(
131 132 133 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
        src_layout, dst_layout, x, out, dev_ctx.GetPlace());
  } else if (src_layout == DataLayout::MKLDNN &&
             dst_layout == DataLayout::MKLDNN) {
    PADDLE_ENFORCE_NE(
        src_layout,
        dst_layout,
        errors::PreconditionNotMet(
            "No layout transform needed between two MKLDNN OPKernels."));
  } else {
    TransferLayoutGeneral<Context>(dev_ctx, x, dst_layout, out);
  }
}
#endif

template <typename Context>
void TransferLayoutKernel(const Context& dev_ctx,
                          const DenseTensor& x,
                          int src_layout,
                          int dst_layout,
                          DenseTensor* out) {
  PADDLE_ENFORCE_NE(src_layout,
                    dst_layout,
                    errors::PreconditionNotMet(
                        "No layout transform needed between same layout."));
  VLOG(10) << "TransDataLayout from " << static_cast<DataLayout>(src_layout)
           << " -> " << static_cast<DataLayout>(dst_layout);

#ifdef PADDLE_WITH_MKLDNN
  TransferLayoutMKLDNN<Context>(dev_ctx,
                                x,
                                static_cast<DataLayout>(src_layout),
                                static_cast<DataLayout>(dst_layout),
                                out);
#else
  TransferLayoutGeneral<Context>(
      dev_ctx, x, static_cast<DataLayout>(dst_layout), out);
#endif
}

170
}  // namespace phi
171

172
PD_REGISTER_GENERAL_KERNEL(transfer_layout,
173 174
                           CPU,
                           ALL_LAYOUT,
175
                           phi::TransferLayoutKernel<phi::CPUContext>,
176
                           ALL_DTYPE) {}