[SYCL][Fusion] API for kernel fusion library (#7465)

This is the second patch in a series of patches to add an implementation of thekernel fusion extension. We have split the implementation into multiple patches to make them more easy to review.

This patch can be reviewed and merged independently of #7416.

This patch adds the first components for the JIT compiler used for implementation of kernel fusion at runtime , concretely:

API definitions
Input translation from SPIR-V to LLVM IR
Insertion of fused kernel function stub and metadata
Supporting infrastructure such as compiler options.
CMake logic to link and code to invoke this JIT from the SYCL runtime will follow in a later patch.

Co-authored-by: Lukas Sommer [email protected]
Co-authored-by: Victor Perez [email protected]
Signed-off-by: Lukas Sommer [email protected]

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,151 @@
+//==------- 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 getKernelFor(KernelName) != nullptr;
+  }
+
+  SYCLKernelInfo *getKernelFor(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;
+  }
+
+private:
+  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