//===-- LanguageLaunch.cpp - Language launch API --------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "LanguageLaunch.h"
#include "LanguageUtils.h"
#include "State.h"
#include <cstdio>

using RuntimeState = llvm::offload::StateTy;
using ThreadState = llvm::offload::ThreadStateTy;

static constexpr ol_error_struct_t InvalidKernelError = {
    OL_ERRC_INVALID_NULL_HANDLE, "kernel is not registered"};

static constexpr ol_error_struct_t InvalidDeviceError = {OL_ERRC_INVALID_DEVICE,
                                                         "invalid device"};

static constexpr ol_error_struct_t InvalidArgumentError = {
    OL_ERRC_INVALID_ARGUMENT, "invalid argument to kernel launch"};

static constexpr ol_error_struct_t InvalidConfigurationError = {
    OL_ERRC_INVALID_SIZE, "invalid kernel launch configuration"};

/// Internal kernel launch implementation
ol_result_t __llvmLaunchKernelImpl(const char *KernelID, dim3 GridDim,
                                   dim3 BlockDim, void *KernelArgsPtr,
                                   size_t DynamicSharedMem, void *Stream) {
  ol_device_handle_t Device = ThreadState::getDefaultDevice();
  ol_symbol_handle_t Kernel = RuntimeState::getKernel(KernelID);
  if (!Device)
    return &InvalidDeviceError;
  if (!KernelID || !KernelArgsPtr)
    return &InvalidArgumentError;
  if (!Kernel)
    return &InvalidKernelError;

  if (GridDim.x == 0 || GridDim.y == 0 || GridDim.z == 0 || BlockDim.x == 0 ||
      BlockDim.y == 0 || BlockDim.z == 0)
    return &InvalidConfigurationError;

  ol_kernel_launch_size_args_t LaunchSizeArgs;
  LaunchSizeArgs.Dimensions =
      1 + (GridDim.y > 1 || BlockDim.y > 1) + (GridDim.z > 1 || BlockDim.z > 1);
  LaunchSizeArgs.NumGroups.x = GridDim.x;
  LaunchSizeArgs.NumGroups.y = GridDim.y;
  LaunchSizeArgs.NumGroups.z = GridDim.z;
  LaunchSizeArgs.GroupSize.x = BlockDim.x;
  LaunchSizeArgs.GroupSize.y = BlockDim.y;
  LaunchSizeArgs.GroupSize.z = BlockDim.z;
  LaunchSizeArgs.DynSharedMemory = DynamicSharedMem;

  ol_queue_handle_t Queue = Stream ? reinterpret_cast<ol_queue_handle_t>(Stream)
                                   : ThreadState::getDefaultQueue();

  struct OffloadKernelArgs {
    void **Args;
    size_t NumArgs;
    size_t *ArgSizes;
  };
  OffloadKernelArgs *OKA = static_cast<OffloadKernelArgs *>(KernelArgsPtr);
  if ((!OKA->Args) != (!OKA->ArgSizes))
    return &InvalidArgumentError;
  if (OKA->NumArgs > 0 && !OKA->Args)
    return &InvalidArgumentError;
  for (size_t I = 0; I < OKA->NumArgs; ++I)
    if (!OKA->Args[I] || OKA->ArgSizes[I] == 0)
      return &InvalidArgumentError;

  return olLaunchKernel(Queue, Device, Kernel, &LaunchSizeArgs,
                        /*Properties=*/nullptr, OKA->NumArgs, OKA->Args,
                        OKA->ArgSizes);
}

extern "C" {

/// Push call configuration for kernel launch
unsigned __llvmPushCallConfiguration(dim3 GridSize, dim3 BlockSize,
                                     size_t SharedMemory, void *Stream) {
  CallConfigurationTy &CC = ThreadState::getCallConfiguration();

  CC.GridSize = GridSize;
  CC.BlockSize = BlockSize;
  CC.SharedMemory = SharedMemory;
  CC.Stream = Stream;
  return 0;
}

/// Pop call configuration for kernel launch
unsigned __llvmPopCallConfiguration(dim3 *GridSize, dim3 *BlockSize,
                                    size_t *SharedMemory, void **Stream) {
  CallConfigurationTy &CC = ThreadState::getCallConfiguration();
  *GridSize = CC.GridSize;
  *BlockSize = CC.BlockSize;
  *SharedMemory = CC.SharedMemory;
  *Stream = CC.Stream;
  return 0;
}

unsigned llvmLaunchKernel(const char *KernelID, dim3 GridDim, dim3 BlockDim,
                          void *KernelArgsPtr, size_t DynamicSharedMem,
                          void *Stream) {
  ol_result_t Result = __llvmLaunchKernelImpl(
      KernelID, GridDim, BlockDim, KernelArgsPtr, DynamicSharedMem, Stream);
  convertAndSetLastError(Result);
  return Result ? Result->Code : 0;
}

} // extern "C"