[SYCL] Implement work group memory extension (#15178)

Implement work group memory extension:
https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/proposed/sycl_ext_oneapi_work_group_memory.asciidoc
Two notes:

- Free function kernel support for work group memory argument will be
added in a future PR.
- When the assignment operator is called in host code, the assigned to
work group memory object does not actually correspond to the same
underlying memory as the one that was assigned from contradicting the
spec. See KhronosGroup/SYCL-Docs#552 for a
similar problem with `local_accessor`

Author: lbushi25

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.");
     }
   }