[SYCL][Fusion] API for kernel fusion library

First components for the JIT compiler used for SYCL kernel fusion:
* API definitions
* Input translation from SPIR-V to LLVM IR
* Insertion of fused kernel function stub and metadata
* Supporting infrastructure such as compiler options.

Co-authored-by: Lukas Sommer <lukas.sommer@codeplay.com>
Co-authored-by: Victor Perez <victor.perez@codeplay.com>
Signed-off-by: Lukas Sommer <lukas.sommer@codeplay.com>

Author: sommerlukas

.github/CODEOWNERS

@@ -70,3 +70,6 @@ llvm/include/llvm/SYCLLowerIR/LowerInvokeSimd.h @kbobrovs @v-klochkov @rolandsch
 .github/workflows/ @intel/dpcpp-devops-reviewers
 buildbot/ @intel/dpcpp-devops-reviewers
 devops/ @intel/dpcpp-devops-reviewers
+
+# Kernel fusion JIT compiler
+sycl-fusion/ @victor-eds @Naghasan @sommerlukas

sycl-fusion/CMakeLists.txt

@@ -0,0 +1,12 @@
+cmake_minimum_required(VERSION 3.14 FATAL_ERROR)
+
+# Define a variable holding the root directory of the JIT compiler project
+# for use in includes etc.
+set(SYCL_JIT_BASE_DIR ${CMAKE_CURRENT_SOURCE_DIR})
+
+# For some reason, the LLVMSPIRVLib does not define any public includes.
+# To link against the library, define the following link to its include
+# directories, similar to how clang/CMakeLists.txt does it.
+set(LLVM_SPIRV_INCLUDE_DIRS "${LLVM_MAIN_SRC_DIR}/../llvm-spirv/include")
+
+add_subdirectory(jit-compiler)

sycl-fusion/README.md

@@ -0,0 +1,9 @@
+# SYCL Kernel Fusion Compiler
+
+Basic JIT compiler infrastructure to perform online kernel fusion of SYCL kernels at runtime.
+
+The experimental SYCL extension for kernel fusion is described in 
+[this proposal](../sycl/doc/extensions/experimental/sycl_ext_codeplay_kernel_fusion.asciidoc).
+
+The design of the JIT compiler is described in this 
+[design document](../sycl/doc/design/KernelFusionJIT.md).

sycl-fusion/common/include/Kernel.h

