//===-- unittests/Runtime/CUDA/Memory.cpp -----------------------*- C++ -*-===//
//
// 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 "flang/Runtime/CUDA/memory.h"
#include "cuda_runtime.h"
#include "../tools.h"
#include "gtest/gtest.h"
#include "flang-rt/runtime/allocator-registry.h"
#include "flang-rt/runtime/terminator.h"
#include "flang/Runtime/CUDA/allocator.h"
#include "flang/Runtime/CUDA/common.h"
#include "flang/Runtime/CUDA/descriptor.h"
#include "flang/Runtime/allocatable.h"
#include "flang/Support/Fortran.h"

using namespace Fortran::runtime;
using namespace Fortran::runtime::cuda;

TEST(MemoryCUFTest, SimpleAllocTramsferFree) {
  int *dev = (int *)RTNAME(CUFMemAlloc)(
      sizeof(int), kMemTypeDevice, __FILE__, __LINE__);
  EXPECT_TRUE(dev != 0);
  int host = 42;
  RTNAME(CUFDataTransferPtrPtr)
  ((void *)dev, (void *)&host, sizeof(int), kHostToDevice, __FILE__, __LINE__);
  host = 0;
  RTNAME(CUFDataTransferPtrPtr)
  ((void *)&host, (void *)dev, sizeof(int), kDeviceToHost, __FILE__, __LINE__);
  EXPECT_EQ(42, host);
  RTNAME(CUFMemFree)((void *)dev, kMemTypeDevice, __FILE__, __LINE__);
}

TEST(MemoryCUFTest, AllocZero) {
  int *dev = (int *)RTNAME(CUFMemAlloc)(0, kMemTypeDevice, __FILE__, __LINE__);
  EXPECT_TRUE(dev != 0);
  RTNAME(CUFMemFree)((void *)dev, kMemTypeDevice, __FILE__, __LINE__);
}

static OwningPtr<Descriptor> createAllocatable(
    Fortran::common::TypeCategory tc, int kind, int rank = 1) {
  return Descriptor::Create(TypeCode{tc, kind}, kind, nullptr, rank, nullptr,
      CFI_attribute_allocatable);
}

TEST(MemoryCUFTest, CUFDataTransferDescDesc) {
  using Fortran::common::TypeCategory;
  RTNAME(CUFRegisterAllocator)();
  // INTEGER(4), DEVICE, ALLOCATABLE :: a(:)
  auto dev{createAllocatable(TypeCategory::Integer, 4)};
  dev->SetAllocIdx(kDeviceAllocatorPos);
  EXPECT_EQ((int)kDeviceAllocatorPos, dev->GetAllocIdx());
  RTNAME(AllocatableSetBounds)(*dev, 0, 1, 10);
  RTNAME(AllocatableAllocate)
  (*dev, /*asyncObject=*/nullptr, /*hasStat=*/false, /*errMsg=*/nullptr,
      __FILE__, __LINE__);
  EXPECT_TRUE(dev->IsAllocated());

  // Create temp array to transfer to device.
  auto x{MakeArray<TypeCategory::Integer, 4>(std::vector<int>{10},
      std::vector<int32_t>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9})};
  RTNAME(CUFDataTransferDescDesc)
  (dev.get(), x.get(), kHostToDevice, __FILE__, __LINE__);

  // Retrieve data from device.
  auto host{MakeArray<TypeCategory::Integer, 4>(std::vector<int>{10},
      std::vector<int32_t>{0, 0, 0, 0, 0, 0, 0, 0, 0, 0})};
  RTNAME(CUFDataTransferDescDesc)
  (host.get(), dev.get(), kDeviceToHost, __FILE__, __LINE__);

  for (unsigned i = 0; i < 10; ++i) {
    EXPECT_EQ(*host->ZeroBasedIndexedElement<std::int32_t>(i), (std::int32_t)i);
  }
}

