paddlepaddle--paddle/paddle/phi/api/lib/data_transform.h

/* 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

#include "paddle/phi/api/include/tensor.h"
#include "paddle/phi/core/distributed/type_defs.h"
#include "paddle/phi/core/kernel_factory.h"
#include "paddle/phi/core/selected_rows.h"
#include "paddle/phi/core/sparse_coo_tensor.h"
#include "paddle/phi/core/sparse_csr_tensor.h"

namespace phi {
class DeviceContext;
namespace distributed {
class DistTensor;
class TensorDistAttr;
}  // namespace distributed
}  // namespace phi

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 set by api[data_transform->skip_transform] in the yaml file.
  bool stop_transform_ = false;

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

  // trans_backend_ and trans_layout_ are true default,
  // and they can only be set by global flag.
  bool trans_backend_ = true;
  bool trans_layout_ = true;
};

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;
}

std::shared_ptr<phi::DenseTensor> PrepareData(
    const Tensor& input,
    const phi::TensorArgDef& target_args_def,
    const TransformFlag& transform_flag,
    bool is_stride_kernel);

paddle::optional<phi::DenseTensor> PrepareData(
    const paddle::optional<Tensor>& input,
    const phi::TensorArgDef& target_args_def,
    const TransformFlag& transform_flag,
    bool is_stride_kernel);

std::unique_ptr<std::vector<phi::DenseTensor>> PrepareData(
    const std::vector<Tensor>& inputs,
    const phi::TensorArgDef& target_args_def,
    const TransformFlag& transform_flag,
    bool is_stride_kernel);

paddle::optional<std::vector<phi::DenseTensor>> PrepareData(
    const paddle::optional<std::vector<Tensor>>& inputs,
    const phi::TensorArgDef& target_args_def,
    const TransformFlag& transform_flag,
    bool is_stride_kernel);

// Only support transferring 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);

// Only support transferring contiguous for SparseCooTensor
std::shared_ptr<phi::SparseCooTensor> PrepareDataForSparseCooTensor(
    const Tensor& input);

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

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

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

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

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

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);

PADDLE_API phi::DenseTensor Trans2Contiguous(const phi::DenseTensor& tensor);

PADDLE_API void CheckAndTrans2Contiguous(phi::DenseTensor* tensor);

phi::DenseTensor CheckAndTrans2NewContiguousTensor(
    const phi::DenseTensor& tensor);

std::vector<phi::DenseTensor> CheckAndTrans2NewContiguousTensor(
    const std::vector<phi::DenseTensor>& tensor);

inline bool NeedTransformPlace(const phi::Place& src_place,
                               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;
  }

#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
  bool ret = src_place.GetType() == AllocationType::GPUPINNED ||
             (target != Backend::ALL_BACKEND &&
              phi::TransToPhiBackend(src_place) !=
                  (target != Backend::GPUDNN ? target : Backend::GPU));
#elif defined(PADDLE_WITH_XPU)
  bool ret = src_place.GetType() == AllocationType::XPUPINNED ||
             (target != Backend::ALL_BACKEND &&
              phi::TransToPhiBackend(src_place) != target);
#elif defined(PADDLE_WITH_IPU)
  bool ret = target != Backend::ALL_BACKEND &&
             phi::TransToPhiBackend(src_place) != target;
#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
  bool ret = target != Backend::ALL_BACKEND;
  if (target == Backend::CUSTOM) {
    ret = ret && !is_custom_place(src_place);
  } else {
    ret =
        ret && phi::TransToPhiBackend(src_place) !=
                   (target != Backend::GPUDNN ? target
                                              : Backend::DEFAULT_CUSTOM_DEVICE);
  }
#else
  bool ret = false;
#endif

#ifdef PADDLE_WITH_DNNL
  if (target == Backend::ONEDNN) {
    ret = src_place.GetType() != AllocationType::CPU;
  }
#endif
  return ret;
}

/* ------------------ for auto parallel ----------------------- */

std::shared_ptr<phi::distributed::DistTensor> ReshardApiInputToKernelInput(
    phi::DeviceContext* dev_ctx,
    const Tensor& tensor,
    const phi::distributed::ArgDistAttr& dist_attr,
    const std::string& arg_name = "");