@@ -0,0 +1,155 @@
+//==------- Kernel.h - Representation of a SYCL kernel for JIT compiler ----==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SYCL_FUSION_COMMON_KERNEL_H
+#define SYCL_FUSION_COMMON_KERNEL_H
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+namespace jit_compiler {
+
+using BinaryAddress = const unsigned char *;
+
+///
+/// Enumerate possible kinds of parameters.
+/// 1:1 correspondence with the definition in kernel_desc.hpp in the DPC++ SYCL
+/// runtime.
+enum class ParameterKind : unsigned {
+  Accessor = 0,
+  StdLayout = 1,
+  Sampler = 2,
+  Pointer = 3,
+  SpecConstBuffer = 4,
+  Stream = 5,
+  Invalid = 0xF,
+};
+
+/// Different binary formats supported as input to the JIT compiler.
+enum class BinaryFormat { INVALID, LLVM, SPIRV };
+
+/// Information about a device intermediate representation module (e.g., SPIR-V,
+/// LLVM IR) from DPC++.
+struct SYCLKernelBinaryInfo {
+
+  BinaryFormat Format = BinaryFormat::INVALID;
+
+  size_t AddressBits = 0;
+
+  BinaryAddress BinaryStart = nullptr;
+
+  size_t BinarySize = 0;
+};
+
+///
+/// Describe a SYCL/OpenCL kernel attribute by its name and values.
+struct SYCLKernelAttribute {
+  using AttributeValueList = std::vector<std::string>;
+
+  // Explicit constructor for compatibility with LLVM YAML I/O.
+  SYCLKernelAttribute() : Values{} {};
+  SYCLKernelAttribute(std::string Name)
+      : AttributeName{std::move(Name)}, Values{} {}
+
+  std::string AttributeName;
+  AttributeValueList Values;
+};
+
+enum ArgUsage : unsigned char {
+  // Used to indicate that an argument is not used by the kernel
+  Unused = 0,
+  // Used to indicate that an argument is used by the kernel
+  Used = 1u,
+  // Used to indicate that the accessor/pointer argument has been promoted to
+  // private memory
+  PromotedPrivate = 1u << 4,
+  // Used to indicate that the accessor/pointer argument has been promoted to
+  // local memory
+  PromotedLocal = 1u << 5,
+};
+
+///
+/// Encode usage of parameters for the actual kernel function.
+// This is a vector of unsigned char, because std::vector<bool> is a weird
+// construct and unlike all other std::vectors, and LLVM YAML I/O is having a
+// hard time coping with it.
+using ArgUsageMask = std::vector<std::underlying_type_t<ArgUsage>>;
+
+///
+/// Describe the list of arguments by their kind.
+struct SYCLArgumentDescriptor {
+
+  // Explicit constructor for compatibility with LLVM YAML I/O.
+  SYCLArgumentDescriptor() : Kinds{}, UsageMask{} {}
+
+  std::vector<ParameterKind> Kinds;
+
+  ArgUsageMask UsageMask;
+};
+
+///
+/// List of SYCL/OpenCL kernel attributes.
+using AttributeList = std::vector<SYCLKernelAttribute>;
+
+/// Information about a kernel from DPC++.
+struct SYCLKernelInfo {
+
+  std::string Name;
+
+  SYCLArgumentDescriptor Args;
+
+  AttributeList Attributes;
+
+  SYCLKernelBinaryInfo BinaryInfo;
+
+  //// Explicit constructor for compatibility with LLVM YAML I/O.
+  SYCLKernelInfo() : Name{}, Args{}, Attributes{}, BinaryInfo{} {}
+
+  SYCLKernelInfo(const std::string &KernelName,
+                 const SYCLArgumentDescriptor &ArgDesc,
+                 const SYCLKernelBinaryInfo &BinInfo)
+      : Name{KernelName}, Args{ArgDesc}, Attributes{}, BinaryInfo{BinInfo} {}
+
+  explicit SYCLKernelInfo(const std::string &KernelName)
+      : Name{KernelName}, Args{}, Attributes{}, BinaryInfo{} {}
+};
+
+///
+/// Represents a SPIR-V translation unit containing SYCL kernels by the
+/// KernelInfo for each of the contained kernels.
+class SYCLModuleInfo {
+public:
+  using KernelInfoList = std::vector<SYCLKernelInfo>;
+
+  void addKernel(SYCLKernelInfo &Kernel) { Kernels.push_back(Kernel); }
+
+  KernelInfoList &kernels() { return Kernels; }
+
+  bool hasKernelFor(const std::string &KernelName) {
+    return findKernelFor(KernelName) != nullptr;
+  }
+
+  SYCLKernelInfo *getKernelFor(const std::string &KernelName) {
+    return findKernelFor(KernelName);
+  }
+
+private:
+  SYCLKernelInfo *findKernelFor(const std::string &KernelName) {
+    auto It =
+        std::find_if(Kernels.begin(), Kernels.end(),
+                     [&](SYCLKernelInfo &K) { return K.Name == KernelName; });
+    return (It != Kernels.end()) ? &*It : nullptr;
+  }
+
+  KernelInfoList Kernels;
+};
+
+} // namespace jit_compiler
+
+#endif // SYCL_FUSION_COMMON_KERNEL_H

sycl-fusion/common/include/KernelIO.h