TEST(MemoryCUFTest, CUFDataTransferDescDescDstNotAllocated) {
  using Fortran::common::TypeCategory;
  RTNAME(CUFRegisterAllocator)();
  // INTEGER(4), DEVICE, ALLOCATABLE :: a(:)
  auto dev{createAllocatable(TypeCategory::Integer, 4)};
  dev->SetAllocIdx(kDeviceAllocatorPos);
  EXPECT_EQ((int)kDeviceAllocatorPos, dev->GetAllocIdx());
  EXPECT_FALSE(dev->IsAllocated());

  // Create temp array to transfer to device.
  auto x{MakeArray<TypeCategory::Integer, 4>(std::vector<int>{10},
      std::vector<int32_t>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9})};
  RTNAME(CUFDataTransferDescDesc)
  (dev.get(), x.get(), kHostToDevice, __FILE__, __LINE__);

  // Retrieve data from device.
  auto host{MakeArray<TypeCategory::Integer, 4>(std::vector<int>{10},
      std::vector<int32_t>{0, 0, 0, 0, 0, 0, 0, 0, 0, 0})};
  RTNAME(CUFDataTransferDescDesc)
  (host.get(), dev.get(), kDeviceToHost, __FILE__, __LINE__);

  for (unsigned i = 0; i < 10; ++i) {
    EXPECT_EQ(*host->ZeroBasedIndexedElement<std::int32_t>(i), (std::int32_t)i);
  }
}

TEST(MemoryCUFTest, CUFDataTransferDescDescStrided) {
  using Fortran::common::TypeCategory;
  static constexpr int elements{10};
  static constexpr int stride{2};
  SubscriptValue extent[]{elements};

  std::int32_t hostStorage[elements * stride]{};
  for (int i{0}; i < elements; ++i) {
    hostStorage[i * stride] = i;
    hostStorage[i * stride + 1] = -1;
  }

  std::int32_t *devStorage{static_cast<std::int32_t *>(RTNAME(CUFMemAlloc)(
      sizeof(hostStorage), kMemTypeDevice, __FILE__, __LINE__))};
  ASSERT_NE(devStorage, nullptr);
  cudaMemset(devStorage, 0xff, sizeof(hostStorage));

  StaticDescriptor<1> hostStaticDesc;
  Descriptor &hostDesc{hostStaticDesc.descriptor()};
  hostDesc.Establish(TypeCode{TypeCategory::Integer, 4}, sizeof(std::int32_t),
      hostStorage, 1, extent);
  hostDesc.GetDimension(0).SetByteStride(stride * sizeof(std::int32_t));

  StaticDescriptor<1> devStaticDesc;
  Descriptor &devDesc{devStaticDesc.descriptor()};
  devDesc.Establish(TypeCode{TypeCategory::Integer, 4}, sizeof(std::int32_t),
      devStorage, 1, extent);
  devDesc.GetDimension(0).SetByteStride(stride * sizeof(std::int32_t));

  RTNAME(CUFDataTransferDescDesc)
  (&devDesc, &hostDesc, kHostToDevice, __FILE__, __LINE__);

  std::int32_t result[elements * stride]{};
  RTNAME(CUFDataTransferPtrPtr)
  (result, devStorage, sizeof(result), kDeviceToHost, __FILE__, __LINE__);

  std::int32_t recvStorage[elements * stride]{};
  for (int i{0}; i < elements * stride; ++i) {
    recvStorage[i] = -2;
  }
  StaticDescriptor<1> recvStaticDesc;
  Descriptor &recvDesc{recvStaticDesc.descriptor()};
  recvDesc.Establish(TypeCode{TypeCategory::Integer, 4}, sizeof(std::int32_t),
      recvStorage, 1, extent);
  recvDesc.GetDimension(0).SetByteStride(stride * sizeof(std::int32_t));
  RTNAME(CUFDataTransferDescDesc)
  (&recvDesc, &devDesc, kDeviceToHost, __FILE__, __LINE__);

  RTNAME(CUFMemFree)(devStorage, kMemTypeDevice, __FILE__, __LINE__);

  for (int i{0}; i < elements; ++i) {
    EXPECT_EQ(result[i * stride], i);
    EXPECT_EQ(result[i * stride + 1], -1);
    EXPECT_EQ(recvStorage[i * stride], i);
    EXPECT_EQ(recvStorage[i * stride + 1], -2);
  }
}

