[SYCL] Add support for eliminated arg masks in SYCLBIN kernel bundles (#19163)

This commit fixes an issue where kernel bundles created from SYCLBIN
would not correctly filter based on the DAE-produced kernel arg masks.

---------

Signed-off-by: Larsen, Steffen <[email protected]>

Author: steffenlarsen

sycl/source/detail/device_image_impl.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include <detail/device_image_impl.hpp>
+#include <detail/kernel_arg_mask.hpp>
 #include <detail/kernel_bundle_impl.hpp>
 
 namespace sycl {
@@ -47,10 +48,14 @@ std::shared_ptr<kernel_impl> device_image_impl::tryGetExtensionKernel(
                                            &UrKernel);
   // Kernel created by urKernelCreate is implicitly retained.
 
+  const KernelArgMask *ArgMask = nullptr;
+  if (auto ArgMaskIt = MEliminatedKernelArgMasks.find(AdjustedName);
+      ArgMaskIt != MEliminatedKernelArgMasks.end())
+    ArgMask = &ArgMaskIt->second;
+
   return std::make_shared<kernel_impl>(
       UrKernel, *detail::getSyclObjImpl(Context), shared_from_this(),
-      OwnerBundle,
-      /*ArgMask=*/nullptr, UrProgram, /*CacheMutex=*/nullptr);
+      OwnerBundle, ArgMask, UrProgram, /*CacheMutex=*/nullptr);
 }
 
 } // namespace detail

sycl/source/detail/device_image_impl.hpp

@@ -147,6 +147,7 @@ class ManagedDeviceBinaries {
 
 using MangledKernelNameMapT = std::map<std::string, std::string, std::less<>>;
 using KernelNameSetT = std::set<std::string, std::less<>>;
+using KernelNameToArgMaskMap = std::unordered_map<std::string, KernelArgMask>;
 
 // Information unique to images compiled at runtime through the
 // ext_oneapi_kernel_compiler extension.
@@ -255,12 +256,23 @@ class device_image_impl
         MKernelIDs(std::move(KernelIDs)),
         MSpecConstsDefValBlob(getSpecConstsDefValBlob()), MOrigins(Origins) {
     updateSpecConstSymMap();
-    // SYCLBIN files have the kernel names embedded in the binaries, so we
-    // collect them.
-    if (BinImage && (MOrigins & ImageOriginSYCLBIN))
+    if (BinImage && (MOrigins & ImageOriginSYCLBIN)) {
+      // SYCLBIN files have the kernel names embedded in the binaries, so we
+      // collect them.
       for (const sycl_device_binary_property &KNProp :
            BinImage->getKernelNames())
         MKernelNames.insert(KNProp->Name);
+
+      KernelArgMask ArgMask;
+      if (BinImage->getKernelParamOptInfo().isAvailable()) {
+        // Extract argument mask from the image.
+        const RTDeviceBinaryImage::PropertyRange &KPOIRange =
+            BinImage->getKernelParamOptInfo();
+        for (const auto &Info : KPOIRange)
+          MEliminatedKernelArgMasks[Info->Name] =
+              createKernelArgMask(DeviceBinaryProperty(Info).asByteArray());
+      }
+    }
   }
 
   device_image_impl(
@@ -271,10 +283,12 @@ class device_image_impl
       const std::vector<unsigned char> &SpecConstsBlob, uint8_t Origins,
       std::optional<KernelCompilerBinaryInfo> &&RTCInfo,
       KernelNameSetT &&KernelNames,
+      KernelNameToArgMaskMap &&EliminatedKernelArgMasks,
       std::unique_ptr<DynRTDeviceBinaryImage> &&MergedImageStorage, private_tag)
       : MBinImage(BinImage), MContext(std::move(Context)),
         MDevices(std::move(Devices)), MState(State), MProgram(Program),
         MKernelIDs(std::move(KernelIDs)), MKernelNames{std::move(KernelNames)},
+        MEliminatedKernelArgMasks{std::move(EliminatedKernelArgMasks)},
         MSpecConstsBlob(SpecConstsBlob),
         MSpecConstsDefValBlob(getSpecConstsDefValBlob()),
         MSpecConstSymMap(SpecConstMap), MOrigins(Origins),
@@ -284,11 +298,14 @@ class device_image_impl
   device_image_impl(const RTDeviceBinaryImage *BinImage, const context &Context,
                     const std::vector<device> &Devices, bundle_state State,
                     ur_program_handle_t Program, syclex::source_language Lang,
-                    KernelNameSetT &&KernelNames, private_tag)
+                    KernelNameSetT &&KernelNames,
+                    KernelNameToArgMaskMap &&EliminatedKernelArgMasks,
+                    private_tag)
       : MBinImage(BinImage), MContext(std::move(Context)),
         MDevices(std::move(Devices)), MState(State), MProgram(Program),
         MKernelIDs(std::make_shared<std::vector<kernel_id>>()),
         MKernelNames{std::move(KernelNames)},