@@ -0,0 +1,95 @@
+//==----- KernelIO.h - YAML output of internal SYCL kernel representation --==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SYCL_FUSION_COMMON_KERNELIO_H
+#define SYCL_FUSION_COMMON_KERNELIO_H
+
+#include "Kernel.h"
+#include "llvm/Support/YAMLParser.h"
+#include "llvm/Support/YAMLTraits.h"
+
+using llvm::yaml::IO;
+using llvm::yaml::MappingTraits;
+using llvm::yaml::ScalarEnumerationTraits;
+
+// Specify how to map std::vectors of different user-defined types to YAML
+// sequences.
+LLVM_YAML_IS_SEQUENCE_VECTOR(jit_compiler::ArgUsageMask)
+LLVM_YAML_IS_FLOW_SEQUENCE_VECTOR(jit_compiler::ParameterKind)
+LLVM_YAML_IS_SEQUENCE_VECTOR(jit_compiler::SYCLArgumentDescriptor)
+LLVM_YAML_IS_SEQUENCE_VECTOR(jit_compiler::SYCLKernelAttribute)
+LLVM_YAML_IS_SEQUENCE_VECTOR(jit_compiler::SYCLKernelInfo)
+
+//
+// Mapping traits for the different elements of KernelInfo.
+namespace llvm {
+namespace yaml {
+
+template <> struct ScalarEnumerationTraits<jit_compiler::ParameterKind> {
+  static void enumeration(IO &IO, jit_compiler::ParameterKind &PK) {
+    IO.enumCase(PK, "Accessor", jit_compiler::ParameterKind::Accessor);
+    IO.enumCase(PK, "StdLayout", jit_compiler::ParameterKind::StdLayout);
+    IO.enumCase(PK, "Sampler", jit_compiler::ParameterKind::Sampler);
+    IO.enumCase(PK, "Pointer", jit_compiler::ParameterKind::Pointer);
+    IO.enumCase(PK, "SpecConstantBuffer",
+                jit_compiler::ParameterKind::SpecConstBuffer);
+    IO.enumCase(PK, "Stream", jit_compiler::ParameterKind::Stream);
+    IO.enumCase(PK, "Invalid", jit_compiler::ParameterKind::Invalid);
+  }
+};
+
+template <> struct ScalarEnumerationTraits<jit_compiler::BinaryFormat> {
+  static void enumeration(IO &IO, jit_compiler::BinaryFormat &BF) {
+    IO.enumCase(BF, "LLVM", jit_compiler::BinaryFormat::LLVM);
+    IO.enumCase(BF, "SPIRV", jit_compiler::BinaryFormat::SPIRV);
+    IO.enumCase(BF, "INVALID", jit_compiler::BinaryFormat::INVALID);
+  }
+};
+
+template <> struct MappingTraits<jit_compiler::SYCLKernelBinaryInfo> {
+  static void mapping(IO &IO, jit_compiler::SYCLKernelBinaryInfo &BI) {
+    IO.mapRequired("Format", BI.Format);
+    IO.mapRequired("AddressBits", BI.AddressBits);
+    // We do not serialize the pointer here on purpose.
+    IO.mapRequired("BinarySize", BI.BinarySize);
+  }
+};
+
+template <> struct MappingTraits<jit_compiler::SYCLArgumentDescriptor> {
+  static void mapping(IO &IO, jit_compiler::SYCLArgumentDescriptor &AD) {
+    IO.mapRequired("Kinds", AD.Kinds);
+    IO.mapRequired("Mask", AD.UsageMask);
+  }
+};
+
+template <> struct MappingTraits<jit_compiler::SYCLKernelAttribute> {
+  static void mapping(IO &IO, jit_compiler::SYCLKernelAttribute &KA) {
+    IO.mapRequired("AttrName", KA.AttributeName);
+    IO.mapRequired("Values", KA.Values);
+  }
+};
+
+template <> struct MappingTraits<jit_compiler::SYCLKernelInfo> {
+  static void mapping(IO &IO, jit_compiler::SYCLKernelInfo &KI) {
+    IO.mapRequired("KernelName", KI.Name);
+    IO.mapRequired("Args", KI.Args);
+    IO.mapOptional("Attributes", KI.Attributes);
+    IO.mapRequired("BinInfo", KI.BinaryInfo);
+  }
+};
+
+template <> struct MappingTraits<jit_compiler::SYCLModuleInfo> {
+  static void mapping(IO &IO, jit_compiler::SYCLModuleInfo &SMI) {
+    IO.mapRequired("Kernels", SMI.kernels());
+  }
+};
+
+} // namespace yaml
+} // namespace llvm
+
+#endif // SYCL_FUSION_COMMON_KERNELIO_H

