From 5ed3e5db82c5d753d6c0206e12bc7bb07505ad22 Mon Sep 17 00:00:00 2001
From: Sergey Semenov <sergey.semenov@intel.com>
Date: Thu, 10 Oct 2024 08:32:29 -0700
Subject: [PATCH] [SYCL] Fix devicelib handling when linking multiple images

Handle devicelib requirements of image dependecies, not just the main
image.
---
 .../program_manager/program_manager.cpp       |  9 ++++---
 .../DeviceDependencies/math_device_lib.cpp    | 27 +++++++++++++++++++
 2 files changed, 33 insertions(+), 3 deletions(-)
 create mode 100644 sycl/test-e2e/DeviceDependencies/math_device_lib.cpp

diff --git a/sycl/source/detail/program_manager/program_manager.cpp b/sycl/source/detail/program_manager/program_manager.cpp
index 6ab92c6e27d04..e48255682749c 100644
--- a/sycl/source/detail/program_manager/program_manager.cpp
+++ b/sycl/source/detail/program_manager/program_manager.cpp
@@ -816,9 +816,10 @@ ur_program_handle_t ProgramManager::getBuiltURProgram(
     // If device image is not SPIR-V, DeviceLibReqMask will be 0 which means
     // no fallback device library will be linked.
     uint32_t DeviceLibReqMask = 0;
-    if (!DeviceCodeWasInCache &&
-        Img.getFormat() == SYCL_DEVICE_BINARY_TYPE_SPIRV &&
-        !SYCLConfig<SYCL_DEVICELIB_NO_FALLBACK>::get())
+    bool UseDeviceLibs = !DeviceCodeWasInCache &&
+                         Img.getFormat() == SYCL_DEVICE_BINARY_TYPE_SPIRV &&
+                         !SYCLConfig<SYCL_DEVICELIB_NO_FALLBACK>::get();
+    if (UseDeviceLibs)
       DeviceLibReqMask = getDeviceLibReqMask(Img);
 
     std::vector<ur_program_handle_t> ProgramsToLink;
@@ -826,6 +827,8 @@ ur_program_handle_t ProgramManager::getBuiltURProgram(
     // program already.
     if (!DeviceCodeWasInCache) {
       for (RTDeviceBinaryImage *BinImg : DeviceImagesToLink) {
+        if (UseDeviceLibs)
+          DeviceLibReqMask |= getDeviceLibReqMask(*BinImg);
         device_image_plain DevImagePlain =
             getDeviceImageFromBinaryImage(BinImg, Context, Device);
         const std::shared_ptr<detail::device_image_impl> &DeviceImageImpl =
diff --git a/sycl/test-e2e/DeviceDependencies/math_device_lib.cpp b/sycl/test-e2e/DeviceDependencies/math_device_lib.cpp
new file mode 100644
index 0000000000000..440935e31e3eb
--- /dev/null
+++ b/sycl/test-e2e/DeviceDependencies/math_device_lib.cpp
@@ -0,0 +1,27 @@
+// REQUIRES: aspect-fp64
+// UNSUPPORTED: hip || cuda
+
+// DEFINE: %{mathflags} = %if cl_options %{/clang:-fno-fast-math%} %else %{-fno-fast-math%}
+
+// RUN: %{build} -fsycl-allow-device-image-dependencies -fsycl-device-lib-jit-link %{mathflags} -o %t.out
+// RUN: %{run} %t.out
+
+#include <cmath>
+#include <sycl/detail/core.hpp>
+
+using namespace sycl;
+
+// Check that device lib dependencies are resolved with
+// -fsycl-allow-device-image-dependencies.
+// TODO this test will become redundant once
+// -fsycl-allow-device-image-dependencies is enabled by default.
+int main() {
+  range<1> Range{1};
+  queue q;
+  buffer<double, 1> buffer1(Range);
+  q.submit([&](sycl::handler &cgh) {
+    auto Acc = buffer1.get_access<access::mode::write>(cgh);
+    cgh.single_task<class DeviceMathTest>([=]() { Acc[0] = std::acosh(1.0); });
+  });
+  return 0;
+}