[SYCL] Implement work group memory extension

clang/include/clang/Basic/Attr.td

@@ -1566,12 +1566,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

@@ -4693,6 +4693,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)) {
@@ -5773,6 +5776,7 @@ static const char *paramKind2Str(KernelParamKind K) {
     CASE(stream);
     CASE(specialization_constants_buffer);
     CASE(pointer);
+    CASE(work_group_memory);
   }
   return "<ERROR>";
 

clang/test/CodeGenSYCL/Inputs/sycl.hpp

@@ -649,6 +649,24 @@ const stream& operator<<(const stream &S, T&&) {
   return S;
 }
 
+// Dummy implementation of work_group_memory for use in CodeGenSYCL tests.
+template <typename DataT>
+class __attribute__((sycl_special_class))
+__SYCL_TYPE(work_group_memory) work_group_memory {
+public:
+  work_group_memory(handler &CGH) {}
+#ifdef __SYCL_DEVICE_ONLY__
+  // Default constructor for objects later initialized with __init member.
+  work_group_memory() = default;
+#endif
+
+  void __init(__attribute((opencl_local)) DataT *Ptr) { this->Ptr = Ptr; }
+  __attribute((opencl_local)) DataT *operator&() const { return Ptr; }
+
+private:
+  __attribute((opencl_local)) DataT *Ptr;
+};
+
 template <typename T, int dimensions = 1,
           typename AllocatorT = int /*fake type as AllocatorT is not used*/>
 class buffer {

clang/test/CodeGenSYCL/work_group_memory.cpp

@@ -0,0 +1,37 @@
+// RUN: %clang_cc1 -fsycl-is-device -triple spir64-unknown-unknown -disable-llvm-passes -emit-llvm %s -o %t.ll
+// RUN: FileCheck < %t.ll %s --check-prefix CHECK-IR
+// RUN: %clang_cc1 -fsycl-is-device -triple spir64-unknown-unknown -fsycl-int-header=%t.h %s
+// RUN: FileCheck < %t.h %s --check-prefix CHECK-INT-HEADER
+//
+// Tests for work_group_memory kernel parameter using the dummy implementation in Inputs/sycl.hpp.
+// The first two RUN commands verify that the init call is generated with the correct arguments in LLVM IR
+// and the second two RUN commands verify the contents of the integration header produced by the frontend.
+//
+// CHECK-IR: define dso_local spir_kernel void @
+// CHECK-IR-SAME: ptr addrspace(3) noundef align 4 [[PTR:%[a-zA-Z0-9_]+]]
+//
+// CHECK-IR: [[PTR]].addr = alloca ptr addrspace(3), align 8
+// CHECK-IR: [[PTR]].addr.ascast = addrspacecast ptr [[PTR]].addr to ptr addrspace(4)
+// CHECK-IR: store ptr addrspace(3) [[PTR]], ptr addrspace(4) [[PTR]].addr.ascast, align 8
+// CHECK-IR: [[PTR_LOAD:%[a-zA-Z0-9_]+]] = load ptr addrspace(3), ptr addrspace(4) [[PTR]].addr.ascast, align 8
+//
+// CHECK-IR: call spir_func void @{{.*}}__init{{.*}}(ptr addrspace(4) noundef align 8 dereferenceable_or_null(8) %{{[a-zA-Z0-9_]+}}, ptr addrspace(3) noundef [[PTR_LOAD]])
+//
+// CHECK-INT-HEADER: const kernel_param_desc_t kernel_signatures[] = {
+// CHECK-INT-HEADER-NEXT: //--- _ZTSZZ4mainENKUlRN4sycl3_V17handlerEE_clES2_EUlNS0_4itemILi1EEEE_
+// CHECK-INT-HEADER-NEXT: { kernel_param_kind_t::kind_work_group_memory, {{[4,8]}}, 0 },
+// CHECK-INT-HEADER-EMPTY:
+// CHECK-INT-HEADER-NEXT: { kernel_param_kind_t::kind_invalid, -987654321, -987654321 },
+// CHECK-INT-HEADER-NEXT: };
+
+#include "Inputs/sycl.hpp"
+
+int main() {
+  sycl::queue Q;
+  Q.submit([&](sycl::handler &CGH) {
+    sycl::work_group_memory<int> mem;
+    sycl::range<1> ndr;
+    CGH.parallel_for(ndr, [=](sycl::item<1> it) { int *ptr = &mem; });
+  });
+  return 0;
+}

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,84 @@
+//===-------------------- 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() : 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 buffer_size) : buffer_size{buffer_size} {}
+
+private:
+  size_t buffer_size;
+  friend class sycl::handler;
+};
+
+} // 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(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(
+            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

