//===------- Offload API tests - olLaunchKernel --------------------===//
//
// 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 "../common/Fixtures.hpp"
#include <OffloadAPI.h>
#include <gtest/gtest.h>

#define KERNEL_TEST(NAME, KERNEL)                                              \
  struct olLaunchKernel##NAME##Test : LaunchSingleKernelTestBase {             \
    void SetUp() override { SetUpKernel(#KERNEL); }                            \
  };                                                                           \
  OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olLaunchKernel##NAME##Test);

KERNEL_TEST(Foo, foo)
KERNEL_TEST(NoArgs, noargs)
KERNEL_TEST(MultiArgs, multiargs)
KERNEL_TEST(Composite, composite)
KERNEL_TEST(Byte, byte)
KERNEL_TEST(LocalMem, localmem)
KERNEL_TEST(LocalMemReduction, localmem_reduction)
KERNEL_TEST(LocalMemStatic, localmem_static)
KERNEL_TEST(SingleCounterSyncEvent, single_counter)
KERNEL_TEST(GlobalCtor, global_ctor)
KERNEL_TEST(GlobalDtor, global_dtor)
KERNEL_TEST(GridSize, gridsize)

struct LaunchMultipleKernelTestBase : LaunchKernelTestBase {
  void SetUpKernels(const char *program, std::vector<const char *> kernels) {
    RETURN_ON_FATAL_FAILURE(SetUpProgram(program));

    Kernels.resize(kernels.size());
    size_t I = 0;
    for (auto K : kernels)
      ASSERT_SUCCESS(
          olGetSymbol(Program, K, OL_SYMBOL_KIND_KERNEL, &Kernels[I++]));
  }

  std::vector<ol_symbol_handle_t> Kernels;
};

struct ArgsSingleCounter {
  int32_t InitLoop;
  int32_t Addend;
  uint32_t *InitVal;
  uint32_t *Out;
};

#define KERNEL_MULTI_TEST(NAME, PROGRAM, ...)                                  \
  struct olLaunchKernel##NAME##Test : LaunchMultipleKernelTestBase {           \
    void SetUp() override { SetUpKernels(#PROGRAM, {__VA_ARGS__}); }           \
  };                                                                           \
  OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olLaunchKernel##NAME##Test);

KERNEL_MULTI_TEST(Global, global, "write", "read")

TEST_P(olLaunchKernelFooTest, Success) {
  void *Mem;
  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_MANAGED,
                            LaunchArgs.GroupSize.x * sizeof(uint32_t), &Mem));

  void *ArgPtrs[] = {&Mem};
  size_t ArgSizes[] = {sizeof(Mem)};

  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));

  ASSERT_SUCCESS(olSyncQueue(Queue));

  uint32_t *Data = (uint32_t *)Mem;
  for (uint32_t i = 0; i < 64; i++) {
    ASSERT_EQ(Data[i], i);
  }

  ASSERT_SUCCESS(olMemFree(Context, Mem));
}

TEST_P(olLaunchKernelFooTest, SuccessThreaded) {
  threadify([&](size_t) {
    void *DevAlloc, *HstAlloc;
    size_t Size = LaunchArgs.GroupSize.x * sizeof(uint32_t);
    ASSERT_SUCCESS(
        olMemAlloc(Context, Device, OL_ALLOC_TYPE_DEVICE, Size, &DevAlloc));
    ASSERT_SUCCESS(olMemAllocHost(Context, Device, Size, &HstAlloc));

    void *ArgPtrs[] = {&DevAlloc};
    size_t ArgSizes[] = {sizeof(DevAlloc)};

    ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                  std::size(ArgPtrs), ArgPtrs, ArgSizes));

    ASSERT_SUCCESS(olMemcpy(Queue, HstAlloc, Host, DevAlloc, Device, Size));

    ASSERT_SUCCESS(olSyncQueue(Queue));

    uint32_t *Data = static_cast<uint32_t *>(HstAlloc);
    for (uint32_t i = 0; i < LaunchArgs.GroupSize.x; i++) {
      ASSERT_EQ(Data[i], i);
    }

    ASSERT_SUCCESS(olMemFree(Context, DevAlloc));
    ASSERT_SUCCESS(olMemFree(Context, HstAlloc));
  });
}

