[SYCL][ESIMD] Support 64-bit offsets for stateless accessor gather/scatter (#9462)

Signed-off-by: Sarnie, Nick <[email protected]>

Author: sarnex

sycl/include/sycl/ext/intel/esimd/memory.hpp

@@ -521,6 +521,7 @@ gather_impl(AccessorTy acc, simd<uint32_t, N> offsets, uint32_t glob_offset,
 ///   \c float.
 /// @tparam N The number of vector elements. Can be \c 1, \c 8, \c 16 or \c 32.
 /// @tparam AccessorTy The accessor type.
+/// @tparam Toffset The offset type.
 /// @param acc The accessor to gather from.
 /// @param offsets Per-element offsets in bytes.
 /// @param glob_offset Offset in bytes added to each individual element's offset
@@ -529,12 +530,17 @@ gather_impl(AccessorTy acc, simd<uint32_t, N> offsets, uint32_t glob_offset,
 ///   predicate are not accessed, their values in the resulting vector are
 ///   undefined.
 ///
-template <typename T, int N, typename AccessorTy>
-__ESIMD_API std::enable_if_t<(sizeof(T) <= 4) &&
-                                 (N == 1 || N == 8 || N == 16 || N == 32) &&
-                                 !std::is_pointer<AccessorTy>::value,
-                             simd<T, N>>
-gather(AccessorTy acc, simd<uint32_t, N> offsets, uint32_t glob_offset = 0,
+template <typename T, int N, typename AccessorTy, typename Toffset>
+__ESIMD_API std::enable_if_t<
+    (sizeof(T) <= 4) && (N == 1 || N == 8 || N == 16 || N == 32) &&
+        !std::is_pointer<AccessorTy>::value && std::is_integral_v<Toffset>,
+    simd<T, N>>
+gather(AccessorTy acc, simd<Toffset, N> offsets,
+#ifdef __ESIMD_FORCE_STATELESS_MEM
+       uint64_t glob_offset = 0,
+#else
+       uint32_t glob_offset = 0,
+#endif
        simd_mask<N> mask = 1) {
 #ifdef __ESIMD_FORCE_STATELESS_MEM
   return gather<T, N>(__ESIMD_DNS::accessorToPointer<T>(acc, glob_offset),
@@ -554,6 +560,7 @@ gather(AccessorTy acc, simd<uint32_t, N> offsets, uint32_t glob_offset = 0,
 ///   \c float.
 /// @tparam N The number of vector elements. Can be \c 1, \c 8, \c 16 or \c 32.
 /// @tparam AccessorTy The accessor type.
+/// @tparam Toffset The offset type.
 /// @param acc The accessor to scatter to.
 /// @param offsets Per-element offsets in bytes.
 /// @param vals Values to write.
@@ -563,12 +570,17 @@ gather(AccessorTy acc, simd<uint32_t, N> offsets, uint32_t glob_offset = 0,
 ///   predicate are not accessed.
 ///
 ///
-template <typename T, int N, typename AccessorTy>
-__ESIMD_API std::enable_if_t<(sizeof(T) <= 4) &&
-                             (N == 1 || N == 8 || N == 16 || N == 32) &&
-                             !std::is_pointer<AccessorTy>::value>
-scatter(AccessorTy acc, simd<uint32_t, N> offsets, simd<T, N> vals,
-        uint32_t glob_offset = 0, simd_mask<N> mask = 1) {
+template <typename T, int N, typename AccessorTy, typename Toffset>
+__ESIMD_API std::enable_if_t<
+    (sizeof(T) <= 4) && (N == 1 || N == 8 || N == 16 || N == 32) &&
+    !std::is_pointer<AccessorTy>::value && std::is_integral_v<Toffset>>
+scatter(AccessorTy acc, simd<Toffset, N> offsets, simd<T, N> vals,
+#ifdef __ESIMD_FORCE_STATELESS_MEM
+        uint64_t glob_offset = 0,
+#else
+        uint32_t glob_offset = 0,
+#endif
+        simd_mask<N> mask = 1) {
 #ifdef __ESIMD_FORCE_STATELESS_MEM
   scatter<T, N>(__ESIMD_DNS::accessorToPointer<T>(acc, glob_offset), offsets,
                 vals, mask);

sycl/test-e2e/ESIMD/accessor_gather_scatter_stateless_64.cpp

@@ -0,0 +1,78 @@
+//=-accessor_gather_scatter_stateless_64.cpp - DPC++ ESIMD on-device test-=//
+//
+// 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
+//
+//=-------------------------------------------------------------------=//
+// REQUIRES: gpu-intel-pvc
+// RUN: %{build} -o %t.out -fsycl-esimd-force-stateless-mem
+// RUN: %{run} %t.out
+
+#include <iostream>
+#include <sycl/ext/intel/esimd.hpp>
+#include <sycl/sycl.hpp>
+
+using namespace sycl;
+using namespace sycl::ext::intel::esimd;
+
+int main(void) {
+  constexpr unsigned VL = 16;
+  constexpr uint64_t Gig = 1024 * 1024 * 1024;
+
+  constexpr uint64_t Size = Gig / 2 + 16;
+  uint64_t *A = new uint64_t[Size];
+
+  for (uint64_t i = 0; i < Size; ++i)
+    A[i] = i;
+
+  buffer<uint64_t, 1> bufa(A, range<1>(Size));
+  queue q;
+
+  auto dev = q.get_device();
+  std::cout << "Running on " << dev.get_info<info::device::name>() << "\n";
+  try {
+    q.submit([&](handler &cgh) {
+       auto PA = bufa.get_access<access::mode::read_write>(cgh);
+       cgh.single_task<class Test>([=]() SYCL_ESIMD_KERNEL {
+         uint64_t offsetStart = (Size - VL) * sizeof(uint64_t);
+         simd<uint64_t, VL> offset(offsetStart, sizeof(uint64_t));
+         simd<uint64_t, VL> beginning(0, sizeof(uint64_t));
+         simd<uint32_t, VL> va = gather<uint32_t, VL>(PA, beginning);
+         simd<uint32_t, VL> vb = gather<uint32_t, VL>(PA, offset);
+         va *= 2;
+         vb *= 5;
+         scatter<uint32_t, VL>(PA, beginning, va);
+         scatter<uint32_t, VL>(PA, offset, vb);
+       });
+     }).wait();
+  } catch (sycl::exception const &e) {
+    std::cout << "SYCL exception caught: " << e.what() << '\n';
+    delete[] A;
+    return 1;
+  }
+
+  bool failed = false;
+  host_accessor A_acc(bufa);
+  for (uint64_t I = 0; I < VL; I++) {
+    uint64_t Expected = static_cast<uint32_t>(I) * 2;
+    if (A_acc[I] != Expected) {
+      std::cout << "FAILED: " << A_acc[I] << " != " << Expected << std::endl;
+      failed = true;
+    }
+  }
+  for (uint64_t I = Size - VL; I < Size; I++) {
+    uint64_t Expected = static_cast<uint32_t>(I) * 5;
+    if (A_acc[I] != Expected) {
+      std::cout << "FAILED: " << A_acc[I] << " != " << Expected << std::endl;
+      failed = true;
+    }
+  }
+
+  if (!failed)
+    std::cout << "PASSED" << std::endl;
+
+  delete[] A;
+
+  return failed;
+}