TEST(MemoryCUFTest, CUFDataTransferDescDescLeadingSliceRank2) {
  using Fortran::common::TypeCategory;
  static constexpr int nx{8};
  static constexpr int ny{4};
  static constexpr int elements{nx * ny};
  SubscriptValue sliceExtent[]{nx - 2, ny};

  std::int32_t hostStorage[elements]{};
  for (int j{0}; j < ny; ++j) {
    for (int i{1}; i < nx - 1; ++i) {
      hostStorage[i + nx * j] = i + 10 * j;
    }
  }

  std::int32_t *devStorage{static_cast<std::int32_t *>(RTNAME(CUFMemAlloc)(
      sizeof(hostStorage), kMemTypeDevice, __FILE__, __LINE__))};
  ASSERT_NE(devStorage, nullptr);
  cudaMemset(devStorage, 0xff, sizeof(hostStorage));

  StaticDescriptor<2> hostStaticDesc;
  Descriptor &hostDesc{hostStaticDesc.descriptor()};
  hostDesc.Establish(TypeCode{TypeCategory::Integer, 4}, sizeof(std::int32_t),
      hostStorage + 1, 2, sliceExtent);
  hostDesc.GetDimension(0).SetByteStride(sizeof(std::int32_t));
  hostDesc.GetDimension(1).SetByteStride(nx * sizeof(std::int32_t));

  StaticDescriptor<2> devStaticDesc;
  Descriptor &devDesc{devStaticDesc.descriptor()};
  devDesc.Establish(TypeCode{TypeCategory::Integer, 4}, sizeof(std::int32_t),
      devStorage + 1, 2, sliceExtent);
  devDesc.GetDimension(0).SetByteStride(sizeof(std::int32_t));
  devDesc.GetDimension(1).SetByteStride(nx * sizeof(std::int32_t));

  RTNAME(CUFDataTransferDescDesc)
  (&devDesc, &hostDesc, kHostToDevice, __FILE__, __LINE__);

  std::int32_t result[elements]{};
  RTNAME(CUFDataTransferPtrPtr)
  (result, devStorage, sizeof(result), kDeviceToHost, __FILE__, __LINE__);

  std::int32_t recvStorage[elements]{};
  for (int i{0}; i < elements; ++i) {
    recvStorage[i] = -2;
  }
  StaticDescriptor<2> recvStaticDesc;
  Descriptor &recvDesc{recvStaticDesc.descriptor()};
  recvDesc.Establish(TypeCode{TypeCategory::Integer, 4}, sizeof(std::int32_t),
      recvStorage + 1, 2, sliceExtent);
  recvDesc.GetDimension(0).SetByteStride(sizeof(std::int32_t));
  recvDesc.GetDimension(1).SetByteStride(nx * sizeof(std::int32_t));
  RTNAME(CUFDataTransferDescDesc)
  (&recvDesc, &devDesc, kDeviceToHost, __FILE__, __LINE__);

  RTNAME(CUFMemFree)(devStorage, kMemTypeDevice, __FILE__, __LINE__);

  for (int j{0}; j < ny; ++j) {
    EXPECT_EQ(result[nx * j], -1);
    EXPECT_EQ(result[nx - 1 + nx * j], -1);
    EXPECT_EQ(recvStorage[nx * j], -2);
    EXPECT_EQ(recvStorage[nx - 1 + nx * j], -2);
    for (int i{1}; i < nx - 1; ++i) {
      const int index{i + nx * j};
      EXPECT_EQ(result[index], hostStorage[index]);
      EXPECT_EQ(recvStorage[index], hostStorage[index]);
    }
  }
}