TEST_P(olLaunchKernelNoArgsTest, Success) {
  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr, 0,
                                nullptr, nullptr));

  ASSERT_SUCCESS(olSyncQueue(Queue));
}

TEST_P(olLaunchKernelMultiArgsTest, Success) {
  void *Mem;
  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_MANAGED,
                            LaunchArgs.GroupSize.x * sizeof(int), &Mem));

  char A = 3;
  int *B = (int *)Mem;
  short C = 5;

  void *ArgPtrs[] = {&A, &B, &C};
  size_t ArgSizes[] = {sizeof(A), sizeof(B), sizeof(C)};

  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));

  ASSERT_SUCCESS(olSyncQueue(Queue));

  int *Data = (int *)Mem;
  for (uint32_t i = 0; i < LaunchArgs.GroupSize.x; i++)
    ASSERT_EQ(Data[i], A + C + static_cast<int>(i));

  ASSERT_SUCCESS(olMemFree(Context, Mem));
}

struct Foo {
  uint32_t a;
  uint32_t b;
};

TEST_P(olLaunchKernelCompositeTest, Success) {
  void *Mem;
  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_MANAGED,
                            LaunchArgs.GroupSize.x * sizeof(uint32_t), &Mem));

  uint8_t N = 1;
  Foo F{2, 3};
  uint32_t *Out = (uint32_t *)Mem;

  void *ArgPtrs[] = {&N, &F, &Out};
  size_t ArgSizes[] = {sizeof(N), sizeof(F), sizeof(Out)};

  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));

  ASSERT_SUCCESS(olSyncQueue(Queue));

  uint32_t *Data = (uint32_t *)Mem;
  for (uint32_t i = 0; i < LaunchArgs.GroupSize.x; i++)
    ASSERT_EQ(Data[i], N + F.a + F.b + i);

  ASSERT_SUCCESS(olMemFree(Context, Mem));
}

TEST_P(olLaunchKernelFooTest, SuccessSynchronous) {
  void *Mem;
  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_MANAGED,
                            LaunchArgs.GroupSize.x * sizeof(uint32_t), &Mem));

  void *ArgPtrs[] = {&Mem};
  size_t ArgSizes[] = {sizeof(Mem)};

  ASSERT_SUCCESS(olLaunchKernel(nullptr, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));

  uint32_t *Data = (uint32_t *)Mem;
  for (uint32_t i = 0; i < 64; i++) {
    ASSERT_EQ(Data[i], i);
  }

  ASSERT_SUCCESS(olMemFree(Context, Mem));
}

TEST_P(olLaunchKernelByteTest, Success) {
  unsigned char C = 42;

  void *ArgPtrs[] = {&C};
  size_t ArgSizes[] = {sizeof(C)};

  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));

  ASSERT_SUCCESS(olSyncQueue(Queue));
}

TEST_P(olLaunchKernelLocalMemTest, Success) {
  SKIP_KNOWN_FAILURE(LevelZero{"unsupported DynSharedMemory"});

  LaunchArgs.NumGroups.x = 4;
  LaunchArgs.DynSharedMemory = 64 * sizeof(uint32_t);

  void *Mem;
  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_MANAGED,
                            LaunchArgs.GroupSize.x * LaunchArgs.NumGroups.x *
                                sizeof(uint32_t),
                            &Mem));

  void *ArgPtrs[] = {&Mem};
  size_t ArgSizes[] = {sizeof(Mem)};

  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));

  ASSERT_SUCCESS(olSyncQueue(Queue));

  uint32_t *Data = (uint32_t *)Mem;
  for (uint32_t i = 0; i < LaunchArgs.GroupSize.x * LaunchArgs.NumGroups.x; i++)
    ASSERT_EQ(Data[i], (i % 64) * 2);

  ASSERT_SUCCESS(olMemFree(Context, Mem));
}

