common.cpp

/*************************************************************************
 * Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See LICENSE for license information.
 ************************************************************************/

#include "common.h"

namespace transformer_engine {
namespace paddle_ext {

TensorWrapper MakeNvteTensor(const void *data_ptr, const std::vector<size_t> &shape,
                             const DType type) {
  return TensorWrapper(const_cast<void *>(data_ptr), shape, type);
}

TensorWrapper MakeNvteTensor(void *data_ptr, const NVTEShape &shape, const DType type) {
  return TensorWrapper(data_ptr, shape, type);
}

TensorWrapper MakeNvteTensor(void *data_ptr, const std::vector<size_t> &shape, const DType type,
                             void *amax_ptr, void *scale_ptr, void *scale_inv_ptr) {
  return TensorWrapper(data_ptr, shape, type, reinterpret_cast<float *>(amax_ptr),
                       reinterpret_cast<float *>(scale_ptr),
                       reinterpret_cast<float *>(scale_inv_ptr));
}

TensorWrapper MakeNvteTensor(paddle::Tensor &tensor) {  // NOLINT
  return MakeNvteTensor(tensor.data(), GetShapeArray(tensor), Paddle2NvteDType(tensor.dtype()));
}

TensorWrapper MakeNvteTensor(const paddle::Tensor &tensor) {
  return MakeNvteTensor(const_cast<void *>(tensor.data()), GetShapeArray(tensor),
                        Paddle2NvteDType(tensor.dtype()));
}

paddle::Tensor AllocateSpace(const NVTEShape &shape, const DType type, const paddle::Place &place,
                             bool init_to_zeros) {
  auto size = shape.ndim;
  if (size == 2 && init_to_zeros) {
    return paddle::zeros({static_cast<int64_t>(shape.data[0]), static_cast<int64_t>(shape.data[1])},
                         Nvte2PaddleDType(type), place);
  } else if (size == 2) {
    return paddle::empty({static_cast<int64_t>(shape.data[0]), static_cast<int64_t>(shape.data[1])},
                         Nvte2PaddleDType(type), place);
  } else if (size == 1 && init_to_zeros) {
    return paddle::zeros({static_cast<int64_t>(shape.data[0])}, Nvte2PaddleDType(type), place);
  } else if (size == 1) {
    return paddle::empty({static_cast<int64_t>(shape.data[0])}, Nvte2PaddleDType(type), place);
  }
  NVTE_CHECK(false, "Should never reach here! func: AllocateSpace");
}

// MHA utils
// convert QKV layout to enum
NVTE_QKV_Layout get_nvte_qkv_layout(const std::string &qkv_layout) {
  static const std::unordered_map<std::string, NVTE_QKV_Layout> layout_map = {
      {"sb3hd", NVTE_QKV_Layout::NVTE_SB3HD},
      {"sbh3d", NVTE_QKV_Layout::NVTE_SBH3D},
      {"sbhd_sb2hd", NVTE_QKV_Layout::NVTE_SBHD_SB2HD},
      {"sbhd_sbh2d", NVTE_QKV_Layout::NVTE_SBHD_SBH2D},
      {"sbhd_sbhd_sbhd", NVTE_QKV_Layout::NVTE_SBHD_SBHD_SBHD},
      {"bs3hd", NVTE_QKV_Layout::NVTE_BS3HD},
      {"bsh3d", NVTE_QKV_Layout::NVTE_BSH3D},
      {"bshd_bs2hd", NVTE_QKV_Layout::NVTE_BSHD_BS2HD},
      {"bshd_bsh2d", NVTE_QKV_Layout::NVTE_BSHD_BSH2D},
      {"bshd_bshd_bshd", NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD},
      {"t3hd", NVTE_QKV_Layout::NVTE_T3HD},
      {"th3d", NVTE_QKV_Layout::NVTE_TH3D},
      {"thd_t2hd", NVTE_QKV_Layout::NVTE_THD_T2HD},
      {"thd_th2d", NVTE_QKV_Layout::NVTE_THD_TH2D},
      {"thd_thd_thd", NVTE_QKV_Layout::NVTE_THD_THD_THD},
  };

  auto it = layout_map.find(qkv_layout);
  if (it != layout_map.end()) {
    return it->second;
  } else {
    NVTE_ERROR("Invalid QKV layout string: " + qkv_layout);
  }
}

}  // namespace paddle_ext
}  // namespace transformer_engine