[SYCL][Fusion] Make target architecture part of JIT cache key (intel#12017)

JIT cache avoids running fusion JIT for previously found fused kernel
sequences with a given set of parameters. Add target architecture to
the set of parameters to avoid returning a previously fused kernel for
a different target.

Fused kernels for SPIR-V targets can be reused regardless of different
architectures.

---------

Signed-off-by: Victor Perez <victor.perez@codeplay.com>

Author: victor-eds

sycl-fusion/common/include/Kernel.h

@@ -61,6 +61,38 @@ enum class ParameterKind : uint32_t {
 /// Different binary formats supported as input to the JIT compiler.
 enum class BinaryFormat : uint32_t { INVALID, LLVM, SPIRV, PTX, AMDGCN };
 
+/// Unique ID for each supported architecture in the SYCL implementation.
+///
+/// Values of this type will only be used in the kernel fusion non-persistent
+/// JIT. There is no guarantee for backwards compatibility, so this should not
+/// be used in persistent caches.
+using DeviceArchitecture = unsigned;
+
+class TargetInfo {
+public:
+  static constexpr TargetInfo get(BinaryFormat Format,
+                                  DeviceArchitecture Arch) {
+    if (Format == BinaryFormat::SPIRV) {
+      /// As an exception, SPIR-V targets have a single common ID (-1), as fused
+      /// kernels will be reused across SPIR-V devices.
+      return {Format, DeviceArchitecture(-1)};
+    }
+    return {Format, Arch};
+  }
+
+  TargetInfo() = default;
+
+  constexpr BinaryFormat getFormat() const { return Format; }
+  constexpr DeviceArchitecture getArch() const { return Arch; }
+
+private:
+  constexpr TargetInfo(BinaryFormat Format, DeviceArchitecture Arch)
+      : Format(Format), Arch(Arch) {}
+
+  BinaryFormat Format;
+  DeviceArchitecture Arch;
+};
+
 /// Information about a device intermediate representation module (e.g., SPIR-V,
 /// LLVM IR) from DPC++.
 struct SYCLKernelBinaryInfo {

sycl-fusion/jit-compiler/include/JITContext.h

@@ -29,8 +29,9 @@ class LLVMContext;
 namespace jit_compiler {
 
 using CacheKeyT =
-    std::tuple<std::vector<std::string>, ParamIdentList, BarrierFlags,
-               std::vector<ParameterInternalization>, std::vector<JITConstant>,
+    std::tuple<DeviceArchitecture, std::vector<std::string>, ParamIdentList,
+               BarrierFlags, std::vector<ParameterInternalization>,
+               std::vector<JITConstant>,
                // This field of the cache is optional because, if all of the
                // ranges are equal, we will perform no remapping, so that fused
                // kernels can be reused with different lists of equal nd-ranges.

sycl-fusion/jit-compiler/include/Options.h

@@ -16,7 +16,7 @@
 
 namespace jit_compiler {
 
-enum OptionID { VerboseOutput, EnableCaching, TargetFormat };
+enum OptionID { VerboseOutput, EnableCaching, TargetDeviceInfo };
 
 class OptionPtrBase {};
 
@@ -80,8 +80,8 @@ struct JITEnableVerbose : public OptionBase<OptionID::VerboseOutput, bool> {};
 
 struct JITEnableCaching : public OptionBase<OptionID::EnableCaching, bool> {};
 
-struct JITTargetFormat
-    : public OptionBase<OptionID::TargetFormat, BinaryFormat> {};
+struct JITTargetInfo
+    : public OptionBase<OptionID::TargetDeviceInfo, TargetInfo> {};
 
 } // namespace option
 } // namespace jit_compiler

sycl-fusion/jit-compiler/lib/KernelFusion.cpp

@@ -94,15 +94,18 @@ FusionResult KernelFusion::fuseKernels(
 
   bool IsHeterogeneousList = jit_compiler::isHeterogeneousList(NDRanges);
 
-  BinaryFormat TargetFormat = ConfigHelper::get<option::JITTargetFormat>();
+  TargetInfo TargetInfo = ConfigHelper::get<option::JITTargetInfo>();
+  BinaryFormat TargetFormat = TargetInfo.getFormat();
+  DeviceArchitecture TargetArch = TargetInfo.getArch();
 
   if (!isTargetFormatSupported(TargetFormat)) {
     return FusionResult(
         "Fusion output target format not supported by this build");
   }
 
   bool CachingEnabled = ConfigHelper::get<option::JITEnableCaching>();
-  CacheKeyT CacheKey{KernelsToFuse,
+  CacheKeyT CacheKey{TargetArch,
+                     KernelsToFuse,
                      Identities,
                      BarriersFlags,
                      Internalization,

sycl/source/detail/device_impl.hpp

@@ -257,11 +257,13 @@ class device_impl {
   std::string
   get_device_info_string(sycl::detail::pi::PiDeviceInfo InfoCode) const;
 
+  /// Get device architecture
+  ext::oneapi::experimental::architecture getDeviceArch() const;
+
 private:
   explicit device_impl(pi_native_handle InteropDevice,
                        sycl::detail::pi::PiDevice Device,
                        PlatformImplPtr Platform, const PluginPtr &Plugin);
-  ext::oneapi::experimental::architecture getDeviceArch() const;
   sycl::detail::pi::PiDevice MDevice = 0;
   sycl::detail::pi::PiDeviceType MType;
   sycl::detail::pi::PiDevice MRootDevice = nullptr;

sycl/source/detail/jit_compiler.cpp

@@ -56,6 +56,13 @@ ::jit_compiler::BinaryFormat getTargetFormat(QueueImplPtr &Queue) {
   }
 }
 
+::jit_compiler::TargetInfo getTargetInfo(QueueImplPtr &Queue) {
+  ::jit_compiler::BinaryFormat Format = getTargetFormat(Queue);
+  return ::jit_compiler::TargetInfo::get(
+      Format, static_cast<::jit_compiler::DeviceArchitecture>(
+                  Queue->getDeviceImplPtr()->getDeviceArch()));
+}
+
 std::pair<const RTDeviceBinaryImage *, sycl::detail::pi::PiProgram>
 retrieveKernelBinary(QueueImplPtr &Queue, CGExecKernel *KernelCG) {
   auto KernelName = KernelCG->getKernelName();
@@ -824,8 +831,9 @@ jit_compiler::fuseKernels(QueueImplPtr Queue,
   JITConfig.set<::jit_compiler::option::JITEnableCaching>(
       detail::SYCLConfig<detail::SYCL_ENABLE_FUSION_CACHING>::get());
 
-  ::jit_compiler::BinaryFormat TargetFormat = getTargetFormat(Queue);
-  JITConfig.set<::jit_compiler::option::JITTargetFormat>(TargetFormat);
+  ::jit_compiler::TargetInfo TargetInfo = getTargetInfo(Queue);
+  ::jit_compiler::BinaryFormat TargetFormat = TargetInfo.getFormat();
+  JITConfig.set<::jit_compiler::option::JITTargetInfo>(TargetInfo);
 
   auto FusionResult = ::jit_compiler::KernelFusion::fuseKernels(
       *MJITContext, std::move(JITConfig), InputKernelInfo, InputKernelNames,

sycl/test-e2e/KernelFusion/jit_caching_multispirv.cpp

@@ -0,0 +1,19 @@
+// REQUIRES: fusion, gpu, (opencl || level_zero)
+// RUN: %{build} -O2 -o %t.out
+// RUN: env SYCL_RT_WARNING_LEVEL=1 %{run-unfiltered-devices} %t.out 2>&1 | FileCheck %s --implicit-check-not "WRONG a VALUE" --implicit-check-not "WRONG b VALUE"
+
+// Test caching for JIT fused kernels when different SPIR-V devices are
+// involved.
+
+#include "./jit_caching_multitarget_common.h"
+
+// Initial invocation
+// CHECK: JIT DEBUG: Compiling new kernel, no suitable cached kernel found
+
+// Identical invocation, should lead to JIT cache hit.
+// CHECK-NEXT: JIT DEBUG: Re-using cached JIT kernel
+// CHECK-NEXT: INFO: Re-using existing device binary for fused kernel
+
+// Invocation with another SPIR-V device involved. Should lead to JIT cache hit.
+// CHECK-NEXT: JIT DEBUG: Re-using cached JIT kernel
+// CHECK-NEXT: INFO: Re-using existing device binary for fused kernel

sycl/test-e2e/KernelFusion/jit_caching_multitarget.cpp

@@ -0,0 +1,23 @@
+// REQUIRES: fusion, gpu, (hip || cuda)
+// RUN: %{build} -fsycl-embed-ir -O2 -o %t.out
+// RUN: env SYCL_RT_WARNING_LEVEL=1 %{run-unfiltered-devices} %t.out 2>&1 | FileCheck %s --implicit-check-not "WRONG a VALUE" --implicit-check-not "WRONG b VALUE"
+// XFAIL: *
+
+// COM: This test is expected to fail on CI, as CUDA and HIP runners do not
+// provide a CPU backend.
+
+// Test caching for JIT fused kernels when devices with different architectures
+// are involved.
+
+#include "./jit_caching_multitarget_common.h"
+
+// Initial invocation
+// CHECK: JIT DEBUG: Compiling new kernel, no suitable cached kernel found
+
+// Identical invocation, should lead to JIT cache hit.
+// CHECK-NEXT: JIT DEBUG: Re-using cached JIT kernel
+// CHECK-NEXT: INFO: Re-using existing device binary for fused kernel
+
+// Invocation with a device with a different architecture. Should trigger JIT
+// compilation.
+// CHECK-NEXT: JIT DEBUG: Compiling new kernel, no suitable cached kernel found

sycl/test-e2e/KernelFusion/jit_caching_multitarget_common.h

@@ -0,0 +1,55 @@
+// Test caching for JIT fused kernels when different devices are involved.
+
+#include <sycl/sycl.hpp>
+
+using namespace sycl;
+
+constexpr inline std::size_t size(1024);
+
+class Kernel0;
+class Kernel1;
+
+void performFusion(queue q, std::size_t *a, std::size_t *b) {
+  {
+    buffer<std::size_t> a_buf(a, size);
+    buffer<std::size_t> b_buf(b, size);
+
+    ext::codeplay::experimental::fusion_wrapper fw{q};
+    fw.start_fusion();
+    q.submit([&](handler &cgh) {
+      accessor a(a_buf, cgh, write_only, no_init);
+      cgh.parallel_for<Kernel0>(size, [=](id<1> i) { a[i] = i; });
+    });
+    q.submit([&](handler &cgh) {
+      accessor a(a_buf, cgh, read_only);
+      accessor b(b_buf, cgh, write_only, no_init);
+      cgh.parallel_for<Kernel1>(size, [=](id<1> i) { b[i] = a[i] * 2; });
+    });
+    fw.complete_fusion();
+  }
+  for (std::size_t i = 0; i < size; ++i) {
+    assert(a[i] == i && "WRONG a VALUE");
+    assert(b[i] == i * 2 && "WRONG b VALUE");
+  }
+}
+
+int main() {
+  queue q{gpu_selector_v,
+          ext::codeplay::experimental::property::queue::enable_fusion{}};
+  queue q_cpu{cpu_selector_v,
+              ext::codeplay::experimental::property::queue::enable_fusion{}};
+
+  std::vector<std::size_t> a(size);
+  std::vector<std::size_t> b(size);
+
+  // Initial invocation
+  performFusion(q, a.data(), b.data());
+
+  // Identical invocation, should lead to JIT cache hit.
+  performFusion(q, a.data(), b.data());
+
+  // Invocation on CPU device.
+  performFusion(q_cpu, a.data(), b.data());
+
+  return 0;
+}