+        MEliminatedKernelArgMasks{std::move(EliminatedKernelArgMasks)},
         MSpecConstsDefValBlob(getSpecConstsDefValBlob()),
         MOrigins(ImageOriginKernelCompiler),
         MRTCBinInfo(KernelCompilerBinaryInfo{Lang}) {
@@ -669,6 +686,10 @@ class device_image_impl
 
   const KernelNameSetT &getKernelNames() const noexcept { return MKernelNames; }
 
+  const KernelNameToArgMaskMap &getEliminatedKernelArgMasks() const noexcept {
+    return MEliminatedKernelArgMasks;
+  }
+
   bool isNonSYCLSourceBased() const noexcept {
     return (getOriginMask() & ImageOriginKernelCompiler) &&
            !isFromSourceLanguage(syclex::source_language::sycl);
@@ -1261,6 +1282,10 @@ class device_image_impl
   // List of known kernel names.
   KernelNameSetT MKernelNames;
 
+  // Map for storing kernel argument masks for kernels. This is currently only
+  // used for images created from SYCLBIN.
+  KernelNameToArgMaskMap MEliminatedKernelArgMasks;
+
   // A mutex for sycnhronizing access to spec constants blob. Mutable because
   // needs to be locked in the const method for getting spec constant value.
   mutable std::mutex MSpecConstAccessMtx;

sycl/source/detail/program_manager/program_manager.cpp

@@ -2871,6 +2871,8 @@ ProgramManager::compile(const DevImgPlainWithDeps &ImgWithDeps,
       setSpecializationConstants(InputImpl, Prog, Adapter);
 
     KernelNameSetT KernelNames = InputImpl.getKernelNames();
+    std::unordered_map<std::string, KernelArgMask> EliminatedKernelArgMasks =
+        InputImpl.getEliminatedKernelArgMasks();
 
     std::optional<detail::KernelCompilerBinaryInfo> RTCInfo =
         InputImpl.getRTCInfo();
@@ -2881,7 +2883,7 @@ ProgramManager::compile(const DevImgPlainWithDeps &ImgWithDeps,
         InputImpl.get_spec_const_data_ref(),
         InputImpl.get_spec_const_blob_ref(), InputImpl.getOriginMask(),
         std::move(RTCInfo), std::move(KernelNames),
-        /*MergedImageStorage = */ nullptr);
+        std::move(EliminatedKernelArgMasks), nullptr);
 
     std::string CompileOptions;
     applyCompileOptionsFromEnvironment(CompileOptions);
@@ -3070,20 +3072,25 @@ ProgramManager::link(const std::vector<device_image_plain> &Imgs,
       RTCInfoPtrs;
   RTCInfoPtrs.reserve(Imgs.size());
   KernelNameSetT MergedKernelNames;
+  std::unordered_map<std::string, KernelArgMask> MergedEliminatedKernelArgMasks;
   for (const device_image_plain &DevImg : Imgs) {
     const DeviceImageImplPtr &DevImgImpl = getSyclObjImpl(DevImg);
     CombinedOrigins |= DevImgImpl->getOriginMask();
     RTCInfoPtrs.emplace_back(&(DevImgImpl->getRTCInfo()));
     MergedKernelNames.insert(DevImgImpl->getKernelNames().begin(),
                              DevImgImpl->getKernelNames().end());
+    MergedEliminatedKernelArgMasks.insert(
+        DevImgImpl->getEliminatedKernelArgMasks().begin(),
+        DevImgImpl->getEliminatedKernelArgMasks().end());
   }
   auto MergedRTCInfo = detail::KernelCompilerBinaryInfo::Merge(RTCInfoPtrs);
 
   DeviceImageImplPtr ExecutableImpl = device_image_impl::create(
       NewBinImg, Context, std::vector<device>{Devs}, bundle_state::executable,
       std::move(KernelIDs), LinkedProg, std::move(NewSpecConstMap),
       std::move(NewSpecConstBlob), CombinedOrigins, std::move(MergedRTCInfo),
-      std::move(MergedKernelNames), std::move(MergedImageStorage));
+      std::move(MergedKernelNames), std::move(MergedEliminatedKernelArgMasks),
+      std::move(MergedImageStorage));
 
   // TODO: Make multiple sets of device images organized by devices they are
   // compiled for.
@@ -3151,11 +3158,15 @@ ProgramManager::build(const DevImgPlainWithDeps &DevImgWithDeps,
       RTCInfoPtrs;
   RTCInfoPtrs.reserve(DevImgWithDeps.size());
   KernelNameSetT MergedKernelNames;
+  std::unordered_map<std::string, KernelArgMask> MergedEliminatedKernelArgMasks;
   for (const device_image_plain &DevImg : DevImgWithDeps) {
     const auto &DevImgImpl = getSyclObjImpl(DevImg);
     RTCInfoPtrs.emplace_back(&(DevImgImpl->getRTCInfo()));
     MergedKernelNames.insert(DevImgImpl->getKernelNames().begin(),
                              DevImgImpl->getKernelNames().end());
+    MergedEliminatedKernelArgMasks.insert(
+        DevImgImpl->getEliminatedKernelArgMasks().begin(),
+        DevImgImpl->getEliminatedKernelArgMasks().end());
   }
   auto MergedRTCInfo = detail::KernelCompilerBinaryInfo::Merge(RTCInfoPtrs);
 
@@ -3164,7 +3175,7 @@ ProgramManager::build(const DevImgPlainWithDeps &DevImgWithDeps,
       bundle_state::executable, std::move(KernelIDs), ResProgram,
       std::move(SpecConstMap), std::move(SpecConstBlob), CombinedOrigins,
       std::move(MergedRTCInfo), std::move(MergedKernelNames),
-      std::move(MergedImageStorage));
+      std::move(MergedEliminatedKernelArgMasks), std::move(MergedImageStorage));
   return createSyclObjFromImpl<device_image_plain>(std::move(ExecImpl));
 }
 

sycl/test-e2e/SYCLBIN/Inputs/dae.hpp

@@ -0,0 +1,49 @@
+#include "common.hpp"
+
+#include <sycl/usm.hpp>
+
+static constexpr size_t NUM = 1024;
+static constexpr size_t WGSIZE = 16;
+static constexpr float EPS = 0.001;
+
+int main(int argc, char *argv[]) {
+  assert(argc == 2);
+
+  sycl::queue Q;
+
+  int Failed = CommonLoadCheck(Q.get_context(), argv[1]);
+
+#if defined(SYCLBIN_INPUT_STATE)
+  auto KBInput = syclexp::get_kernel_bundle<sycl::bundle_state::input>(
+      Q.get_context(), std::string{argv[1]});
+  auto KBExe = sycl::build(KBInput);
+#elif defined(SYCLBIN_OBJECT_STATE)
+  auto KBObj = syclexp::get_kernel_bundle<sycl::bundle_state::object>(
+      Q.get_context(), std::string{argv[1]});
+  auto KBExe = sycl::link(KBObj);
+#else // defined(SYCLBIN_EXECUTABLE_STATE)
+  auto KBExe = syclexp::get_kernel_bundle<sycl::bundle_state::executable>(
+      Q.get_context(), std::string{argv[1]});
+#endif
+
+  assert(KBExe.ext_oneapi_has_kernel("iota"));
+  sycl::kernel IotaKern = KBExe.ext_oneapi_get_kernel("iota");
+
+  float *Ptr = sycl::malloc_shared<float>(NUM, Q);
+  Q.submit([&](sycl::handler &CGH) {
+     // First arugment is unused, but should still be passed, even if eliminated
+     // by DAE.
+     CGH.set_args(3.14f, Ptr);
+     CGH.parallel_for(sycl::nd_range{{NUM}, {WGSIZE}}, IotaKern);
+   }).wait_and_throw();
+
+  for (int I = 0; I < NUM; I++) {
+    const float Truth = static_cast<float>(I);
+    if (std::abs(Ptr[I] - Truth) > EPS) {
+      std::cout << "Result: " << Ptr[I] << " expected " << I << "\n";
+      ++Failed;
+    }
+  }
+  sycl::free(Ptr, Q);
+  return Failed;
+}

sycl/test-e2e/SYCLBIN/Inputs/dae_kernel.cpp

@@ -0,0 +1,10 @@
+#include <sycl/sycl.hpp>
+
+namespace syclexp = sycl::ext::oneapi::experimental;
+namespace syclext = sycl::ext::oneapi;
+
+extern "C" SYCL_EXT_ONEAPI_FUNCTION_PROPERTY(
+    (syclexp::nd_range_kernel<1>)) void iota(float, float *ptr) {
+  size_t id = syclext::this_work_item::get_nd_item<1>().get_global_linear_id();
+  ptr[id] = static_cast<float>(id);
+}

sycl/test-e2e/SYCLBIN/dae_executable.cpp

@@ -0,0 +1,23 @@
+//==----------- dae_executable.cpp --- SYCLBIN extension tests -------------==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// REQUIRES: aspect-usm_device_allocations
+
+// -- Test for using a kernel from a SYCLBIN with a dead argument.
+
+// SYCLBIN currently only properly detects SPIR-V binaries.
+// XFAIL: !target-spir
+// XFAIL-TRACKER: CMPLRLLVM-68811
+
+// RUN: %clangxx --offload-new-driver -fsyclbin=executable %{sycl_target_opts} %S/Inputs/dae_kernel.cpp -o %t.syclbin
+// RUN: %{build} -o %t.out
+// RUN: %{l0_leak_check} %{run} %t.out %t.syclbin
+
+#define SYCLBIN_EXECUTABLE_STATE
+
+#include "Inputs/dae.hpp"

sycl/test-e2e/SYCLBIN/dae_input.cpp

@@ -0,0 +1,23 @@
+//==----------- dae_input.cpp --- SYCLBIN extension tests ------------------==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// REQUIRES: aspect-usm_device_allocations
+
+// -- Test for using a kernel from a SYCLBIN with a dead argument.
+
+// SYCLBIN currently only properly detects SPIR-V binaries.
+// XFAIL: !target-spir
+// XFAIL-TRACKER: CMPLRLLVM-68811
+
+// RUN: %clangxx --offload-new-driver -fsyclbin=input %{sycl_target_opts} %S/Inputs/dae_kernel.cpp -o %t.syclbin
+// RUN: %{build} -o %t.out
+// RUN: %{l0_leak_check} %{run} %t.out %t.syclbin
+
+#define SYCLBIN_INPUT_STATE
+
+#include "Inputs/dae.hpp"

sycl/test-e2e/SYCLBIN/dae_object.cpp

@@ -0,0 +1,23 @@
+//==----------- dae_object.cpp --- SYCLBIN extension tests -----------------==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// REQUIRES: aspect-usm_device_allocations
+
+// -- Test for using a kernel from a SYCLBIN with a dead argument.
+
+// SYCLBIN currently only properly detects SPIR-V binaries.
+// XFAIL: !target-spir
+// XFAIL-TRACKER: CMPLRLLVM-68811
+
+// RUN: %clangxx --offload-new-driver -fsyclbin=object %{sycl_target_opts} %S/Inputs/dae_kernel.cpp -o %t.syclbin
+// RUN: %{build} -o %t.out
+// RUN: %{l0_leak_check} %{run} %t.out %t.syclbin
+
+#define SYCLBIN_OBJECT_STATE
+
+#include "Inputs/dae.hpp"