[SYCL] Add support for work group memory free function kernel parameter

clang/lib/Sema/SemaSYCL.cpp

@@ -1522,7 +1522,7 @@ class KernelObjVisitor {
   void visitParam(ParmVarDecl *Param, QualType ParamTy,
                   HandlerTys &...Handlers) {
     if (isSyclSpecialType(ParamTy, SemaSYCLRef))
-      KP_FOR_EACH(handleOtherType, Param, ParamTy);
+      KP_FOR_EACH(handleSyclSpecialType, Param, ParamTy);
     else if (ParamTy->isStructureOrClassType()) {
       if (KP_FOR_EACH(handleStructType, Param, ParamTy)) {
         CXXRecordDecl *RD = ParamTy->getAsCXXRecordDecl();
@@ -2070,8 +2070,11 @@ class SyclKernelFieldChecker : public SyclKernelFieldHandler {
   }
 
   bool handleSyclSpecialType(ParmVarDecl *PD, QualType ParamTy) final {
-    Diag.Report(PD->getLocation(), diag::err_bad_kernel_param_type) << ParamTy;
-    IsInvalid = true;
+    if (!SemaSYCL::isSyclType(ParamTy, SYCLTypeAttr::work_group_memory)) {
+      Diag.Report(PD->getLocation(), diag::err_bad_kernel_param_type)
+          << ParamTy;
+      IsInvalid = true;
+    }
     return isValid();
   }
 
@@ -2223,8 +2226,8 @@ class SyclKernelUnionChecker : public SyclKernelFieldHandler {
   }
 
   bool handleSyclSpecialType(ParmVarDecl *PD, QualType ParamTy) final {
-    // TODO
-    unsupportedFreeFunctionParamType();
+    if (!SemaSYCL::isSyclType(ParamTy, SYCLTypeAttr::work_group_memory))
+      unsupportedFreeFunctionParamType(); // TODO
     return true;
   }
 
@@ -3008,9 +3011,26 @@ class SyclKernelDeclCreator : public SyclKernelFieldHandler {
     return handleSpecialType(FD, FieldTy);
   }
 
-  bool handleSyclSpecialType(ParmVarDecl *, QualType) final {
-    // TODO
-    unsupportedFreeFunctionParamType();
+  bool handleSyclSpecialType(ParmVarDecl *PD, QualType ParamTy) final {
+    if (SemaSYCL::isSyclType(ParamTy, SYCLTypeAttr::work_group_memory)) {
+      const auto *RecordDecl = ParamTy->getAsCXXRecordDecl();
+      assert(RecordDecl && "The type must be a RecordDecl");
+      CXXMethodDecl *InitMethod = getMethodByName(RecordDecl, InitMethodName);
+      assert(InitMethod && "The type must have the __init method");
+      // Don't do -1 here because we count on this to be the first parameter
+      // added (if any).
+      size_t ParamIndex = Params.size();
+      for (const ParmVarDecl *Param : InitMethod->parameters()) {
+        QualType ParamTy = Param->getType();
+        addParam(Param, ParamTy.getCanonicalType());
+        // Propagate add_ir_attributes_kernel_parameter attribute.
+        if (const auto *AddIRAttr =
+                Param->getAttr<SYCLAddIRAttributesKernelParameterAttr>())
+          Params.back()->addAttr(AddIRAttr->clone(SemaSYCLRef.getASTContext()));
+      }
+      LastParamIndex = ParamIndex;
+    } else // TODO
+      unsupportedFreeFunctionParamType();
     return true;
   }
 
@@ -3286,9 +3306,7 @@ class SyclKernelArgsSizeChecker : public SyclKernelFieldHandler {
   }
 
   bool handleSyclSpecialType(ParmVarDecl *PD, QualType ParamTy) final {
-    // TODO
-    unsupportedFreeFunctionParamType();
-    return true;
+    return handleSpecialType(ParamTy);
   }
 
   bool handleSyclSpecialType(const CXXRecordDecl *, const CXXBaseSpecifier &BS,
@@ -4416,6 +4434,45 @@ class FreeFunctionKernelBodyCreator : public SyclKernelFieldHandler {
                                 {});
   }
 
+  MemberExpr *buildMemberExpr(Expr *Base, ValueDecl *Member) {
+    DeclAccessPair MemberDAP = DeclAccessPair::make(Member, AS_none);
+    MemberExpr *Result = SemaSYCLRef.SemaRef.BuildMemberExpr(
+        Base, /*IsArrow */ false, FreeFunctionSrcLoc, NestedNameSpecifierLoc(),
+        FreeFunctionSrcLoc, Member, MemberDAP,
+        /*HadMultipleCandidates*/ false,
+        DeclarationNameInfo(Member->getDeclName(), FreeFunctionSrcLoc),
+        Member->getType(), VK_LValue, OK_Ordinary);
+    return Result;
+  }
+
+  void createSpecialMethodCall(const CXXRecordDecl *RD, StringRef MethodName,
+                               Expr *MemberBaseExpr,
+                               SmallVectorImpl<Stmt *> &AddTo) {
+    CXXMethodDecl *Method = getMethodByName(RD, MethodName);
+    if (!Method)
+      return;
+    unsigned NumParams = Method->getNumParams();
+    llvm::SmallVector<Expr *, 4> ParamDREs(NumParams);
+    llvm::ArrayRef<ParmVarDecl *> KernelParameters =
+        DeclCreator.getParamVarDeclsForCurrentField();
+    for (size_t I = 0; I < NumParams; ++I) {
+      QualType ParamType = KernelParameters[I]->getOriginalType();
+      ParamDREs[I] = SemaSYCLRef.SemaRef.BuildDeclRefExpr(
+          KernelParameters[I], ParamType, VK_LValue, FreeFunctionSrcLoc);
+    }
+    MemberExpr *MethodME = buildMemberExpr(MemberBaseExpr, Method);
+    QualType ResultTy = Method->getReturnType();
+    ExprValueKind VK = Expr::getValueKindForType(ResultTy);
+    ResultTy = ResultTy.getNonLValueExprType(SemaSYCLRef.getASTContext());
+    llvm::SmallVector<Expr *, 4> ParamStmts;
+    const auto *Proto = cast<FunctionProtoType>(Method->getType());
+    SemaSYCLRef.SemaRef.GatherArgumentsForCall(FreeFunctionSrcLoc, Method,
+                                               Proto, 0, ParamDREs, ParamStmts);
+    AddTo.push_back(CXXMemberCallExpr::Create(
+        SemaSYCLRef.getASTContext(), MethodME, ParamStmts, ResultTy, VK,
+        FreeFunctionSrcLoc, FPOptionsOverride()));
+  }
+
 public:
   static constexpr const bool VisitInsideSimpleContainers = false;
 
@@ -4435,9 +4492,37 @@ class FreeFunctionKernelBodyCreator : public SyclKernelFieldHandler {
     return true;
   }
 
-  bool handleSyclSpecialType(ParmVarDecl *, QualType) final {
-    // TODO
-    unsupportedFreeFunctionParamType();
+  bool handleSyclSpecialType(ParmVarDecl *PD, QualType ParamTy) final {
+    if (SemaSYCL::isSyclType(ParamTy, SYCLTypeAttr::work_group_memory)) {
+      const auto *RecordDecl = ParamTy->getAsCXXRecordDecl();
+      QualType Ty = PD->getOriginalType();
+      ASTContext &Ctx = SemaSYCLRef.SemaRef.getASTContext();
+      VarDecl *WorkGroupMemoryClone = VarDecl::Create(
+          Ctx, DeclCreator.getKernelDecl(), FreeFunctionSrcLoc,
+          FreeFunctionSrcLoc, PD->getIdentifier(), PD->getType(),
+          Ctx.getTrivialTypeSourceInfo(Ty), SC_None);
+      InitializedEntity VarEntity =
+          InitializedEntity::InitializeVariable(WorkGroupMemoryClone);
+      InitializationKind InitKind =
+          InitializationKind::CreateDefault(FreeFunctionSrcLoc);
+      InitializationSequence InitSeq(SemaSYCLRef.SemaRef, VarEntity, InitKind,
+                                     std::nullopt);
+      ExprResult Init = InitSeq.Perform(SemaSYCLRef.SemaRef, VarEntity,
+                                        InitKind, std::nullopt);
+      WorkGroupMemoryClone->setInit(
+          SemaSYCLRef.SemaRef.MaybeCreateExprWithCleanups(Init.get()));
+      WorkGroupMemoryClone->setInitStyle(VarDecl::CallInit);
+      Stmt *DS = new (SemaSYCLRef.getASTContext())
+          DeclStmt(DeclGroupRef(WorkGroupMemoryClone), FreeFunctionSrcLoc,
+                   FreeFunctionSrcLoc);
+      BodyStmts.push_back(DS);
+      Expr *MemberBaseExpr = SemaSYCLRef.SemaRef.BuildDeclRefExpr(
+          WorkGroupMemoryClone, Ty, VK_PRValue, FreeFunctionSrcLoc);
+      createSpecialMethodCall(RecordDecl, InitMethodName, MemberBaseExpr,
+                              BodyStmts);
+      ArgExprs.push_back(MemberBaseExpr);
+    } else // TODO
+      unsupportedFreeFunctionParamType();
     return true;
   }
 
@@ -4717,9 +4802,11 @@ class SyclKernelIntHeaderCreator : public SyclKernelFieldHandler {
     return true;
   }
 
-  bool handleSyclSpecialType(ParmVarDecl *, QualType) final {
-    // TODO
-    unsupportedFreeFunctionParamType();
+  bool handleSyclSpecialType(ParmVarDecl *PD, QualType ParamTy) final {
+    if (SemaSYCL::isSyclType(ParamTy, SYCLTypeAttr::work_group_memory))
+      addParam(PD, ParamTy, SYCLIntegrationHeader::kind_work_group_memory);
+    else
+      unsupportedFreeFunctionParamType(); // TODO
     return true;
   }
 
@@ -6196,7 +6283,16 @@ void SYCLIntegrationHeader::emit(raw_ostream &O) {
   O << "#include <sycl/detail/defines_elementary.hpp>\n";
   O << "#include <sycl/detail/kernel_desc.hpp>\n";
   O << "#include <sycl/ext/oneapi/experimental/free_function_traits.hpp>\n";
-
+  // When using work group memory parameters in free kernel functions, the
+  // integration header emits incorrect forward declarations for the work group
+  // memory type because it ignores default arguments. This means the user
+  // cannot use work group memory types with parameters omitted such as
+  // work_group_memory<int> where the hidden second parameter has a default
+  // value. To circumvent this, we include the correct forward declaration
+  // ourselves.
+  O << "#include <tuple>\n";
+  O << "#include "
+       "<sycl/ext/oneapi/experimental/work_group_memory_forward_decl.hpp>\n";
   O << "\n";
 
   LangOptions LO;

clang/test/CodeGenSYCL/free_function_int_header.cpp

@@ -2,7 +2,7 @@
 // RUN: FileCheck -input-file=%t.h %s
 // 
 // This test checks integration header contents for free functions with scalar,
-// pointer and non-decomposed struct parameters.
+// pointer, non-decomposed struct parameters and work group memory parameters.
 
 #include "mock_properties.hpp"
 #include "sycl.hpp"
@@ -96,6 +96,12 @@ void ff_7(KArgWithPtrArray<ArrSize> KArg) {
 
 template void ff_7(KArgWithPtrArray<TestArrSize> KArg);
 
+__attribute__((sycl_device))
+[[__sycl_detail__::add_ir_attributes_function("sycl-nd-range-kernel", 0)]] 
+void ff_8(sycl::work_group_memory<int>) {
+}
+
+
 // CHECK:      const char* const kernel_names[] = {
 // CHECK-NEXT:   {{.*}}__sycl_kernel_ff_2Piii
 // CHECK-NEXT:   {{.*}}__sycl_kernel_ff_2Piiii
@@ -105,6 +111,7 @@ template void ff_7(KArgWithPtrArray<TestArrSize> KArg);
 // CHECK-NEXT:   {{.*}}__sycl_kernel_ff_410NoPointers8Pointers3Agg
 // CHECK-NEXT:   {{.*}}__sycl_kernel_ff_6I3Agg7DerivedEvT_T0_i
 // CHECK-NEXT:   {{.*}}__sycl_kernel_ff_7ILi3EEv16KArgWithPtrArrayIXT_EE
+// CHECK-NEXT:   {{.*}}__sycl_kernel_ff_8N4sycl3_V117work_group_memoryIiEE
 // CHECK-NEXT:   ""
 // CHECK-NEXT: };
 
@@ -148,6 +155,9 @@ template void ff_7(KArgWithPtrArray<TestArrSize> KArg);
 // CHECK:  //--- _Z18__sycl_kernel_ff_7ILi3EEv16KArgWithPtrArrayIXT_EE
 // CHECK-NEXT:  { kernel_param_kind_t::kind_std_layout, 48, 0 },
 
+// CHECK:  //--- _Z18__sycl_kernel_ff_8N4sycl3_V117work_group_memoryIiEE
+// CHECK-NEXT:  { kernel_param_kind_t::kind_work_group_memory, 8, 0 },
+
 // CHECK:        { kernel_param_kind_t::kind_invalid, -987654321, -987654321 },
 // CHECK-NEXT: };
 
@@ -294,6 +304,26 @@ template void ff_7(KArgWithPtrArray<TestArrSize> KArg);
 // CHECK-NEXT: };
 // CHECK-NEXT: }
 
+// CHECK: Definition of _Z18__sycl_kernel_ff_8N4sycl3_V117work_group_memoryIiEE as a free function kernel
+
+// CHECK: Forward declarations of kernel and its argument types:
+// CHECK: template <typename DataT> class work_group_memory;
+
+// CHECK: void ff_8(sycl::work_group_memory<int>);
+// CHECK-NEXT: static constexpr auto __sycl_shim9() {
+// CHECK-NEXT: return (void (*)(class sycl::work_group_memory<int>))ff_8;
+// CHECK-NEXT: }
+// CHECK-NEXT: namespace sycl {
+// CHECK-NEXT: template <>
+// CHECK-NEXT: struct ext::oneapi::experimental::is_kernel<__sycl_shim9()> {
+// CHECK-NEXT: static constexpr bool value = true;
+// CHECK-NEXT: };
+// CHECK-NEXT: template <>
+// CHECK-NEXT: struct ext::oneapi::experimental::is_single_task_kernel<__sycl_shim9()> {
+// CHECK-NEXT: static constexpr bool value = true;
+// CHECK-NEXT: };
+// CHECK-NEXT: }
+
 // CHECK: #include <sycl/kernel_bundle.hpp>
 
 // CHECK: Definition of kernel_id of _Z18__sycl_kernel_ff_2Piii
@@ -359,3 +389,11 @@ template void ff_7(KArgWithPtrArray<TestArrSize> KArg);
 // CHECK-NEXT:   return sycl::detail::get_kernel_id_impl(std::string_view{"_Z18__sycl_kernel_ff_7ILi3EEv16KArgWithPtrArrayIXT_EE"});
 // CHECK-NEXT: }
 // CHECK-NEXT: }
+
+// CHECK: Definition of kernel_id of _Z18__sycl_kernel_ff_8N4sycl3_V117work_group_memoryIiEE
+// CHECK-NEXT: namespace sycl {
+// CHECK-NEXT: template <>
+// CHECK-NEXT: kernel_id ext::oneapi::experimental::get_kernel_id<__sycl_shim9()>() {
+// CHECK-NEXT: return sycl::detail::get_kernel_id_impl(std::string_view{"_Z18__sycl_kernel_ff_8N4sycl3_V117work_group_memoryIiEE"});
+// CHECK-NEXT: }
+// CHECK-NEXT: }

clang/test/CodeGenSYCL/free_function_kernel_params.cpp

@@ -1,7 +1,7 @@
 // RUN: %clang_cc1 -internal-isystem %S/Inputs -fsycl-is-device -triple spir64 \
 // RUN: -emit-llvm %s -o - | FileCheck %s
 // This test checks parameter IR generation for free functions with parameters
-// of non-decomposed struct type.
+// of non-decomposed struct type and work group memory type.
 
 #include "sycl.hpp"
 
@@ -56,3 +56,18 @@ template void ff_6(KArgWithPtrArray<TestArrSize> KArg);
 // CHECK: %struct.KArgWithPtrArray = type { [3 x ptr addrspace(4)], [3 x i32], [3 x i32] }
 // CHECK: define dso_local spir_kernel void @{{.*}}__sycl_kernel{{.*}}(ptr noundef byval(%struct.NoPointers) align 4 %__arg_S1, ptr noundef byval(%struct.__generated_Pointers) align 8 %__arg_S2, ptr noundef byval(%struct.__generated_Agg) align 8 %__arg_S3)
 // CHECK: define dso_local spir_kernel void @{{.*}}__sycl_kernel_ff_6{{.*}}(ptr noundef byval(%struct.__generated_KArgWithPtrArray) align 8 %__arg_KArg)
+
+__attribute__((sycl_device))
+[[__sycl_detail__::add_ir_attributes_function("sycl-nd-range-kernel", 0)]]
+void ff_7(sycl::work_group_memory<int> mem) {
+}
+
+// CHECK:  define dso_local spir_kernel void @{{.*}}__sycl_kernel_ff_7{{.*}}(ptr addrspace(3) noundef align 4 %__arg_Ptr)
+// CHECK:  %__arg_Ptr.addr = alloca ptr addrspace(3), align 8
+// CHECK-NEXT:  %mem = alloca %"class.sycl::_V1::work_group_memory", align 8
+// CHECK:  %__arg_Ptr.addr.ascast = addrspacecast ptr %__arg_Ptr.addr to ptr addrspace(4)
+// CHECK-NEXT:  %mem.ascast = addrspacecast ptr %mem to ptr addrspace(4)
+// CHECK:  store ptr addrspace(3) %__arg_Ptr, ptr addrspace(4) %__arg_Ptr.addr.ascast, align 8
+// CHECK-NEXT:  [[REGISTER:%[a-zA-Z0-9_]+]] = load ptr addrspace(3), ptr addrspace(4) %__arg_Ptr.addr.ascast, align 8
+// CHECK-NEXT:  call spir_func void @{{.*}}work_group_memory{{.*}}__init{{.*}}(ptr addrspace(4) noundef align 8 dereferenceable_or_null(8) %mem.ascast, ptr addrspace(3) noundef [[REGISTER]])
+

clang/test/CodeGenSYCL/work_group_memory.cpp

@@ -19,7 +19,7 @@
 //
 // 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-NEXT: { kernel_param_kind_t::kind_work_group_memory, 8, 0 },
 // CHECK-INT-HEADER-EMPTY:
 // CHECK-INT-HEADER-NEXT: { kernel_param_kind_t::kind_invalid, -987654321, -987654321 },
 // CHECK-INT-HEADER-NEXT: };

clang/test/SemaSYCL/Inputs/sycl.hpp

@@ -448,6 +448,23 @@ class __SYCL_TYPE(multi_ptr) multi_ptr<T, AS, access::decorated::legacy> {
   pointer_t m_Pointer;
 };
 
+// Dummy implementation of work_group_memory for use in SemaSYCL 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; }
+
+private:
+  __attribute((opencl_local)) DataT *Ptr;
+};
+
 namespace ext {
 namespace oneapi {
 namespace experimental {

clang/test/SemaSYCL/Inputs/sycl/ext/oneapi/experimental/work_group_memory_forward_decl.hpp

@@ -0,0 +1,16 @@
+#pragma once
+// Dummy header file for the purpose of SemaSYCL testing.
+// It shadows the file 
+// sycl/include/sycl/ext/oneapi/experimental/work_group_memory_forward_decl.hpp
+namespace sycl {
+inline namespace _V1 {
+namespace ext {
+namespace oneapi {
+namespace experimental {
+template <typename DataT, typename PropertiesT = int>
+class work_group_memory;
+} // namespace experimental
+} // namespace oneapi
+} // namespace ext
+} // namespace _V1
+} // namespace sycl

clang/test/SemaSYCL/free_function_kernel_params.cpp

@@ -1,7 +1,7 @@
 // RUN: %clang_cc1 -internal-isystem %S/Inputs -fsycl-is-device -ast-dump \
 // RUN: %s -o - | FileCheck %s
 // This test checks parameter rewriting for free functions with parameters
-// of type scalar, pointer and non-decomposed struct.
+// of type scalar, pointer, non-decomposed struct and work group memory.
 
 #include "sycl.hpp"
 
@@ -171,3 +171,23 @@ template void ff_6(Agg S1, Derived1 S2, int);
 // CHECK-NEXT: DeclRefExpr {{.*}} '__generated_Derived1' lvalue ParmVar {{.*}} '__arg_S2' '__generated_Derived1'
 // CHECK-NEXT: ImplicitCastExpr {{.*}} 'int' <LValueToRValue>
 // CHECK-NEXT: DeclRefExpr {{.*}} 'int' lvalue ParmVar {{.*}} '__arg_end' 'int'
+
+__attribute__((sycl_device))
+[[__sycl_detail__::add_ir_attributes_function("sycl-nd-range-kernel", 0)]]
+void ff_7(sycl::work_group_memory<int> mem) {
+}
+// CHECK: FunctionDecl {{.*}}__sycl_kernel{{.*}}'void (__local int *)'
+// CHECK-NEXT: ParmVarDecl {{.*}} used __arg_Ptr '__local int *'
+// CHECK-NEXT: CompoundStmt
+// CHECK-NEXT: DeclStmt
+// CHECK-NEXT: VarDecl {{.*}} used mem 'sycl::work_group_memory<int>' callinit
+// CHECK-NEXT: CXXConstructExpr {{.*}} 'sycl::work_group_memory<int>' 'void () noexcept'
+// CHECK-NEXT: CXXMemberCallExpr {{.*}} 'void'
+// CHECK-NEXT: MemberExpr {{.*}} 'void (__local int *)' lvalue .__init
+// CHECK-NEXT: DeclRefExpr {{.*}} 'sycl::work_group_memory<int>' Var {{.*}} 'mem' 'sycl::work_group_memory<int>'
+// CHECK-NEXT: ImplicitCastExpr {{.*}} '__local int *' <LValueToRValue>
+// CHECK-NEXT: DeclRefExpr {{.*}} '__local int *' lvalue ParmVar {{.*}} '__arg_Ptr' '__local int *'
+// CHECK-NEXT: CallExpr {{.*}} 'void'
+// CHECK-NEXT: ImplicitCastExpr {{.*}} 'void (*)(sycl::work_group_memory<int>)' <FunctionToPointerDecay>
+// CHECK-NEXT: DeclRefExpr {{.*}} 'void (sycl::work_group_memory<int>)' lvalue Function {{.*}} 'ff_7' 'void (sycl::work_group_memory<int>)'
+// CHECK-NEXT: DeclRefExpr {{.*}} 'sycl::work_group_memory<int>' Var {{.*}} 'mem' 'sycl::work_group_memory<int>'