[SYCL] Fixup SYCL metadata when dead argument is optimized and enable DAE for ESIMD (#10644)

sarnex · web-flow · commit d4807ba75c56 · 2023-08-04T13:41:30.000-05:00
This change enables DAE for ESIMD kernels and adds code to fixup
metadata inserted
by the SYCL FE when an argument is removed.

I found the metadata problem when running our E2E tests after enabling
DAE for ESIMD but it does not seem like an ESIMD-specific issue. This
change is tested by existing E2E tests.

The metadata fixup is required because downstream code in llvm-spirv
assumes
the number of operands for the metadata will match the actual number of
arguments
for the kernel.

Another option would be to fix the downstream code, but making the
metadata be accurate
to the actual kernel the downstream code will see seemed to be a better
root cause fix to me.

Let me know if you disagree with the above, or if you have a way to
prevent explicitly
listing all of the metadata to fix, which I don't love.

---------

Signed-off-by: Sarnie, Nick &lt;nick.sarnie@intel.com&gt;
diff --git a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
@@ -575,12 +575,10 @@ void DeadArgumentEliminationPass::surveyFunction(const Function &F) {
 
   // We can't modify arguments if the function is not local
   // but we can do so for SYCL kernel functions.
-  // DAE is not currently supported for ESIMD kernels.
-  bool FuncIsSyclNonEsimdKernel =
+  bool FuncIsSyclKernel =
       CheckSYCLKernels &&
-      (F.getCallingConv() == CallingConv::SPIR_KERNEL || IsNVPTXKernel(&F)) &&
-      !F.getMetadata("sycl_explicit_simd");
-  bool FuncIsLive = !F.hasLocalLinkage() && !FuncIsSyclNonEsimdKernel;
+      (F.getCallingConv() == CallingConv::SPIR_KERNEL || IsNVPTXKernel(&F));
+  bool FuncIsLive = !F.hasLocalLinkage() && !FuncIsSyclKernel;
   if (FuncIsLive && (!ShouldHackArguments || F.isIntrinsic())) {
     markLive(F);
     return;
@@ -820,6 +818,37 @@ bool DeadArgumentEliminationPass::removeDeadStuffFromFunction(Function *F) {
       MDOmitArgs.push_back(AliveArg ? MDOmitArgFalse : MDOmitArgTrue);
     F->setMetadata("sycl_kernel_omit_args",
                    llvm::MDNode::get(F->getContext(), MDOmitArgs));
+
+    // Update metadata inserted by the SYCL FE to match the new kernel
+    // signature.
+    auto FixupMetadata = [&](StringRef MDName) {
+      auto MDToFixup = F->getMetadata(MDName);
+      if (MDToFixup) {
+        assert(MDToFixup->getNumOperands() == MDOmitArgs.size() &&
+               "Unexpected metadata operands");
+        SmallVector<Metadata *, 10> NewMDOps;
+        for (unsigned int i = 0; i < MDToFixup->getNumOperands(); i++) {
+          const auto *MDConst = cast<ConstantAsMetadata>(MDOmitArgs[i]);
+          bool ArgWasRemoved =
+              static_cast<bool>(cast<ConstantInt>(MDConst->getValue())
+                                    ->getValue()
+                                    .getZExtValue());
+          if (!ArgWasRemoved)
+            NewMDOps.push_back(MDToFixup->getOperand(i));
+        }
+        F->setMetadata(MDName, llvm::MDNode::get(F->getContext(), NewMDOps));
+      }
+    };
+    FixupMetadata("kernel_arg_buffer_location");
+    FixupMetadata("kernel_arg_runtime_aligned");
+    FixupMetadata("kernel_arg_exclusive_ptr");
+    FixupMetadata("kernel_arg_addr_space");
+    FixupMetadata("kernel_arg_access_qual");
+    FixupMetadata("kernel_arg_type");
+    FixupMetadata("kernel_arg_base_type");
+    FixupMetadata("kernel_arg_type_qual");
+    FixupMetadata("kernel_arg_accessor_ptr");
+    FixupMetadata("kernel_arg_name");
   }
 
   // Find out the new return value.
diff --git a/sycl/test/check_device_code/buffer_location_codegen.cpp b/sycl/test/check_device_code/buffer_location_codegen.cpp
@@ -1,7 +1,7 @@
 // RUN: %clangxx -fsycl -c -fsycl-device-only -S -emit-llvm %s -o - | FileCheck %s
 
-// CHECK: define {{.*}}spir_kernel void @_ZTSZZ4mainENKUlRN4sycl3_V17handlerEE_clES2_E15kernel_function{{.*}} !kernel_arg_buffer_location ![[MDBL:[0-9]+]]
-// CHECK: ![[MDBL]] = !{i32 3, i32 -1, i32 -1, i32 -1, i32 -1, i32 -1, i32 2, i32 -1, i32 -1, i32 -1, i32 2, i32 -1, i32 -1, i32 -1, i32 -1}
+// CHECK: define {{.*}}spir_kernel void @_ZTSZZ4mainENKUlRN4sycl3_V17handlerEE_clES2_E15kernel_function(){{.*}} !kernel_arg_buffer_location ![[MDBL:[0-9]+]]
+// CHECK: ![[MDBL]] = !{}
 
 #include <sycl/sycl.hpp>
 
diff --git a/sycl/test/esimd/dae.cpp b/sycl/test/esimd/dae.cpp
@@ -0,0 +1,24 @@
+// RUN: %clangxx -fsycl-device-only -Xclang -fsycl-is-device -emit-llvm %s -S -o %t.ll -I %sycl_include
+// RUN: FileCheck %s --input-file %t.ll
+
+// Check SYCL FE metadata is updated when dead argument elimination removes an
+// argument
+
+#include <sycl/ext/intel/esimd.hpp>
+#include <sycl/sycl.hpp>
+using namespace sycl;
+
+template <typename name, typename Func>
+__attribute__((sycl_kernel)) void my_kernel(Func kernelFunc) {
+  kernelFunc();
+}
+
+SYCL_EXTERNAL SYCL_ESIMD_FUNCTION ESIMD_NOINLINE void callee(int x) {}
+
+// CHECK: define dso_local spir_kernel {{.*}} !kernel_arg_addr_space ![[#MD:]]
+// CHECK: !kernel_arg_access_qual ![[#MD]] !kernel_arg_type ![[#MD]] !kernel_arg_base_type ![[#MD]] !kernel_arg_type_qual ![[#MD]] !kernel_arg_accessor_ptr ![[#MD]]
+SYCL_EXTERNAL void __attribute__((noinline)) caller(int x) {
+  my_kernel<class kernel_abc>([=]() SYCL_ESIMD_KERNEL { callee(x); });
+}
+
+//CHECK: [[#MD]] = !{}
diff --git a/sycl/test/esimd/genx_func_attr.cpp b/sycl/test/esimd/genx_func_attr.cpp
@@ -24,7 +24,7 @@ SYCL_ESIMD_FUNCTION SYCL_EXTERNAL ESIMD_NOINLINE void callee(int x) {
 // inherits SLMSize and NBarrierCount from callee
 void caller_abc(int x) {
   kernel<class kernel_abc>([=]() SYCL_ESIMD_KERNEL { callee(x); });
-  // CHECK: define dso_local spir_kernel void @_ZTSZ10caller_abciE10kernel_abc(i32 noundef "VCArgumentIOKind"="0" %{{.*}}) local_unnamed_addr #2
+  // CHECK: define dso_local spir_kernel void @_ZTSZ10caller_abciE10kernel_abc() local_unnamed_addr #2
 }
 
 // inherits only NBarrierCount from callee
@@ -33,7 +33,7 @@ void caller_xyz(int x) {
     slm_init(1235); // also works in non-O0
     callee(x);
   });
-  // CHECK: define dso_local spir_kernel void @_ZTSZ10caller_xyziE10kernel_xyz(i32 noundef "VCArgumentIOKind"="0" %{{.*}}) local_unnamed_addr #2
+  // CHECK: define dso_local spir_kernel void @_ZTSZ10caller_xyziE10kernel_xyz() local_unnamed_addr #2
 }
 
 // CHECK: attributes #2 = { {{.*}} "VCNamedBarrierCount"="13" "VCSLMSize"="2469"
diff --git a/sycl/test/esimd/slm_init_specconst_size.cpp b/sycl/test/esimd/slm_init_specconst_size.cpp
@@ -1,9 +1,9 @@
 // RUN: %clangxx -O2 -fsycl -fsycl-device-only -Xclang -no-opaque-pointers -emit-llvm %s -o %t
 // RUN: sycl-post-link -split-esimd -lower-esimd -O2 -S %t -o %t.table
-// RUN: FileCheck --check-prefixes=CHECK,CHECK-TYPED %s -input-file=%t_esimd_0.ll
+// RUN: FileCheck %s -input-file=%t_esimd_0.ll
 // RUN: %clangxx -O2 -fsycl -fsycl-device-only -Xclang -opaque-pointers -emit-llvm %s -o %t
 // RUN: sycl-post-link -split-esimd -lower-esimd -O2 -S %t -o %t.table
-// RUN: FileCheck --check-prefixes=CHECK,CHECK-OPAQUE %s -input-file=%t_esimd_0.ll
+// RUN: FileCheck %s -input-file=%t_esimd_0.ll
 // Checks that we set 0 as VCSLMSize when slm_init is used with
 // non-constant operand, like with specialization constant.
 
@@ -24,8 +24,7 @@ int main() {
           [=](sycl::kernel_handler kh) SYCL_ESIMD_KERNEL {
             slm_init(kh.get_specialization_constant<Size>());
           });
-      // CHECK-TYPED: define weak_odr dso_local spir_kernel void @{{.*}}(i8 addrspace(1)* noundef align 1 "VCArgumentIOKind"="0" %{{.*}}) local_unnamed_addr #1
-      // CHECK-OPAQUE: define weak_odr dso_local spir_kernel void @{{.*}}(ptr addrspace(1) noundef align 1 "VCArgumentIOKind"="0" %{{.*}}) local_unnamed_addr #1
+      // CHECK: define weak_odr dso_local spir_kernel void @{{.*}}() local_unnamed_addr #1
     });
   }
 