TEST(MemoryCUFTest, CUFDataTransferDescDescLeadingSlice) {
  using Fortran::common::TypeCategory;
  static constexpr int nx{8};
  static constexpr int ny{4};
  static constexpr int nz{3};
  static constexpr int elements{nx * ny * nz};
  SubscriptValue sliceExtent[]{nx - 2, ny, nz};

  std::int32_t hostStorage[elements]{};
  for (int k{0}; k < nz; ++k) {
    for (int j{0}; j < ny; ++j) {
      for (int i{1}; i < nx - 1; ++i) {
        hostStorage[i + nx * (j + ny * k)] = i + 10 * j + 100 * k;
      }
    }
  }

  std::int32_t *devStorage{static_cast<std::int32_t *>(RTNAME(CUFMemAlloc)(
      sizeof(hostStorage), kMemTypeDevice, __FILE__, __LINE__))};
  ASSERT_NE(devStorage, nullptr);
  cudaMemset(devStorage, 0xff, sizeof(hostStorage));

  StaticDescriptor<3> hostStaticDesc;
  Descriptor &hostDesc{hostStaticDesc.descriptor()};
  hostDesc.Establish(TypeCode{TypeCategory::Integer, 4}, sizeof(std::int32_t),
      hostStorage + 1, 3, sliceExtent);
  hostDesc.GetDimension(0).SetByteStride(sizeof(std::int32_t));
  hostDesc.GetDimension(1).SetByteStride(nx * sizeof(std::int32_t));
  hostDesc.GetDimension(2).SetByteStride(nx * ny * sizeof(std::int32_t));

  StaticDescriptor<3> devStaticDesc;
  Descriptor &devDesc{devStaticDesc.descriptor()};
  devDesc.Establish(TypeCode{TypeCategory::Integer, 4}, sizeof(std::int32_t),
      devStorage + 1, 3, sliceExtent);
  devDesc.GetDimension(0).SetByteStride(sizeof(std::int32_t));
  devDesc.GetDimension(1).SetByteStride(nx * sizeof(std::int32_t));
  devDesc.GetDimension(2).SetByteStride(nx * ny * sizeof(std::int32_t));

  RTNAME(CUFDataTransferDescDesc)
  (&devDesc, &hostDesc, kHostToDevice, __FILE__, __LINE__);

  std::int32_t result[elements]{};
  RTNAME(CUFDataTransferPtrPtr)
  (result, devStorage, sizeof(result), kDeviceToHost, __FILE__, __LINE__);

  std::int32_t recvStorage[elements]{};
  for (int i{0}; i < elements; ++i) {
    recvStorage[i] = -2;
  }
  StaticDescriptor<3> recvStaticDesc;
  Descriptor &recvDesc{recvStaticDesc.descriptor()};
  recvDesc.Establish(TypeCode{TypeCategory::Integer, 4}, sizeof(std::int32_t),
      recvStorage + 1, 3, sliceExtent);
  recvDesc.GetDimension(0).SetByteStride(sizeof(std::int32_t));
  recvDesc.GetDimension(1).SetByteStride(nx * sizeof(std::int32_t));
  recvDesc.GetDimension(2).SetByteStride(nx * ny * sizeof(std::int32_t));
  RTNAME(CUFDataTransferDescDesc)
  (&recvDesc, &devDesc, kDeviceToHost, __FILE__, __LINE__);

  RTNAME(CUFMemFree)(devStorage, kMemTypeDevice, __FILE__, __LINE__);

  for (int k{0}; k < nz; ++k) {
    for (int j{0}; j < ny; ++j) {
      EXPECT_EQ(result[nx * (j + ny * k)], -1);
      EXPECT_EQ(result[nx - 1 + nx * (j + ny * k)], -1);
      EXPECT_EQ(recvStorage[nx * (j + ny * k)], -2);
      EXPECT_EQ(recvStorage[nx - 1 + nx * (j + ny * k)], -2);
      for (int i{1}; i < nx - 1; ++i) {
        const int index{i + nx * (j + ny * k)};
        EXPECT_EQ(result[index], hostStorage[index]);
        EXPECT_EQ(recvStorage[index], hostStorage[index]);
      }
    }
  }
}