[SYCL] Implement work group memory extension

clang/include/clang/Basic/Attr.td

@@ -1520,12 +1520,12 @@ def SYCLType: InheritableAttr {
   let Subjects = SubjectList<[CXXRecord, Enum], ErrorDiag>;
   let LangOpts = [SYCLIsDevice, SilentlyIgnoreSYCLIsHost];
   let Args = [EnumArgument<"Type", "SYCLType", /*is_string=*/true,
-                           ["accessor", "local_accessor",
+                           ["accessor", "local_accessor", "work_group_memory",
                             "specialization_id", "kernel_handler", "buffer_location",
                             "no_alias", "accessor_property_list", "group",
                             "private_memory", "aspect", "annotated_ptr", "annotated_arg",
                             "stream", "sampler", "host_pipe", "multi_ptr"],
-                           ["accessor", "local_accessor",
+                           ["accessor", "local_accessor", "work_group_memory",
                             "specialization_id", "kernel_handler", "buffer_location",
                             "no_alias", "accessor_property_list", "group",
                             "private_memory", "aspect", "annotated_ptr", "annotated_arg",

clang/include/clang/Sema/SemaSYCL.h

@@ -62,7 +62,8 @@ class SYCLIntegrationHeader {
     kind_pointer,
     kind_specialization_constants_buffer,
     kind_stream,
-    kind_last = kind_stream
+    kind_work_group_memory,
+    kind_last = kind_work_group_memory
   };
 
 public:

clang/lib/Sema/SemaSYCL.cpp

@@ -4522,6 +4522,9 @@ class SyclKernelIntHeaderCreator : public SyclKernelFieldHandler {
                           CurOffset + offsetOf(FD, FieldTy));
     } else if (SemaSYCL::isSyclType(FieldTy, SYCLTypeAttr::stream)) {
       addParam(FD, FieldTy, SYCLIntegrationHeader::kind_stream);
+    } else if (SemaSYCL::isSyclType(FieldTy, SYCLTypeAttr::work_group_memory)) {
+      addParam(FieldTy, SYCLIntegrationHeader::kind_work_group_memory,
+               offsetOf(FD, FieldTy));
     } else if (SemaSYCL::isSyclType(FieldTy, SYCLTypeAttr::sampler) ||
                SemaSYCL::isSyclType(FieldTy, SYCLTypeAttr::annotated_ptr) ||
                SemaSYCL::isSyclType(FieldTy, SYCLTypeAttr::annotated_arg)) {
@@ -5655,6 +5658,7 @@ static const char *paramKind2Str(KernelParamKind K) {
     CASE(stream);
     CASE(specialization_constants_buffer);
     CASE(pointer);
+    CASE(work_group_memory);
   }
   return "<ERROR>";
 

clang/test/SemaSYCL/Inputs/sycl/detail/kernel_desc.hpp

@@ -18,6 +18,7 @@ namespace detail {
     kind_pointer = 3,
     kind_specialization_constants_buffer = 4,
     kind_stream = 5,
+    kind_work_group_memory = 6,
     kind_invalid = 0xf, // not a valid kernel kind
   };
 

sycl-jit/common/include/Kernel.h

@@ -58,6 +58,7 @@ enum class ParameterKind : uint32_t {
   Pointer = 3,
   SpecConstBuffer = 4,
   Stream = 5,
+  WorkGroupMemory = 6,
   Invalid = 0xF,
 };
 

sycl/include/sycl/detail/kernel_desc.hpp

@@ -58,6 +58,7 @@ enum class kernel_param_kind_t {
   kind_pointer = 3,
   kind_specialization_constants_buffer = 4,
   kind_stream = 5,
+  kind_work_group_memory = 6,
   kind_invalid = 0xf, // not a valid kernel kind
 };
 

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

@@ -0,0 +1,88 @@
+//===-------------------- work_group_memory.hpp ---------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <type_traits>
+
+namespace sycl {
+inline namespace _V1 {
+namespace detail {
+template <typename T> struct is_unbounded_array : std::false_type {};
+
+template <typename T> struct is_unbounded_array<T[]> : std::true_type {};
+
+template <typename T>
+inline constexpr bool is_unbounded_array_v = is_unbounded_array<T>::value;
+
+class work_group_memory_impl {
+public:
+  work_group_memory_impl() : wgm_size{0}, buffer_size{0} {}
+  work_group_memory_impl(const work_group_memory_impl &rhs) = default;
+  work_group_memory_impl &
+  operator=(const work_group_memory_impl &rhs) = default;
+  work_group_memory_impl(size_t wgm_size, size_t buffer_size)
+      : wgm_size{wgm_size}, buffer_size{buffer_size} {}
+  size_t wgm_size;
+  size_t buffer_size;
+};
+
+inline size_t getWorkGroupMemoryOwnSize(detail::work_group_memory_impl *wgm) {
+  return wgm->wgm_size;
+}
+} // namespace detail
+
+namespace ext::oneapi::experimental {
+template <typename DataT, typename PropertyListT = empty_properties_t>
+class __SYCL_SPECIAL_CLASS __SYCL_TYPE(work_group_memory) work_group_memory
+    : sycl::detail::work_group_memory_impl {
+public:
+  using value_type = std::remove_all_extents_t<DataT>;
+
+private:
+  using decoratedPtr = typename sycl::detail::DecoratedType<
+      value_type, access::address_space::local_space>::type *;
+
+public:
+  work_group_memory() = default;
+  work_group_memory(const work_group_memory &rhs) = default;
+  work_group_memory &operator=(const work_group_memory &rhs) = default;
+  template <typename T = DataT,
+            typename = std::enable_if_t<!sycl::detail::is_unbounded_array_v<T>>>
+  work_group_memory(handler &)
+      : sycl::detail::work_group_memory_impl(sizeof(work_group_memory),
+                                             sizeof(DataT)) {}
+  template <typename T = DataT,
+            typename = std::enable_if_t<sycl::detail::is_unbounded_array_v<T>>>
+  work_group_memory(size_t num, handler &)
+      : sycl::detail::work_group_memory_impl(
+            sizeof(work_group_memory),
+            num * sizeof(std::remove_extent_t<DataT>)) {}
+  template <access::decorated IsDecorated = access::decorated::no>
+  multi_ptr<value_type, access::address_space::local_space, IsDecorated>
+  get_multi_ptr() const {
+    return sycl::address_space_cast<access::address_space::local_space,
+                                    IsDecorated, value_type>(ptr);
+  }
+  DataT *operator&() const { return reinterpret_cast<DataT *>(ptr); }
+  operator DataT &() const { return *reinterpret_cast<DataT *>(ptr); }
+  template <typename T = DataT,
+            typename = std::enable_if_t<!std::is_array_v<T>>>
+  const work_group_memory &operator=(const DataT &value) const {
+    *ptr = value;
+    return *this;
+  }
+#ifdef __SYCL_DEVICE_ONLY__
+  void __init(decoratedPtr ptr) { this->ptr = ptr; }
+#endif
+private:
+  decoratedPtr ptr;
+};
+} // namespace ext::oneapi::experimental
+} // namespace _V1
+} // namespace sycl

sycl/include/sycl/handler.hpp

@@ -161,6 +161,8 @@ class pipe;
 }
 
 namespace ext ::oneapi ::experimental {
+template <typename, typename> class work_group_memory;
+
 struct image_descriptor;
 } // namespace ext::oneapi::experimental
 
@@ -169,6 +171,8 @@ class graph_impl;
 } // namespace ext::oneapi::experimental::detail
 namespace detail {
 
+class work_group_memory_impl;
+size_t getWorkGroupMemoryOwnSize(work_group_memory_impl *);
 class handler_impl;
 class kernel_impl;
 class queue_impl;
@@ -535,8 +539,10 @@ class __SYCL_EXPORT handler {
   /// According to section 4.7.6.11. of the SYCL specification, a local accessor
   /// must not be used in a SYCL kernel function that is invoked via single_task
   /// or via the simple form of parallel_for that takes a range parameter.
+  /// The same restriction is in place for work group memory objects. Refer to
+  /// its spec for more details.
   template <typename KernelName, typename KernelType>
-  void throwOnLocalAccessorMisuse() const {
+  void throwOnKernelParameterMisuse() const {
     using NameT =
         typename detail::get_kernel_name_t<KernelName, KernelType>::name;
     for (unsigned I = 0; I < detail::getKernelNumParams<NameT>(); ++I) {
@@ -552,6 +558,13 @@ class __SYCL_EXPORT handler {
             "A local accessor must not be used in a SYCL kernel function "
             "that is invoked via single_task or via the simple form of "
             "parallel_for that takes a range parameter.");
+      if (Kind == detail::kernel_param_kind_t::kind_work_group_memory)
+        throw sycl::exception(
+            make_error_code(errc::kernel_argument),
+            "A work group memory object must not be used in a SYCL kernel "
+            "function that is invoked via single_task or via the simple form "
+            "of "
+            "parallel_for that takes a range parameter.");
     }
   }
 
@@ -701,6 +714,17 @@ class __SYCL_EXPORT handler {
 #endif
   }
 
+  template <typename DataT, typename PropertyListT>
+  void setArgHelper(
+      int ArgIndex,
+      const ext::oneapi::experimental::work_group_memory<DataT, PropertyListT>
+          &Arg) {
+    addArg(detail::kernel_param_kind_t::kind_work_group_memory, &Arg,
+           detail::getWorkGroupMemoryOwnSize(
+               static_cast<detail::work_group_memory_impl *>(&Arg)),
+           ArgIndex);
+  }
+
   // setArgHelper for non local accessor argument.
   template <typename DataT, int Dims, access::mode AccessMode,
             access::target AccessTarget, access::placeholder IsPlaceholder>
@@ -1316,7 +1340,7 @@ class __SYCL_EXPORT handler {
   void parallel_for_lambda_impl(range<Dims> UserRange, PropertiesT Props,
                                 KernelType KernelFunc) {
     throwIfActionIsCreated();
-    throwOnLocalAccessorMisuse<KernelName, KernelType>();
+    throwOnKernelParameterMisuse<KernelName, KernelType>();
     if (!range_size_fits_in_size_t(UserRange))
       throw sycl::exception(make_error_code(errc::runtime),
                             "The total number of work-items in "
@@ -1841,7 +1865,7 @@ class __SYCL_EXPORT handler {
                                _KERNELFUNCPARAM(KernelFunc)) {
     (void)Props;
     throwIfActionIsCreated();
-    throwOnLocalAccessorMisuse<KernelName, KernelType>();
+    throwOnKernelParameterMisuse<KernelName, KernelType>();
     // TODO: Properties may change the kernel function, so in order to avoid
     //       conflicts they should be included in the name.
     using NameT =
@@ -2046,6 +2070,13 @@ class __SYCL_EXPORT handler {
     setArgHelper(ArgIndex, std::move(Arg));
   }
 
+  template <typename DataT, typename PropertyListT =
+                                ext::oneapi::experimental::empty_properties_t>
+  void set_arg(int ArgIndex, const ext::oneapi::experimental::work_group_memory<
+                                 DataT, PropertyListT> &Arg) {
+    setArgHelper(ArgIndex, Arg);
+  }
+
   // set_arg for graph dynamic_parameters
   template <typename T>
   void set_arg(int argIndex,

sycl/include/sycl/sycl.hpp

@@ -100,6 +100,7 @@
 #include <sycl/ext/oneapi/experimental/raw_kernel_arg.hpp>
 #include <sycl/ext/oneapi/experimental/root_group.hpp>
 #include <sycl/ext/oneapi/experimental/tangle_group.hpp>
+#include <sycl/ext/oneapi/experimental/work_group_memory.hpp>
 #include <sycl/ext/oneapi/filter_selector.hpp>
 #include <sycl/ext/oneapi/free_function_queries.hpp>
 #include <sycl/ext/oneapi/functional.hpp>

sycl/source/detail/jit_compiler.cpp

@@ -132,6 +132,8 @@ translateArgType(kernel_param_kind_t Kind) {
     return PK::SpecConstBuffer;
   case kind::kind_stream:
     return PK::Stream;
+  case kind::kind_work_group_memory:
+    return PK::WorkGroupMemory;
   case kind::kind_invalid:
     return PK::Invalid;
   }

sycl/source/detail/scheduler/commands.cpp

@@ -2273,6 +2273,8 @@ void SetArgBasedOnType(
     const std::function<void *(Requirement *Req)> &getMemAllocationFunc,
     const sycl::context &Context, detail::ArgDesc &Arg, size_t NextTrueIndex) {
   switch (Arg.MType) {
+  case kernel_param_kind_t::kind_work_group_memory:
+    break;
   case kernel_param_kind_t::kind_stream:
     break;
   case kernel_param_kind_t::kind_accessor: {

sycl/source/handler.cpp

@@ -32,6 +32,7 @@
 #include <sycl/stream.hpp>
 
 #include <sycl/ext/oneapi/bindless_images_memory.hpp>
+#include <sycl/ext/oneapi/experimental/work_group_memory.hpp>
 #include <sycl/ext/oneapi/memcpy2d.hpp>
 
 namespace sycl {
@@ -789,6 +790,12 @@ void handler::processArg(void *Ptr, const detail::kernel_param_kind_t &Kind,
     }
     break;
   }
+  case kernel_param_kind_t::kind_work_group_memory: {
+    addArg(kernel_param_kind_t::kind_std_layout, nullptr,
+           static_cast<detail::work_group_memory_impl *>(Ptr)->buffer_size,
+           Index + IndexShift);
+    break;
+  }
   case kernel_param_kind_t::kind_sampler: {
     addArg(kernel_param_kind_t::kind_sampler, Ptr, sizeof(sampler),
            Index + IndexShift);

sycl/test-e2e/WorkGroupMemory/UR_argument_test.cpp

@@ -0,0 +1,20 @@
+// RUN: %{build} -o %t.out
+// RUN: env SYCL_UR_TRACE=2 %{run} %t.out | FileCheck %s
+
+#include <sycl/detail/core.hpp>
+#include <sycl/ext/oneapi/experimental/work_group_memory.hpp>
+
+// Check that the work group memory object is mapped to exactly one backend
+// kernel argument.
+
+int main() {
+  sycl::queue q;
+  q.submit([&](sycl::handler &cgh) {
+    sycl::ext::oneapi::experimental::work_group_memory<int[2]> data{cgh};
+    cgh.parallel_for(sycl::nd_range<1>{1, 1},
+                     [=](sycl::nd_item<1> it) { data[0] = 42; });
+  });
+}
+
+// CHECK-COUNT-1: ---> urKernelSetArg
+// CHECK-NOT: ---> urKernelSetArg

sycl/test-e2e/WorkGroupMemory/invalid_usage.cpp

@@ -0,0 +1,38 @@
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+#include <cassert>
+#include <sycl/sycl.hpp>
+using namespace sycl::ext::oneapi::experimental;
+
+// As per the spec, a work_group_memory object cannot be used in a single task
+// kernel or in a sycl::range kernel. An exception with error code
+// errc::kernel_argument must be thrown in such cases. This test verifies this.
+
+int main() {
+
+  sycl::queue q;
+  try {
+    q.submit([&](sycl::handler &cgh) {
+      work_group_memory<int> mem{cgh};
+      cgh.single_task([=]() { mem = 42; });
+    });
+    assert(false && "Work group memory was used in a single_task kernel and an "
+                    "exception was not seen"); // Fail, exception was not seen
+  } catch (sycl::exception &e) {
+    // Exception seen but must verify that the error code is correct
+    assert(e.code() == sycl::errc::kernel_argument);
+  }
+  // Same thing but with a range kernel
+  try {
+    q.submit([&](sycl::handler &cgh) {
+      work_group_memory<int> mem{cgh};
+      cgh.parallel_for(sycl::range{1}, [=](sycl::id<> it) { mem = 42; });
+    });
+    assert(false && "Work group memory was used in a range kernel and an "
+                    "exception was not seen"); // Fail, exception was not seen
+  } catch (sycl::exception &e) {
+    // Exception seen but must verify that the error code is correct
+    assert(e.code() == sycl::errc::kernel_argument);
+  }
+  return 0;
+}