/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/
#include <cnnl.h>
#include <torch/extension.h>
#include "kernels/embedding.mluh"
#include "register_api.h"
#include "utils.h"

namespace tmo {
namespace kernel_test_api {
at::Tensor embedding_test(const at::Tensor &input,
                          const at::Tensor &weight,
                          const int vocab_offset,
                          const int vocab_part) {
  MLU_TENSOR_CHECK_FATAL(input);
  MLU_TENSOR_CHECK_FATAL(weight);

  cnnlHandle_t handle = torch_mlu::getCurrentHandle();
  cnrtQueue_t queue;
  cnnlGetQueue(handle, &queue);
  std::shared_ptr<torch::Device> dev = std::make_shared<torch::Device>(input.device());
  int batch = input.size(0);
  int seq = input.size(1);
  int total_vocab_size = weight.size(0);
  int hidden_size = weight.size(1);
  cnnlDataType_t dtype = torch2cnnlDataType(torch::typeMetaToScalarType(weight.dtype()));
  torch::Tensor output = torch::zeros(
      {batch, seq, hidden_size},
      torch::dtype(torch::typeMetaToScalarType(weight.dtype())).device(*dev).requires_grad(false));

  TMO_KERNEL_CHECK_FATAL(invokeEmbedding(queue, weight.data_ptr(), input.data_ptr(),
                                         output.data_ptr(), dtype, vocab_offset, vocab_part,
                                         total_vocab_size, hidden_size, batch * seq));

  return output;
}

}  // namespace kernel_test_api
}  // namespace tmo