@@ -163,6 +163,8 @@ class pipe;
 }
 
 namespace ext ::oneapi ::experimental {
+template <typename, typename>
+class work_group_memory;
 struct image_descriptor;
 } // namespace ext::oneapi::experimental
 
@@ -171,6 +173,7 @@ class graph_impl;
 } // namespace ext::oneapi::experimental::detail
 namespace detail {
 
+class work_group_memory_impl;
 class handler_impl;
 class kernel_impl;
 class queue_impl;
@@ -564,8 +567,8 @@ class __SYCL_EXPORT handler {
   // The version for regular(standard layout) argument.
   template <typename T, typename... Ts>
   void setArgsHelper(int ArgIndex, T &&Arg, Ts &&...Args) {
-    set_arg(ArgIndex, std::move(Arg));
-    setArgsHelper(++ArgIndex, std::move(Args)...);
+    set_arg(ArgIndex, std::forward<T>(Arg));
+    setArgsHelper(++ArgIndex, std::forward<Ts>(Args)...);
   }
 
   void setArgsHelper(int) {}
@@ -603,6 +606,8 @@ class __SYCL_EXPORT handler {
 #endif
   }
 
+  void setArgHelper(int ArgIndex, detail::work_group_memory_impl &Arg);
+
   // setArgHelper for non local accessor argument.
   template <typename DataT, int Dims, access::mode AccessMode,
             access::target AccessTarget, access::placeholder IsPlaceholder>
@@ -1096,7 +1101,7 @@ class __SYCL_EXPORT handler {
                                 KernelType KernelFunc) {
 #ifndef __SYCL_DEVICE_ONLY__
     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 "
@@ -1641,7 +1646,7 @@ class __SYCL_EXPORT handler {
     kernel_single_task_wrapper<NameT, KernelType, PropertiesT>(KernelFunc);
 #ifndef __SYCL_DEVICE_ONLY__
     throwIfActionIsCreated();
-    throwOnLocalAccessorMisuse<KernelName, KernelType>();
+    throwOnKernelParameterMisuse<KernelName, KernelType>();
     verifyUsedKernelBundleInternal(
         detail::string_view{detail::getKernelName<NameT>()});
     // No need to check if range is out of INT_MAX limits as it's compile-time
@@ -1840,6 +1845,14 @@ 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,
+      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,
@@ -1858,9 +1871,8 @@ class __SYCL_EXPORT handler {
   ///
   /// \param Args are argument values to be set.
   template <typename... Ts> void set_args(Ts &&...Args) {
-    setArgsHelper(0, std::move(Args)...);
+    setArgsHelper(0, std::forward<Ts>(Args)...);
   }
-
   /// Defines and invokes a SYCL kernel function as a function object type.
   ///
   /// If it is a named function object and the function object type is
@@ -3233,7 +3245,6 @@ class __SYCL_EXPORT handler {
 private:
   std::shared_ptr<detail::handler_impl> impl;
   std::shared_ptr<detail::queue_impl> MQueue;
-
   std::vector<detail::LocalAccessorImplPtr> MLocalAccStorage;
   std::vector<std::shared_ptr<detail::stream_impl>> MStreamStorage;
   detail::string MKernelName;
@@ -3554,7 +3565,7 @@ class __SYCL_EXPORT handler {
   /// 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.
   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) {
@@ -3570,6 +3581,12 @@ 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.");
     }
   }
 

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/handler_impl.hpp

@@ -197,6 +197,9 @@ class handler_impl {
 
   /// True if MCodeLoc is sycl entry point code location
   bool MIsTopCodeLoc = true;
+
+  /// List of work group memory objects associated with this handler
+  std::vector<std::shared_ptr<detail::work_group_memory_impl>> MWorkGroupMemoryObjects;
 };
 
 } // namespace detail

sycl/source/detail/jit_compiler.cpp

@@ -133,6 +133,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

@@ -2297,6 +2297,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

@@ -34,6 +34,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 {
@@ -795,6 +796,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);
@@ -812,6 +819,13 @@ void handler::processArg(void *Ptr, const detail::kernel_param_kind_t &Kind,
   }
 }
 
+void handler::setArgHelper(int ArgIndex, detail::work_group_memory_impl &Arg) {
+  impl->MWorkGroupMemoryObjects.push_back(
+      std::make_shared<detail::work_group_memory_impl>(Arg));
+  addArg(detail::kernel_param_kind_t::kind_work_group_memory,
+         impl->MWorkGroupMemoryObjects.back().get(), 0, ArgIndex);
+}
+
 // The argument can take up more space to store additional information about
 // MAccessRange, MMemoryRange, and MOffset added with addArgsForGlobalAccessor.
 // We use the worst-case estimate because the lifetime of the vector is short.