Merge branch 'free_function_kernel_scratch_memory' of https://github.com/lbushi25/llvm into free_function_kernel_scratch_memory

Author: lbushi25

llvm/lib/SYCLLowerIR/LowerWGLocalMemory.cpp

@@ -71,10 +71,10 @@ sycl::getKernelNamesUsingImplicitLocalMem(const Module &M) {
       return -1;
     };
     llvm::for_each(M.functions(), [&](const Function &F) {
-      if (F.getCallingConv() == CallingConv::SPIR_KERNEL) {
+      if (F.getCallingConv() == CallingConv::SPIR_KERNEL &&
+          F.hasFnAttribute(WORK_GROUP_STATIC_ATTR)) {
         int ArgPos = GetArgumentPos(F);
-        if (ArgPos >= 0 || F.hasFnAttribute(WORK_GROUP_STATIC_ATTR))
-          SPIRKernelNames.emplace_back(F.getName(), ArgPos);
+        SPIRKernelNames.emplace_back(F.getName(), ArgPos);
       }
     });
   }

llvm/test/SYCLLowerIR/work_group_static.ll

@@ -22,9 +22,29 @@ entry:
   ret void
 }
 
+; Function Attrs: convergent norecurse
+; CHECK: @__sycl_kernel_B{{.*}} #[[ATTRS:[0-9]+]]
+define weak_odr dso_local spir_kernel void @__sycl_kernel_B(ptr addrspace(1) %0) local_unnamed_addr #1 !kernel_arg_addr_space !5 {
+entry:
+  %1 = tail call spir_func ptr addrspace(3) @__sycl_dynamicLocalMemoryPlaceholder(i64 128) #1
+  ret void
+}
+
+; Function Attrs: convergent norecurse
+; CHECK: @__sycl_kernel_C{{.*}} #[[ATTRS]]
+define weak_odr dso_local spir_kernel void @__sycl_kernel_C(ptr addrspace(1) %0) local_unnamed_addr #1 !kernel_arg_addr_space !5 {
+entry:
+  %1 = tail call spir_func ptr addrspace(3) @__sycl_allocateLocalMemory(i64 128, i64 4) #1
+  ret void
+}
+
+; Function Attrs: convergent
+declare dso_local spir_func ptr addrspace(3) @__sycl_allocateLocalMemory(i64, i64) local_unnamed_addr #1
+
 ; Function Attrs: convergent
 declare dso_local spir_func ptr addrspace(3) @__sycl_dynamicLocalMemoryPlaceholder(i64) local_unnamed_addr #1
 
+; CHECK: #[[ATTRS]] = {{.*}} "sycl-work-group-static"
 attributes #0 = { convergent norecurse "disable-tail-calls"="false" "frame-pointer"="all" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "uniform-work-group-size"="true" "unsafe-fp-math"="false" "use-soft-float"="false" "sycl-work-group-static"="1" }
 attributes #1 = { convergent norecurse }
 

sycl/test-e2e/FreeFunctionKernels/free_function_kernel_local_memory.cpp

