[SYCL] Fix bug in enqueue free functions implementation (#21461)

There is a bug in some of the enqueue free functions that causes them to
work only with 1-dimensional kernels. This PR fixes this bug and adds a
couple of test cases to test it.

Author: lbushi25

sycl/include/sycl/ext/oneapi/experimental/enqueue_functions.hpp

@@ -416,14 +416,16 @@ void nd_launch(queue Q, nd_range<Dimensions> Range,
                [[maybe_unused]] kernel_function_s<Func> KernelFunc,
                ArgsT &&...Args) {
   detail::submit_kernel_direct_parallel_for(
-      std::move(Q), Range, [Args...](sycl::nd_item<>) { Func(Args...); });
+      std::move(Q), Range,
+      [Args...](sycl::nd_item<Dimensions>) { Func(Args...); });
 }
 
 template <auto *Func, int Dimensions, typename... ArgsT>
 void nd_launch(handler &CGH, nd_range<Dimensions> Range,
                [[maybe_unused]] kernel_function_s<Func> KernelFunc,
                ArgsT &&...Args) {
-  CGH.parallel_for(Range, [Args...](sycl::nd_item<>) { Func(Args...); });
+  CGH.parallel_for(Range,
+                   [Args...](sycl::nd_item<Dimensions>) { Func(Args...); });
 }
 
 template <auto *Func, int Dimensions, typename Properties, typename... ArgsT>
@@ -436,7 +438,7 @@ void nd_launch(queue Q, launch_config<nd_range<Dimensions>, Properties> Config,
       ConfigAccess(Config);
   detail::submit_kernel_direct_parallel_for(
       std::move(Q), ConfigAccess.getRange(),
-      [Args...](sycl::nd_item<>) { Func(Args...); }, {},
+      [Args...](sycl::nd_item<Dimensions>) { Func(Args...); }, {},
       ConfigAccess.getProperties());
 }
 
@@ -449,7 +451,7 @@ void nd_launch(handler &CGH,
                                                         Properties>
       ConfigAccess(Config);
   CGH.parallel_for(ConfigAccess.getRange(), ConfigAccess.getProperties(),
-                   [Args...](sycl::nd_item<>) { Func(Args...); });
+                   [Args...](sycl::nd_item<Dimensions>) { Func(Args...); });
 }
 
 inline void memcpy(handler &CGH, void *Dest, const void *Src, size_t NumBytes) {

sycl/test-e2e/FreeFunctionKernels/free_function_kernels_enqueue.cpp

@@ -45,6 +45,18 @@ void square(int *src, int *dst) {
   dst[Lid] = src[Lid] * src[Lid];
 }
 
+SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<2>))
+void square2D(int *src, int *dst) {
+  size_t Gid = syclext::this_work_item::get_nd_item<2>().get_global_linear_id();
+  dst[Gid] = src[Gid] * src[Gid];
+}
+
+SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<3>))
+void square3D(int *src, int *dst) {
+  size_t Gid = syclext::this_work_item::get_nd_item<3>().get_global_linear_id();
+  dst[Gid] = src[Gid] * src[Gid];
+}
+
 template <typename T>
 SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<1>))
 void squareWithScratchMemoryTemplated(T *src, T *dst) {
@@ -60,7 +72,7 @@ void squareWithAccessor(accType src, accType dst) {
   dst[Lid] = src[Lid] * src[Lid];
 }
 
-constexpr int SIZE = 16;
+constexpr int SIZE = 8;
 
 int main() {
   sycl::queue Q;
@@ -175,6 +187,34 @@ int main() {
     assert(Dst[I] == Src[I] * Src[I]);
   }
 
+  int *Src2D = sycl::malloc_shared<int>(SIZE * SIZE, Q);
+  int *Dst2D = sycl::malloc_shared<int>(SIZE * SIZE, Q);
+
+  Q.submit([&](sycl::handler &CGH) {
+     syclexp::nd_launch(CGH,
+                        ::sycl::nd_range<2>(::sycl::range<2>(SIZE, SIZE),
+                                            ::sycl::range<2>(SIZE, SIZE)),
+                        syclexp::kernel_function<square2D>, Src2D, Dst2D);
+   }).wait();
+
+  for (int I = 0; I < SIZE * SIZE; I++) {
+    assert(Dst2D[I] == Src2D[I] * Src2D[I]);
+  }
+
+  int *Src3D = sycl::malloc_shared<int>(SIZE * SIZE * SIZE, Q);
+  int *Dst3D = sycl::malloc_shared<int>(SIZE * SIZE * SIZE, Q);
+
+  Q.submit([&](sycl::handler &CGH) {
+     syclexp::nd_launch(CGH,
+                        ::sycl::nd_range<3>(::sycl::range<3>(SIZE, SIZE, SIZE),
+                                            ::sycl::range<3>(SIZE, SIZE, SIZE)),
+                        syclexp::kernel_function<square3D>, Src3D, Dst3D);
+   }).wait();
+
+  for (int I = 0; I < SIZE * SIZE * SIZE; I++) {
+    assert(Dst3D[I] == Src3D[I] * Src3D[I]);
+  }
+
   Q.submit([&](sycl::handler &CGH) {
      static_assert(std::is_same_v<decltype(syclexp::single_task(
                                       CGH, syclexp::kernel_function<successor>,