data_transform.h 6.4 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* 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. */

#pragma once

17 18
#include "paddle/phi/api/include/tensor.h"
#include "paddle/phi/core/kernel_factory.h"
19
#include "paddle/phi/core/selected_rows.h"
W
wanghuancoder 已提交
20 21
#include "paddle/phi/core/sparse_coo_tensor.h"
#include "paddle/phi/core/sparse_csr_tensor.h"
22

23 24 25 26 27 28
namespace phi {
namespace distributed {
class DistTensor;
}  // namespace distributed
}  // namespace phi

29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66
namespace paddle {
namespace experimental {

class TransformFlag {
 public:
  TransformFlag(bool stop_transform = false,
                bool trans_dtype = false,
                bool trans_backend = true,
                bool trans_layout = true)
      : stop_transform_(stop_transform),
        trans_data_type_(trans_dtype),
        trans_backend_(trans_backend),
        trans_layout_(trans_layout) {}

  bool NeedTransform() const {
    return !stop_transform_ &&
           (trans_data_type_ || trans_backend_ || trans_layout_);
  }

  bool need_trans_data_type() const {
    return !stop_transform_ && trans_data_type_;
  }

  bool need_trans_backend() const { return !stop_transform_ && trans_backend_; }

  bool need_trans_layout() const { return !stop_transform_ && trans_layout_; }

 private:
  // This is the highest priority in flags,
  // and can be setted by api[data_transform->skip_transform] in the yaml file.
  bool stop_transform_ = false;

  // trans_data_type_ can be setted by api[data_transform->support_trans_dtype]
  // in the yaml file.
  // trans_data_type_ only affect the non complex types,
  // the complex is always transferd, except stop_transform_ is true.
  bool trans_data_type_ = false;

67
  // trans_backend_ and trans_layout_ are true defaultly,
68 69 70 71 72
  // and they can only be setted by global flag.
  bool trans_backend_ = true;
  bool trans_layout_ = true;
};

73 74 75 76 77 78 79 80 81 82 83 84 85
static inline phi::TensorArgDef GetKernelInputArgDef(
    const phi::TensorArgDef& input_def, phi::Backend kernel_backend) {
  phi::TensorArgDef input_actual_def = input_def;
#ifdef PADDLE_WITH_CUSTOM_DEVICE
  // When the backend of input tensor arg_def is CUSTOM, we need to set it to
  // the actual backend by expected_kernel_key.
  if (input_actual_def.backend == phi::Backend::CUSTOM) {
    input_actual_def.SetBackend(kernel_backend);
  }
#endif
  return input_actual_def;
}

86
std::shared_ptr<phi::DenseTensor> PrepareData(
87
    const Tensor& input,
88
    const phi::TensorArgDef& target_args_def,
W
wanghuancoder 已提交
89 90
    const TransformFlag& transform_flag,
    bool is_stride_kernel);
91

92
paddle::optional<phi::DenseTensor> PrepareData(
93 94
    const paddle::optional<Tensor>& input,
    const phi::TensorArgDef& target_args_def,
W
wanghuancoder 已提交
95 96
    const TransformFlag& transform_flag,
    bool is_stride_kernel);
97

98
std::unique_ptr<std::vector<phi::DenseTensor>> PrepareData(
99
    const std::vector<Tensor>& inputs,
100
    const phi::TensorArgDef& target_args_def,
W
wanghuancoder 已提交
101 102
    const TransformFlag& transform_flag,
    bool is_stride_kernel);
103

104 105 106
paddle::optional<std::vector<phi::DenseTensor>> PrepareData(
    const paddle::optional<std::vector<Tensor>>& inputs,
    const phi::TensorArgDef& target_args_def,
W
wanghuancoder 已提交
107 108
    const TransformFlag& transform_flag,
    bool is_stride_kernel);
109

110 111 112 113 114 115 116 117 118 119 120
// Only support transfering place for SelectedRows
std::shared_ptr<phi::SelectedRows> PrepareDataForSelectedRows(
    const Tensor& input,
    const phi::TensorArgDef& target_args_def,
    const TransformFlag& transform_flag);

paddle::optional<phi::SelectedRows> PrepareDataForSelectedRows(
    const paddle::optional<Tensor>& input,
    const phi::TensorArgDef& target_args_def,
    const TransformFlag& transform_flag);

W
wanghuancoder 已提交
121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141
// Only support transfering contiguous for SparseCooTensor
std::shared_ptr<phi::SparseCooTensor> PrepareDataForSparseCooTensor(
    const Tensor& input);

paddle::optional<phi::SparseCooTensor> PrepareDataForSparseCooTensor(
    const paddle::optional<Tensor>& input);

// Only support transfering contiguous for SparseCsrTensor
std::shared_ptr<phi::SparseCsrTensor> PrepareDataForSparseCsrTensor(
    const Tensor& input);

paddle::optional<phi::SparseCsrTensor> PrepareDataForSparseCsrTensor(
    const paddle::optional<Tensor>& input);

// Only support transfering contiguous
std::shared_ptr<phi::DenseTensor> PrepareDataForDenseTensorInSparse(
    const Tensor& input);

paddle::optional<phi::DenseTensor> PrepareDataForDenseTensorInSparse(
    const paddle::optional<Tensor>& input);

142 143 144 145 146 147 148 149 150 151 152 153
void TransDataBackend(const phi::DenseTensor* tensor,
                      Backend target_backend,
                      phi::DenseTensor* out);

void TransDataBackend(const std::vector<phi::DenseTensor*>& tensor,
                      Backend target_backend,
                      std::vector<phi::DenseTensor*> out);

void TransDataBackend(const phi::SelectedRows* tensor,
                      Backend target_backend,
                      phi::SelectedRows* out);

W
wanghuancoder 已提交
154 155 156
phi::DenseTensor Trans2Contiguous(const phi::DenseTensor& tensor);

void CheckAndTrans2Contiguous(phi::DenseTensor* tensor);
H
hong 已提交
157
inline bool NeedTransformPlace(const phi::Place& src_place,
H
hong 已提交
158 159 160 161 162 163 164 165 166
                               const Backend& target,
                               const TransformFlag& transform_flag) {
  // NOTE(dev): The default value of TransformFlag is True, if it is set with
  // False
  // somewhere such as ops.yaml or backward.yaml that means we should skip data
  // transform. Because "stop_transform_" has highest priority.
  if (!transform_flag.need_trans_backend()) {
    return false;
  }
H
hong 已提交
167
  bool ret = src_place.GetType() == AllocationType::GPUPINNED ||
H
hong 已提交
168
             (target != Backend::ALL_BACKEND &&
H
hong 已提交
169
              phi::TransToPhiBackend(src_place) !=
H
hong 已提交
170 171 172 173
                  (target != Backend::GPUDNN ? target : Backend::GPU));
  return ret;
}

174 175 176 177 178 179 180 181 182
/* ------------------ for auto parallel ----------------------- */

// TODO(chenweihang): impl Reshard input and output function
std::shared_ptr<phi::distributed::DistTensor> PrepareDataForDistTensor(
    const Tensor& input,
    const phi::TensorArgDef& target_args_def,
    const TransformFlag& transform_flag,
    bool is_stride_kernel);

183 184
}  // namespace experimental
}  // namespace paddle