TEST_P(olLaunchKernelLocalMemReductionTest, Success) {
  SKIP_KNOWN_FAILURE(LevelZero{"unsupported DynSharedMemory"});

  LaunchArgs.NumGroups.x = 4;
  LaunchArgs.DynSharedMemory = 64 * sizeof(uint32_t);

  void *Mem;
  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_MANAGED,
                            LaunchArgs.NumGroups.x * sizeof(uint32_t), &Mem));

  void *ArgPtrs[] = {&Mem};
  size_t ArgSizes[] = {sizeof(Mem)};

  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));

  ASSERT_SUCCESS(olSyncQueue(Queue));

  uint32_t *Data = (uint32_t *)Mem;
  for (uint32_t i = 0; i < LaunchArgs.NumGroups.x; i++)
    ASSERT_EQ(Data[i], 2 * LaunchArgs.GroupSize.x);

  ASSERT_SUCCESS(olMemFree(Context, Mem));
}

TEST_P(olLaunchKernelLocalMemStaticTest, Success) {
  LaunchArgs.NumGroups.x = 4;
  LaunchArgs.DynSharedMemory = 0;

  void *Mem;
  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_MANAGED,
                            LaunchArgs.NumGroups.x * sizeof(uint32_t), &Mem));

  void *ArgPtrs[] = {&Mem};
  size_t ArgSizes[] = {sizeof(Mem)};

  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));

  ASSERT_SUCCESS(olSyncQueue(Queue));

  uint32_t *Data = (uint32_t *)Mem;
  for (uint32_t i = 0; i < LaunchArgs.NumGroups.x; i++)
    ASSERT_EQ(Data[i], 2 * LaunchArgs.GroupSize.x);

  ASSERT_SUCCESS(olMemFree(Context, Mem));
}

// The test intends to verify the correctness of the current implementation of
// the event synchronisation.
TEST_P(olLaunchKernelSingleCounterSyncEventTest, SuccessSyncEvent) {
  void *InitValuePassed;
  void *ResNum;

  size_t Size = sizeof(uint32_t);

  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_DEVICE, Size,
                            &InitValuePassed));
  ASSERT_SUCCESS(
      olMemAlloc(Context, Device, OL_ALLOC_TYPE_DEVICE, Size, &ResNum));

  uint32_t HostInitVal = 0;
  ASSERT_SUCCESS(
      olMemcpy(Queue, InitValuePassed, Device, &HostInitVal, Host, Size));
  ASSERT_SUCCESS(olMemcpy(Queue, ResNum, Device, &HostInitVal, Host, Size));
  ASSERT_SUCCESS(olSyncQueue(Queue));

  // The execution time of the provided kernel should be high enough to ensure
  // that the read value of the final result is not correct without explicit
  // synchronization: explicit waiting for the event following the submitted
  // operation. LoopRange is arbitrarily set with the goal of prolonging the
  // kernel execution through a high number of loop iterations. In each
  // iteration, NumberToAdd is added to the final sum, which is stored in
  // ResNum.
  int32_t LoopRange = 1000000;
  int32_t NumberToAdd = 2;

  ArgsSingleCounter Args{LoopRange, NumberToAdd, (uint32_t *)InitValuePassed,
                         (uint32_t *)ResNum};

  void *ArgPtrs[] = {&Args.InitLoop, &Args.Addend, &Args.InitVal, &Args.Out};
  size_t ArgSizes[] = {sizeof(Args.InitLoop), sizeof(Args.Addend),
                       sizeof(Args.InitVal), sizeof(Args.Out)};

  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));

  uint32_t FinalResVal = 0;
  ASSERT_SUCCESS(olMemcpy(Queue, &FinalResVal, Host, ResNum, Device, Size));

  ol_event_handle_t Event = nullptr;
  ASSERT_SUCCESS(olCreateEvent(Queue, OL_EVENT_FLAGS_NONE, &Event));
  ASSERT_SUCCESS(olSyncEvent(Event));

  ASSERT_EQ(FinalResVal, NumberToAdd * LoopRange);

  ASSERT_SUCCESS(olMemFree(Context, InitValuePassed));
  ASSERT_SUCCESS(olMemFree(Context, ResNum));
}

