// REQUIRES: gpu && (level_zero || opencl)
// RUN: %{build} -o %t.out
// RUN: %{run} %t.out

#include <numeric>
#include <sycl/detail/core.hpp>
#include <sycl/ext/oneapi/experimental/prefetch.hpp>
#include <sycl/usm.hpp>

using namespace sycl;
namespace syclex = sycl::ext::oneapi::experimental;

constexpr size_t N = 128;
constexpr size_t NumWI = 4;
constexpr size_t arrSize = N / NumWI;

#define COMMA ,

#define TEST_PREFETCH_WO_COUNT(PREFETCH_ARG)                                   \
  {                                                                            \
    std::vector<int> res(N);                                                   \
    {                                                                          \
      buffer<int, 1> buf(res.data(), N);                                       \
      q.submit([&](handler &h) {                                               \
        auto acc = buf.get_access<access_mode::write>(h);                      \
        h.parallel_for(                                                        \
            nd_range<1>(range<1>(NumWI), range<1>(NumWI)),                     \
            [=](nd_item<1> idx) {                                              \
              syclex::joint_prefetch(                                          \
                  idx.get_group(), PREFETCH_ARG,                               \
                  syclex::properties{syclex::prefetch_hint_L1});               \
              for (int i = idx.get_local_linear_id() * arrSize;                \
                   i < idx.get_local_linear_id() * arrSize + arrSize; i++)     \
                acc[i] = dataChar[0] * dataChar[0];                            \
            });                                                                \
      });                                                                      \
      q.wait();                                                                \
    }                                                                          \
    for (int i = 0; i < N; i++)                                                \
      assert(res[i] == dataChar[0] * dataChar[0]);                             \
  }

#define TEST_PREFETCH_W_COUNT(PREFETCH_ARG)                                    \
  {                                                                            \
    std::vector<int> res(N);                                                   \
    {                                                                          \
      buffer<int, 1> buf(res.data(), N);                                       \
      q.submit([&](handler &h) {                                               \
        auto acc = buf.get_access<access_mode::write>(h);                      \
        h.parallel_for(                                                        \
            nd_range<1>(range<1>(NumWI), range<1>(NumWI)),                     \
            [=](nd_item<1> idx) {                                              \
              syclex::joint_prefetch(                                          \
                  idx.get_group(), PREFETCH_ARG, arrSize,                      \
                  syclex::properties{syclex::prefetch_hint_L1});               \
              for (int i = idx.get_local_linear_id() * arrSize;                \
                   i < idx.get_local_linear_id() * arrSize + arrSize; i++)     \
                acc[i] = dataChar[i] * dataChar[i];                            \
            });                                                                \
      });                                                                      \
      q.wait();                                                                \
    }                                                                          \
    for (int i = 0; i < N; i++)                                                \
      assert(res[i] == dataChar[i] * dataChar[i]);                             \
  }

void testPrefetchWithAcc(queue q, const std::vector<int> &data,
                         bool prefetchOneElem = true) {
  std::vector<int> res(N);
  {
    buffer<int, 1> bufRes(res.data(), N);
    buffer<int, 1> bufData(data.data(), N);
    q.submit([&](handler &h) {
      auto accRes = bufRes.get_access<access_mode::write>(h);
      auto accData = bufData.get_access<access_mode::read>(h);
      h.parallel_for(
          nd_range<1>(range<1>(NumWI), range<1>(NumWI)), [=](nd_item<1> idx) {
            if (prefetchOneElem)
              syclex::joint_prefetch(
                  idx.get_group(), accData, id(0),
                  syclex::properties{syclex::prefetch_hint_L1});
            else
              syclex::joint_prefetch(
                  idx.get_group(), accData,
                  id(idx.get_local_linear_id() * arrSize), arrSize,
                  syclex::properties{syclex::prefetch_hint_L1});
            for (int i = idx.get_local_linear_id() * arrSize;
                 i < idx.get_local_linear_id() * arrSize + arrSize; i++)
              accRes[i] = prefetchOneElem ? accData[0] * accData[0]
                                          : accData[i] * accData[i];
          });
    });
    q.wait();
  }
  for (int i = 0; i < N; i++)
    assert(res[i] == (prefetchOneElem ? data[0] * data[0] : data[i] * data[i]));
}

int main() {
  queue q;

  if (q.get_device().has(aspect::usm_shared_allocations)) {
    auto *dataChar = malloc_shared<char>(N, q);
    auto *dataVoid = reinterpret_cast<void *>(dataChar);
    auto mPtrChar = address_space_cast<access::address_space::global_space,
                                       access::decorated::yes>(dataChar);
    auto mPtrVoid = address_space_cast<access::address_space::global_space,
                                       access::decorated::yes>(dataVoid);

    std::iota(dataChar, dataChar + N, 0);

    // void prefetch(void* ptr, Properties properties = {});
    TEST_PREFETCH_WO_COUNT(dataVoid)
    // void prefetch(T* ptr, Properties properties = {});
    TEST_PREFETCH_WO_COUNT(dataChar)
    // void prefetch(void* ptr, size_t bytes, Properties properties = {});
    TEST_PREFETCH_W_COUNT(reinterpret_cast<void *>(
        &dataChar[idx.get_local_linear_id() * arrSize]))
    // void prefetch(T* ptr, size_t count, Properties properties = {});
    TEST_PREFETCH_W_COUNT(&dataChar[idx.get_local_linear_id() * arrSize])
    // void prefetch(multi_ptr<void, AddressSpace, IsDecorated> ptr,
    //               Properties properties = {});
    TEST_PREFETCH_WO_COUNT(mPtrVoid)
    // void prefetch(multi_ptr<T, AddressSpace, IsDecorated> ptr,
    //               Properties properties = {});
    TEST_PREFETCH_WO_COUNT(mPtrChar)
    // void prefetch(multi_ptr<void, AddressSpace, IsDecorated> ptr,
    //               size_t bytes, Properties properties = {});
    TEST_PREFETCH_W_COUNT(
        address_space_cast<access::address_space::global_space COMMA
                               access::decorated::yes>(reinterpret_cast<void *>(
            &dataChar[idx.get_local_linear_id() * arrSize])))
    // void prefetch(multi_ptr<T, AddressSpace, IsDecorated> ptr, size_t count,
    //               Properties properties = {});
    TEST_PREFETCH_W_COUNT(
        address_space_cast<
            access::address_space::global_space COMMA access::decorated::yes>(
            &dataChar[idx.get_local_linear_id() * arrSize]))

    free(dataChar, q);
  }
  {
    std::vector<int> data(N);
    std::iota(data.begin(), data.end(), 0);

    // void prefetch(accessor<DataT, Dimensions, AccessMode, target::device,
    //               IsPlaceholder> acc, id<Dimensions> offset,
    //               Properties properties = {});
    testPrefetchWithAcc(q, data);
    // void prefetch(accessor<DataT, Dimensions, AccessMode, target::device,
    //               IsPlaceholder> acc, id<Dimensions> offset, size_t count,
    //               Properties properties = {});
    testPrefetchWithAcc(q, data, false);
  }
}