sycl-fusion/jit-compiler/CMakeLists.txt

@@ -0,0 +1,42 @@
+
+add_llvm_library(sycl-fusion
+  SHARED
+   lib/KernelFusion.cpp
+   lib/JITContext.cpp
+   lib/translation/SPIRVLLVMTranslation.cpp
+   lib/fusion/FusionHelper.cpp
+   lib/fusion/ModuleHelper.cpp
+   lib/helper/ConfigHelper.cpp
+
+   LINK_COMPONENTS
+   Core
+   Support
+   Analysis
+   TransformUtils
+   Passes
+   Linker
+   ScalarOpts
+   InstCombine
+)
+
+target_include_directories(sycl-fusion
+  PUBLIC
+  $<INSTALL_INTERFACE:include>
+  $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
+  $<BUILD_INTERFACE:${SYCL_JIT_BASE_DIR}/common/include>
+  PRIVATE
+  ${CMAKE_CURRENT_SOURCE_DIR}/lib
+  ${LLVM_SPIRV_INCLUDE_DIRS}
+)
+
+find_package(Threads REQUIRED)
+
+target_link_libraries(sycl-fusion
+  PRIVATE
+  LLVMSPIRVLib
+  ${CMAKE_THREAD_LIBS_INIT}
+)
+
+install(TARGETS sycl-fusion
+  LIBRARY DESTINATION "lib${LLVM_LIBDIR_SUFFIX}" COMPONENT sycl-fusion
+  RUNTIME DESTINATION "bin" COMPONENT sycl-fusion)

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

@@ -0,0 +1,65 @@
+//==------- JITContext.h - Context holding data for the JIT compiler -------==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SYCL_FUSION_JIT_COMPILER_JITCONTEXT_H
+#define SYCL_FUSION_JIT_COMPILER_JITCONTEXT_H
+
+#include "llvm/IR/LLVMContext.h"
+#include <shared_mutex>
+#include <unordered_map>
+
+#include "Kernel.h"
+#include "Parameter.h"
+
+namespace jit_compiler {
+
+///
+/// Wrapper around a SPIR-V binary.
+class SPIRVBinary {
+public:
+  explicit SPIRVBinary(std::string Binary);
+
+  jit_compiler::BinaryAddress address() const;
+
+  size_t size() const;
+
+private:
+  std::string Blob;
+};
+
+///
+/// Context to persistenly store information across invocations of the JIT
+/// compiler and manage lifetimes of binaries.
+class JITContext {
+
+public:
+  JITContext();
+
+  ~JITContext();
+
+  llvm::LLVMContext *getLLVMContext();
+
+  SPIRVBinary &emplaceSPIRVBinary(std::string Binary);
+
+private:
+  // FIXME: Change this to std::shared_mutex after switching to C++17.
+  using MutexT = std::shared_timed_mutex;
+
+  using ReadLockT = std::shared_lock<MutexT>;
+
+  using WriteLockT = std::unique_lock<MutexT>;
+
+  std::unique_ptr<llvm::LLVMContext> LLVMCtx;
+
+  MutexT BinariesMutex;
+
+  std::vector<SPIRVBinary> Binaries;
+};
+} // namespace jit_compiler
+
+#endif // SYCL_FUSION_JIT_COMPILER_JITCONTEXT_H