// The test checks the correctness of the synchronization between queues using
// events. Enqueueing the kernel on the queue `Q1` should produce `Result1`,
// which is used as an initial value for the queue `Q2`. Therefore, before
// executing any future work, `Q2` should wait for the event `Event1`, which is
// created after submitting work on `Q1`. The required synchronization is
// ensured by `olWaitEvents(Q2, &Event1, 1)`. If `Q2` uses the value passed as
// `Result1` before the work on `Q1` has completed, the result of the kernel
// enqueued on `Q2` would be incorrect.
TEST_P(olLaunchKernelSingleCounterSyncEventTest, SuccessTwoQueues) {
  void *InitValuePassed;
  void *ResNum1;
  void *ResNum2;

  size_t Size = sizeof(uint32_t);

  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_DEVICE, Size,
                            &InitValuePassed));
  ASSERT_SUCCESS(
      olMemAlloc(Context, Device, OL_ALLOC_TYPE_DEVICE, Size, &ResNum1));
  ASSERT_SUCCESS(
      olMemAlloc(Context, Device, OL_ALLOC_TYPE_DEVICE, Size, &ResNum2));

  uint32_t HostInitVal = 0;
  ASSERT_SUCCESS(
      olMemcpy(Queue, InitValuePassed, Device, &HostInitVal, Host, Size));
  ASSERT_SUCCESS(olMemcpy(Queue, ResNum1, Device, &HostInitVal, Host, Size));
  ASSERT_SUCCESS(olMemcpy(Queue, ResNum2, Device, &HostInitVal, Host, Size));
  ASSERT_SUCCESS(olSyncQueue(Queue));

  ol_queue_handle_t Queue2 = nullptr;
  ASSERT_SUCCESS(olCreateQueue(Context, Device, &Queue2));

  // For the explanation of the reasoning behind particular values assigned to
  // parameters, see the comment in the Success test from the same test suite
  int32_t LoopRange = 1000000;
  int32_t NumberToAdd = 2;

  ArgsSingleCounter Args{LoopRange, NumberToAdd, (uint32_t *)InitValuePassed,
                         (uint32_t *)ResNum1};
  void *ArgPtrs[] = {&Args.InitLoop, &Args.Addend, &Args.InitVal, &Args.Out};
  size_t ArgSizes[] = {sizeof(Args.InitLoop), sizeof(Args.Addend),
                       sizeof(Args.InitVal), sizeof(Args.Out)};
  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));

  ol_event_handle_t Event = nullptr;
  ASSERT_SUCCESS(olCreateEvent(Queue, OL_EVENT_FLAGS_NONE, &Event));
  ASSERT_SUCCESS(olWaitEvents(Queue2, &Event, 1));
  ArgsSingleCounter Args2{LoopRange, NumberToAdd, (uint32_t *)ResNum1,
                          (uint32_t *)ResNum2};
  void *ArgPtrs2[] = {&Args2.InitLoop, &Args2.Addend, &Args2.InitVal,
                      &Args2.Out};

  // At the beginning of the kernel, ResNum from the first queue is saved
  // locally as the initial value. If operations enqueued before Event have not
  // been completed by the time the kernel is executed using the second queue,
  // the FinalResVal would be incorrect.
  ASSERT_SUCCESS(olLaunchKernel(Queue2, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs2), ArgPtrs2, ArgSizes));

  ASSERT_SUCCESS(olSyncQueue(Queue2));
  uint32_t FinalResVal = 0;
  ASSERT_SUCCESS(olMemcpy(Queue2, &FinalResVal, Host, ResNum2, Device, Size));
  ASSERT_SUCCESS(olSyncQueue(Queue2));

  ASSERT_EQ(FinalResVal, 2 * NumberToAdd * LoopRange);

  ASSERT_SUCCESS(olMemFree(Context, InitValuePassed));
  ASSERT_SUCCESS(olMemFree(Context, ResNum1));
  ASSERT_SUCCESS(olMemFree(Context, ResNum2));
}

TEST_P(olLaunchKernelGlobalTest, Success) {
  void *Mem;
  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_MANAGED,
                            LaunchArgs.GroupSize.x * sizeof(uint32_t), &Mem));

  void *ArgPtrs[] = {&Mem};
  size_t ArgSizes[] = {sizeof(Mem)};

  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernels[0], &LaunchArgs, nullptr,
                                0, nullptr, nullptr));
  ASSERT_SUCCESS(olSyncQueue(Queue));
  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernels[1], &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));
  ASSERT_SUCCESS(olSyncQueue(Queue));

  uint32_t *Data = (uint32_t *)Mem;
  for (uint32_t i = 0; i < 64; i++) {
    ASSERT_EQ(Data[i], i * 2);
  }

  ASSERT_SUCCESS(olMemFree(Context, Mem));
}