@@ -0,0 +1,112 @@
+// REQUIRES: aspect-usm_shared_allocations
+// UNSUPPORTED: target-amd
+// UNSUPPORTED-TRACKER: https://github.com/intel/llvm/issues/16072
+
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+
+// This test verifies that we can compile, run and get correct results when
+// using a free function kernel that allocates shared local memory in a kernel
+// either by way of the work group scratch memory extension or the work group
+// static memory extension.
+
+#include <sycl/ext/oneapi/work_group_static.hpp>
+
+#include "helpers.hpp"
+#include <cassert>
+#include <sycl/ext/oneapi/experimental/enqueue_functions.hpp>
+#include <sycl/ext/oneapi/free_function_queries.hpp>
+#include <sycl/group_barrier.hpp>
+#include <sycl/usm.hpp>
+
+namespace syclext = sycl::ext::oneapi;
+namespace syclexp = sycl::ext::oneapi::experimental;
+
+constexpr int SIZE = 16;
+
+SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<1>))
+void scratch_kernel(float *src, float *dst) {
+  size_t lid = syclext::this_work_item::get_nd_item<1>().get_local_linear_id();
+  float *local_mem = (float *)syclexp::get_work_group_scratch_memory();
+  local_mem[lid] = 2 * src[lid];
+  dst[lid] = local_mem[lid];
+}
+
+SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<1>))
+void static_kernel(float *src, float *dst) {
+  sycl::nd_item<1> item = syclext::this_work_item::get_nd_item<1>();
+  size_t lid = item.get_local_linear_id();
+  syclexp::work_group_static<float[SIZE]> local_mem;
+  local_mem[lid] = src[lid] * src[lid];
+  sycl::group_barrier(item.get_group());
+  if (item.get_group().leader()) { // Check that memory is indeed shared between
+                                   // the work group
+    for (int i = 0; i < SIZE; ++i)
+      assert(local_mem[i] == src[i] * src[i]);
+  }
+  dst[lid] = local_mem[lid];
+}
+
+SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<1>))
+void scratch_static_kernel(float *src, float *dst) {
+  size_t lid = syclext::this_work_item::get_nd_item<1>().get_local_linear_id();
+  float *scratch_mem = (float *)syclexp::get_work_group_scratch_memory();
+  syclexp::work_group_static<float[SIZE]> static_mem;
+  scratch_mem[lid] = src[lid];
+  static_mem[lid] = src[lid];
+  dst[lid] = scratch_mem[lid] + static_mem[lid];
+}
+
+int main() {
+  sycl::queue q;
+  float *src = sycl::malloc_shared<float>(SIZE, q);
+  float *dst = sycl::malloc_shared<float>(SIZE, q);
+
+  for (int i = 0; i < SIZE; i++) {
+    src[i] = i;
+  }
+
+  auto scratchbndl = syclexp::get_kernel_bundle<scratch_kernel,
+                                                sycl::bundle_state::executable>(
+      q.get_context());
+  auto staticbndl =
+      syclexp::get_kernel_bundle<static_kernel, sycl::bundle_state::executable>(
+          q.get_context());
+  auto scratchstaticbndl = syclexp::get_kernel_bundle<
+      scratch_static_kernel, sycl::bundle_state::executable>(q.get_context());
+
+  sycl::kernel ScratchKernel =
+      scratchbndl.template ext_oneapi_get_kernel<scratch_kernel>();
+  sycl::kernel StaticKernel =
+      staticbndl.template ext_oneapi_get_kernel<static_kernel>();
+  sycl::kernel ScratchStaticKernel =
+      scratchstaticbndl.template ext_oneapi_get_kernel<scratch_static_kernel>();
+  syclexp::launch_config ScratchKernelcfg{
+      ::sycl::nd_range<1>(::sycl::range<1>(SIZE), ::sycl::range<1>(SIZE)),
+      syclexp::properties{
+          syclexp::work_group_scratch_size(SIZE * sizeof(float))}};
+  syclexp::launch_config StaticKernelcfg{
+      ::sycl::nd_range<1>(::sycl::range<1>(SIZE), ::sycl::range<1>(SIZE))};
+
+  syclexp::nd_launch(q, ScratchKernelcfg, ScratchKernel, src, dst);
+  q.wait();
+  for (int i = 0; i < SIZE; i++) {
+    assert(dst[i] == 2 * src[i]);
+  }
+
+  syclexp::nd_launch(q, StaticKernelcfg, StaticKernel, src, dst);
+  q.wait();
+  for (int i = 0; i < SIZE; i++) {
+    assert(dst[i] == src[i] * src[i]);
+  }
+
+  syclexp::nd_launch(q, ScratchKernelcfg, ScratchStaticKernel, src, dst);
+  q.wait();
+  for (int i = 0; i < SIZE; i++) {
+    assert(dst[i] == 2 * src[i]);
+  }
+
+  sycl::free(src, q);
+  sycl::free(dst, q);
+  return 0;
+}