inline syclcompat::local_mem as well

Author: wenju-he

llvm/lib/SYCLLowerIR/LowerWGLocalMemory.cpp

@@ -9,6 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/SYCLLowerIR/LowerWGLocalMemory.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/Demangle/Demangle.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
@@ -91,29 +92,35 @@ ModulePass *llvm::createSYCLLowerWGLocalMemoryLegacyPass() {
 // inlined first before each __sycl_allocateLocalMemory call can be lowered to a
 // unique global variable. Inlining them here so that this pass doesn't have
 // implicit dependency on AlwaysInlinerPass.
+//
+// syclcompat::local_mem, which represents a unique allocation, calls
+// group_local_memory_for_overwrite. So local_mem should be inlined as well.
 static bool inlineGroupLocalMemoryFunc(Module &M) {
   Function *ALMFunc = M.getFunction(SYCL_ALLOCLOCALMEM_CALL);
   if (!ALMFunc || ALMFunc->use_empty())
     return false;
 
-  bool Changed = false;
-  for (auto *U : ALMFunc->users()) {
-    auto *Caller = cast<CallInst>(U)->getFunction();
-    if (!Caller->hasFnAttribute("sycl_forceinline")) {
-      // Already inlined.
-      continue;
-    }
-    for (auto *U2 : make_early_inc_range(Caller->users())) {
-      auto *CI = cast<CallInst>(U2);
-      InlineFunctionInfo IFI;
-      [[maybe_unused]] auto Result = InlineFunction(*CI, IFI);
-      assert(Result.isSuccess() && "inlining failed");
+  SmallVector<Function *, 4> WorkList{ALMFunc};
+  DenseSet<Function *> Visited;
+  while (!WorkList.empty()) {
+    auto *F = WorkList.pop_back_val();
+    for (auto *U : make_early_inc_range(F->users())) {
+      auto *CI = cast<CallInst>(U);
+      auto *Caller = CI->getFunction();
+      if (Caller->hasFnAttribute("sycl-forceinline") &&
+          Visited.insert(Caller).second)
+        WorkList.push_back(Caller);
+      if (F != ALMFunc) {
+        InlineFunctionInfo IFI;
+        [[maybe_unused]] auto Result = InlineFunction(*CI, IFI);
+        assert(Result.isSuccess() && "inlining failed");
+      }
     }
-    Caller->eraseFromParent();
-    Changed = true;
   }
+  for (auto *F : Visited)
+    F->eraseFromParent();
 
-  return Changed;
+  return !Visited.empty();
 }
 
 // TODO: It should be checked that __sycl_allocateLocalMemory (or its source

llvm/test/SYCLLowerIR/group_local_memory_inline.ll

@@ -63,4 +63,4 @@ entry:
 declare spir_func ptr addrspace(3) @__sycl_allocateLocalMemory(i64 noundef, i64 noundef)
 
 attributes #0 = { alwaysinline }
-attributes #1 = { "sycl_forceinline"="true" }
+attributes #1 = { "sycl-forceinline"="true" }

sycl/include/sycl/ext/oneapi/group_local_memory.hpp

@@ -22,7 +22,7 @@ inline namespace _V1 {
 namespace ext::oneapi {
 template <typename T, typename Group>
 #ifdef __SYCL_DEVICE_ONLY__
-[[__sycl_detail__::add_ir_attributes_function("sycl_forceinline", true)]]
+[[__sycl_detail__::add_ir_attributes_function("sycl-forceinline", true)]]
 #endif
 std::enable_if_t<
     std::is_trivially_destructible_v<T> && sycl::detail::is_group<Group>::value,
@@ -48,7 +48,7 @@ group_local_memory_for_overwrite(Group g) {
 
 template <typename T, typename Group, typename... Args>
 #ifdef __SYCL_DEVICE_ONLY__
-[[__sycl_detail__::add_ir_attributes_function("sycl_forceinline", true)]]
+[[__sycl_detail__::add_ir_attributes_function("sycl-forceinline", true)]]
 #endif
 std::enable_if_t<
     std::is_trivially_destructible_v<T> && sycl::detail::is_group<Group>::value,

sycl/include/syclcompat/memory.hpp

@@ -68,7 +68,11 @@
 
 namespace syclcompat {
 
-template <typename AllocT> auto *local_mem() {
+template <typename AllocT>
+#ifdef __SYCL_DEVICE_ONLY__
+[[__sycl_detail__::add_ir_attributes_function("sycl-forceinline", true)]]
+#endif
+auto *local_mem() {
   sycl::multi_ptr<AllocT, sycl::access::address_space::local_space>
       As_multi_ptr =
           sycl::ext::oneapi::group_local_memory_for_overwrite<AllocT>(

sycl/test/check_device_code/extensions/group_local_memory.cpp

@@ -5,6 +5,11 @@
 // The test checks that multiple calls to the same template instantiation of a
 // group local memory function result in separate allocations.
 
+// CHECK: @WGLocalMem{{.*}} = internal addrspace(3) global [4 x i8] undef, align 4
+// CHECK-NEXT: @WGLocalMem{{.*}} = internal addrspace(3) global [4 x i8] undef, align 4
+// CHECK-NEXT: @WGLocalMem{{.*}} = internal addrspace(3) global [4 x i8] undef, align 4
+// CHECK-NEXT: @WGLocalMem{{.*}} = internal addrspace(3) global [4 x i8] undef, align 4
+
 #include <sycl/detail/core.hpp>
 #include <sycl/ext/oneapi/group_local_memory.hpp>
 #include <sycl/usm.hpp>
@@ -31,8 +36,3 @@ int main() {
     });
   });
 }
-
-// CHECK: @WGLocalMem{{.*}} = internal addrspace(3) global [4 x i8] undef, align 4
-// CHECK-NEXT: @WGLocalMem{{.*}} = internal addrspace(3) global [4 x i8] undef, align 4
-// CHECK-NEXT: @WGLocalMem{{.*}} = internal addrspace(3) global [4 x i8] undef, align 4
-// CHECK-NEXT: @WGLocalMem{{.*}} = internal addrspace(3) global [4 x i8] undef, align 4

sycl/test/check_device_code/syclcompat_local_mem.cpp

@@ -0,0 +1,27 @@
+// RUN: %clangxx -fsycl -fsycl-device-only -S -emit-llvm %s -o - | FileCheck %s
+
+// The test checks that multiple calls to the same template instantiation of
+// syclcompat local_mem function result in separate allocations.
+
+// CHECK: @WGLocalMem{{.*}} = internal addrspace(3) global [4 x i8] undef, align 4
+// CHECK-NEXT: @WGLocalMem{{.*}} = internal addrspace(3) global [4 x i8] undef, align 4
+
+#include <sycl/detail/core.hpp>
+#include <syclcompat/memory.hpp>
+
+using namespace sycl;
+
+int main() {
+  queue Q;
+
+  int **Out = malloc_shared<int *>(4, Q);
+
+  Q.submit([&](handler &Cgh) {
+    Cgh.parallel_for(nd_range<1>({1}, {1}), [=](nd_item<1> Item) {
+      auto Ptr0 = syclcompat::local_mem<int[1]>();
+      auto Ptr1 = syclcompat::local_mem<int[1]>();
+      Out[0] = Ptr0;
+      Out[1] = Ptr1;
+    });
+  });
+}