[SYCL] Disable dead arg elimination for free function kernels (#19776)

This PR disables dead argument elimination for free function kernels.
When using a free function kernel, the user sets the arguments manually
using `handler::set_arg` but if certain arguments are optimized away in
case they are not used, the implementation needs to be aware of this and
remove the relevant `set_arg` calls. Instead, its much more feasible to
just disable dead arg elimination in this case.

---------

Co-authored-by: Alexey Sachkov <[email protected]>

Author: lbushi25

llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp

@@ -575,6 +575,17 @@ void DeadArgumentEliminationPass::surveyFunction(const Function &F) {
     return;
   }
 
+  // Do not modify arguments when the SYCL kernel is a free function kernel.
+  // In this case, the user sets the arguments of the kernel by themselves
+  // and dead argument elimination may interfere with their expectations.
+  const bool FuncIsSyclFreeFunctionKernel =
+      F.hasFnAttribute("sycl-single-task-kernel") ||
+      F.hasFnAttribute("sycl-nd-range-kernel");
+  if (FuncIsSyclFreeFunctionKernel) {
+    markFrozen(F);
+    return;
+  }
+
   LLVM_DEBUG(
       dbgs() << "DeadArgumentEliminationPass - Inspecting callers for fn: "
              << F.getName() << "\n");

llvm/test/Transforms/DeadArgElim/sycl-kernels-neg.ll

@@ -24,7 +24,30 @@ define weak_odr void @NotASpirKernel(float %arg1, float %arg2) {
 
 define weak_odr void @ESIMDKernel(float %arg1, float %arg2) !sycl_explicit_simd !0 {
 ; CHECK-LABEL: define {{[^@]+}}@ESIMDKernel
-; CHECK-SAME: (float [[ARG1:%.*]], float [[ARG2:%.*]]) !sycl_explicit_simd !0 {
+; CHECK-SAME: (float [[ARG1:%.*]], float [[ARG2:%.*]]) {{.*}}{
+; CHECK-NEXT:    call void @foo(float [[ARG1]])
+; CHECK-NEXT:    ret void
+;
+  call void @foo(float %arg1)
+  ret void
+}
+
+; The following two tests ensure that dead arguments are not eliminated
+; from a free function kernel.
+
+define weak_odr spir_kernel void @FreeFuncKernelSingleTask(float %arg1, float %arg2) "sycl-single-task-kernel"="0" {
+; CHECK-LABEL: define {{[^@]+}}@FreeFuncKernelSingleTask
+; CHECK-SAME: (float [[ARG1:%.*]], float [[ARG2:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    call void @foo(float [[ARG1]])
+; CHECK-NEXT:    ret void
+;
+  call void @foo(float %arg1)
+  ret void
+}
+
+define weak_odr spir_kernel void @FreeFuncKernelNdRange(float %arg1, float %arg2) "sycl-nd-range-kernel"="0" {
+; CHECK-LABEL: define {{[^@]+}}@FreeFuncKernelNdRange
+; CHECK-SAME: (float [[ARG1:%.*]], float [[ARG2:%.*]]) #[[ATTR1:[0-9]+]] {
 ; CHECK-NEXT:    call void @foo(float [[ARG1]])
 ; CHECK-NEXT:    ret void
 ;

sycl/test-e2e/FreeFunctionKernels/template_specialization.cpp

@@ -173,7 +173,10 @@ int main() {
   test_func_custom_type<A::B::C::TestClass>();
   test_func<F<float>, float>();
   test_func<F<uint32_t>, uint32_t>();
-  test_func<variadic_templated<double>, double>();
+  // Variadic template functions do not work with free function kernels. See
+  // CMPLRLLVM-69528.
+  // TODO: Uncomment the following line once the tracker is resolved.
+  //  test_func<variadic_templated<double>, double>();
   test_func<sum1<3, float>, float>();
   test_accessor();
   test_shared();