TEST_P(olLaunchKernelGlobalTest, InvalidNotAKernel) {
  ol_symbol_handle_t Global = nullptr;
  ASSERT_SUCCESS(
      olGetSymbol(Program, "global", OL_SYMBOL_KIND_GLOBAL_VARIABLE, &Global));
  ASSERT_ERROR(OL_ERRC_SYMBOL_KIND,
               olLaunchKernel(Queue, Device, Global, &LaunchArgs, nullptr, 0,
                              nullptr, nullptr));
}

TEST_P(olLaunchKernelGlobalCtorTest, Success) {
  SKIP_KNOWN_FAILURE(LevelZero{"unsupported feature"});

  void *Mem;
  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_MANAGED,
                            LaunchArgs.GroupSize.x * sizeof(uint32_t), &Mem));

  void *ArgPtrs[] = {&Mem};
  size_t ArgSizes[] = {sizeof(Mem)};

  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                std::size(ArgPtrs), ArgPtrs, ArgSizes));
  ASSERT_SUCCESS(olSyncQueue(Queue));

  uint32_t *Data = (uint32_t *)Mem;
  for (uint32_t i = 0; i < 64; i++) {
    ASSERT_EQ(Data[i], i + 100);
  }

  ASSERT_SUCCESS(olMemFree(Context, Mem));
}

TEST_P(olLaunchKernelGlobalDtorTest, Success) {
  // TODO: We can't inspect the result of a destructor yet, once we
  // find/implement a way, update this test. For now we just check that nothing
  // crashes
  ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr, 0,
                                nullptr, nullptr));
  ASSERT_SUCCESS(olSyncQueue(Queue));
}

TEST_P(olLaunchKernelGridSizeTest, Success) {
  void *Mem;
  ASSERT_SUCCESS(olMemAlloc(Context, Device, OL_ALLOC_TYPE_MANAGED,
                            6 * sizeof(uint32_t), &Mem));

  uint32_t *NumBlocks = static_cast<uint32_t *>(Mem);
  uint32_t *NumThreads = static_cast<uint32_t *>(Mem) + 3;

  void *ArgPtrs[] = {&NumBlocks, &NumThreads};
  size_t ArgSizes[] = {sizeof(NumBlocks), sizeof(NumThreads)};

  const uint32_t BaseBlocks[3] = {64, 16, 8};
  const uint32_t BaseThreads[3] = {32, 4, 2};

  for (uint32_t Dim = 1; Dim <= 3; ++Dim) {
    NumBlocks[0] = NumBlocks[1] = NumBlocks[2] = 0;
    NumThreads[0] = NumThreads[1] = NumThreads[2] = 0;

    LaunchArgs.Dimensions = 3;
    LaunchArgs.NumGroups = {BaseBlocks[0], Dim >= 2 ? BaseBlocks[1] : 1,
                            Dim >= 3 ? BaseBlocks[2] : 1};
    LaunchArgs.GroupSize = {BaseThreads[0], Dim >= 2 ? BaseThreads[1] : 1,
                            Dim >= 3 ? BaseThreads[2] : 1};

    ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
                                  std::size(ArgPtrs), ArgPtrs, ArgSizes));

    ASSERT_SUCCESS(olSyncQueue(Queue));

    ASSERT_EQ(NumBlocks[0], LaunchArgs.NumGroups.x);
    ASSERT_EQ(NumBlocks[1], LaunchArgs.NumGroups.y);
    ASSERT_EQ(NumBlocks[2], LaunchArgs.NumGroups.z);
    ASSERT_EQ(NumThreads[0], LaunchArgs.GroupSize.x);
    ASSERT_EQ(NumThreads[1], LaunchArgs.GroupSize.y);
    ASSERT_EQ(NumThreads[2], LaunchArgs.GroupSize.z);
  }

  ASSERT_SUCCESS(olMemFree(Context, Mem));
}