std::vector<std::shared_ptr<phi::distributed::DistTensor>>
ReshardApiInputToKernelInput(phi::DeviceContext* dev_ctx,
                             const std::vector<Tensor>& tensor,
                             const phi::distributed::ArgDistAttr& dist_attr,
                             const std::string& arg_name = "");

paddle::optional<std::shared_ptr<phi::distributed::DistTensor>>
ReshardApiInputToKernelInput(phi::DeviceContext* dev_ctx,
                             const paddle::optional<Tensor>& tensor,
                             const phi::distributed::ArgDistAttr& dist_attr,
                             const std::string& arg_name = "");

paddle::optional<std::vector<std::shared_ptr<phi::distributed::DistTensor>>>
ReshardApiInputToKernelInput(
    phi::DeviceContext* dev_ctx,
    const paddle::optional<std::vector<Tensor>>& tensors,
    const phi::distributed::ArgDistAttr& dist_attr,
    const std::string& arg_name = "");

void SetInplaceOutputCorrectDistAttr(
    phi::DeviceContext* dev_ctx,
    Tensor& tensor,  // NOLINT
    const phi::distributed::TensorDistAttr& dist_attr,
    bool use_general_spmd_rule = true);

void SetInplaceOutputCorrectDistAttr(
    phi::DeviceContext* dev_ctx,
    Tensor& tensor,  // NOLINT
    const phi::distributed::ArgDistAttr& dist_attr,
    bool use_general_spmd_rule = true);

void SetInplaceOutputCorrectDistAttr(
    phi::DeviceContext* dev_ctx,
    std::vector<Tensor>& tensors,  // NOLINT
    const std::vector<phi::distributed::TensorDistAttr>& dist_attr,
    bool use_general_spmd_rule = true);

void SetInplaceOutputCorrectDistAttr(
    phi::DeviceContext* dev_ctx,
    std::vector<Tensor>& tensors,  // NOLINT
    const phi::distributed::ArgDistAttr& dist_attr,
    bool use_general_spmd_rule = true);

void ReshardOutputPartialAxisToReplicated(
    phi::DeviceContext* dev_ctx, phi::distributed::DistTensor* out_tensor);

void ReshardKernelOutputToApiOutput(
    phi::DeviceContext* dev_ctx,
    const std::shared_ptr<phi::distributed::DistTensor>& src_tensor,
    Tensor* dst_tensor,
    const std::string& arg_name = "");

void ReshardKernelOutputToApiOutput(
    phi::DeviceContext* dev_ctx,
    const std::vector<std::shared_ptr<phi::distributed::DistTensor>>&
        src_tensors,
    const std::vector<Tensor*>& dst_tensors,
    const std::string& arg_name = "");

std::shared_ptr<phi::distributed::DistTensor> PrepareDataForDistTensor(
    std::shared_ptr<phi::distributed::DistTensor> input,
    const phi::TensorArgDef& target_args_def,
    const TransformFlag& transform_flag,
    bool is_stride_kernel);

std::vector<std::shared_ptr<phi::distributed::DistTensor>>
PrepareDataForDistTensor(
    std::vector<std::shared_ptr<phi::distributed::DistTensor>> input,
    const phi::TensorArgDef& target_args_def,
    const TransformFlag& transform_flag,
    bool is_stride_kernel);

paddle::optional<std::shared_ptr<phi::distributed::DistTensor>>
PrepareDataForDistTensor(
    paddle::optional<std::shared_ptr<phi::distributed::DistTensor>> input,
    const phi::TensorArgDef& target_args_def,
    const TransformFlag& transform_flag,
    bool is_stride_kernel);

paddle::optional<std::vector<std::shared_ptr<phi::distributed::DistTensor>>>
PrepareDataForDistTensor(
    paddle::optional<std::vector<std::shared_ptr<phi::distributed::DistTensor>>>
        input,
    const phi::TensorArgDef& target_args_def,
    const TransformFlag& transform_flag,
    bool is_stride_kernel);

std::string ReshardDebugInfo(const phi::distributed::DistTensor& src_tensor,
                             const phi::distributed::TensorDistAttr& dist_attr);

}  // namespace experimental
}  // namespace paddle