[SYCL][Fusion] Enable fusion of kernels with different ND-ranges (#8209)

All kernels with the same (or unspecified) local size and offset can be
fused. In order to make this work, some builtins getting index space
information must be remapped and the resulting ND-range of the fused
kernel, calculated.

The ND-range of the fused kernel will have:
1. The same number of dimensions as the input ND-range with the higher
number of dimensions;
2. The same local size as the shared local size (or unspecified)
3. The same offset as the shared offset
4. The global size will be the **greatest** input global size as per the
following ordering:
  i. Number of work items (enforces correctness);
  ii. Number of occurrences (less remappings needed);
  iii. Lexical order of the dimensions (introduces determinism).

Builtins obtaining the local/global size/id, work-group id, number of
work-groups or offset are remapped introducing as per an alwaysinline
function that can be reused along the fusion pass. More information can
be found in the Builtins.cpp file, where the remapping logic is
implemented.

---------

Signed-off-by: Victor Perez <[email protected]>

Author: victor-eds

sycl-fusion/CMakeLists.txt

@@ -12,6 +12,7 @@ set(LLVM_SPIRV_INCLUDE_DIRS "${LLVM_MAIN_SRC_DIR}/../llvm-spirv/include")
 if(WIN32)
   message(WARNING "Kernel fusion not yet supported on Windows")
 else(WIN32)
+  add_subdirectory(common)
   add_subdirectory(jit-compiler)
   add_subdirectory(passes)
   add_subdirectory(test)

sycl-fusion/common/CMakeLists.txt

@@ -0,0 +1,23 @@
+add_llvm_library(sycl-fusion-common
+  lib/NDRangesHelper.cpp
+)
+
+target_include_directories(sycl-fusion-common
+  PUBLIC
+  ${CMAKE_CURRENT_SOURCE_DIR}/include
+  ${CMAKE_CURRENT_SOURCE_DIR}/lib
+)
+
+if (BUILD_SHARED_LIBS)
+  if(NOT MSVC AND NOT APPLE)
+    # Manage symbol visibility through the linker to make sure no LLVM symbols
+    # are exported and confuse the drivers.
+    set(linker_script "${CMAKE_CURRENT_SOURCE_DIR}/ld-version-script.txt")
+    target_link_libraries(
+      sycl-fusion-common PRIVATE "-Wl,--version-script=${linker_script}")
+    set_target_properties(sycl-fusion-common
+      PROPERTIES
+      LINK_DEPENDS
+      ${linker_script})
+  endif()
+endif()

sycl-fusion/common/include/Kernel.h

@@ -10,6 +10,8 @@
 #define SYCL_FUSION_COMMON_KERNEL_H
 
 #include <algorithm>
+#include <array>
+#include <cassert>
 #include <string>
 #include <vector>
 
@@ -97,6 +99,77 @@ struct SYCLArgumentDescriptor {
 /// List of SYCL/OpenCL kernel attributes.
 using AttributeList = std::vector<SYCLKernelAttribute>;
 
+using Indices = std::array<size_t, 3>;
+
+///
+/// Class to model SYCL nd_range
+class NDRange {
+public:
+  constexpr static Indices AllZeros{0, 0, 0};
+
+  ///
+  /// Return the product of each index in an indices array.
+  constexpr static size_t linearize(const jit_compiler::Indices &I) {
+    return I[0] * I[1] * I[2];
+  }
+
+  NDRange() : NDRange{1, {1, 1, 1}} {}
+
+  NDRange(int Dimensions, const Indices &GlobalSize,
+          const Indices &LocalSize = {1, 1, 1},
+          const Indices &Offset = {0, 0, 0})
+      : Dimensions{Dimensions},
+        GlobalSize{GlobalSize}, LocalSize{LocalSize}, Offset{Offset} {
+#ifndef NDEBUG
+    const auto CheckDim = [Dimensions](const Indices &Range) {
+      return std::all_of(Range.begin() + Dimensions, Range.end(),
+                         [](auto D) { return D == 1; });
+    };
+    const auto CheckOffsetDim = [Dimensions](const Indices &Offset) {
+      return std::all_of(Offset.begin() + Dimensions, Offset.end(),
+
+                         [](auto D) { return D == 0; });
+    };
+#endif // NDEBUG
+    assert(CheckDim(GlobalSize) &&
+           "Invalid global range for number of dimensions");
+    assert(
+        (CheckDim(LocalSize) || std::all_of(LocalSize.begin(), LocalSize.end(),
+                                            [](auto D) { return D == 0; })) &&
+        "Invalid local range for number of dimensions");
+    assert(CheckOffsetDim(Offset) && "Invalid offset for number of dimensions");
+  }
+
+  constexpr const Indices &getGlobalSize() const { return GlobalSize; }
+  constexpr const Indices &getLocalSize() const { return LocalSize; }
+  constexpr const Indices &getOffset() const { return Offset; }
+  constexpr int getDimensions() const { return Dimensions; }
+
+  bool hasSpecificLocalSize() const { return LocalSize != AllZeros; }
+
+  friend constexpr bool operator==(const NDRange &LHS, const NDRange &RHS) {
+    return LHS.Dimensions == RHS.Dimensions &&
+           LHS.GlobalSize == RHS.GlobalSize &&
+           (!LHS.hasSpecificLocalSize() || !RHS.hasSpecificLocalSize() ||
+            LHS.LocalSize == RHS.LocalSize) &&
+           LHS.Offset == RHS.Offset;
+  }
+
+  friend constexpr bool operator!=(const NDRange &LHS, const NDRange &RHS) {
+    return !(LHS == RHS);
+  }
+
+private:
+  /** @brief The number of dimensions. */
+  int Dimensions;
+  /** @brief The local range. */
+  Indices GlobalSize;
+  /** @brief The local range. */
+  Indices LocalSize;
+  /** @brief The offet. */
+  Indices Offset;
+};
+
 /// Information about a kernel from DPC++.
 struct SYCLKernelInfo {
 
@@ -106,18 +179,21 @@ struct SYCLKernelInfo {
 
   AttributeList Attributes;
 
+  NDRange NDR;
+
   SYCLKernelBinaryInfo BinaryInfo;
 
   //// Explicit constructor for compatibility with LLVM YAML I/O.
-  SYCLKernelInfo() : Name{}, Args{}, Attributes{}, BinaryInfo{} {}
+  SYCLKernelInfo() : Name{}, Args{}, Attributes{}, NDR{}, BinaryInfo{} {}
 
   SYCLKernelInfo(const std::string &KernelName,
-                 const SYCLArgumentDescriptor &ArgDesc,
+                 const SYCLArgumentDescriptor &ArgDesc, const NDRange &NDR,
                  const SYCLKernelBinaryInfo &BinInfo)
-      : Name{KernelName}, Args{ArgDesc}, Attributes{}, BinaryInfo{BinInfo} {}
+      : Name{KernelName}, Args{ArgDesc}, Attributes{}, NDR{NDR}, BinaryInfo{
+                                                                     BinInfo} {}
 
   explicit SYCLKernelInfo(const std::string &KernelName)
-      : Name{KernelName}, Args{}, Attributes{}, BinaryInfo{} {}
+      : Name{KernelName}, Args{}, Attributes{}, NDR{}, BinaryInfo{} {}
 };
 
 ///

sycl-fusion/common/ld-version-script.txt

@@ -0,0 +1,8 @@
+{
+  global:
+    /* Export everything from jit_compiler namespace */
+    _ZN12jit_compiler*;
+
+  local:
+    *;
+};

sycl-fusion/common/include/KernelIO.h renamed to sycl-fusion/common/lib/KernelIO.h

@@ -0,0 +1,99 @@
+//==-------------------------- NDRangesHelper.cpp --------------------------==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "NDRangesHelper.h"
+
+#include <map>
+
+using namespace jit_compiler;
+using namespace llvm;
+
+///
+/// Return the maximal global size using the following order:
+/// 1. A range is greater than another range if it contains more elements (see
+/// linearize());
+/// 2. Else if it appears more times in the input list of ranges;
+/// 3. Else if it is greater in lexicographical order.
+static Indices getMaximalGlobalSize(ArrayRef<NDRange> NDRanges) {
+  size_t NumElements{0};
+  std::map<Indices, size_t> FreqMap;
+  for (const auto &ND : NDRanges) {
+    const auto &GS = ND.getGlobalSize();
+    const auto N = NDRange::linearize(GS);
+    if (N < NumElements) {
+      continue;
+    }
+    if (N > NumElements) {
+      NumElements = N;
+      FreqMap.clear();
+    }
+    ++FreqMap[GS];
+  }
+  return std::max_element(FreqMap.begin(), FreqMap.end(),
+                          [](const auto &LHS, const auto &RHS) {
+                            const auto LHSN = LHS.second;
+                            const auto RHSN = RHS.second;
+                            if (LHSN < RHSN) {
+                              return true;
+                            }
+                            if (LHSN > RHSN) {
+                              return false;
+                            }
+                            return LHS.first < RHS.first;
+                          })
+      ->first;
+}
+
+static bool compatibleRanges(const NDRange &LHS, const NDRange &RHS) {
+  const auto Dimensions = std::max(LHS.getDimensions(), RHS.getDimensions());
+  const auto EqualIndices = [Dimensions](const Indices &LHS,
+                                         const Indices &RHS) {
+    return std::equal(LHS.begin(), LHS.begin() + Dimensions, RHS.begin());
+  };
+  return (!LHS.hasSpecificLocalSize() || !RHS.hasSpecificLocalSize() ||
+          EqualIndices(LHS.getLocalSize(), RHS.getLocalSize())) &&
+         EqualIndices(LHS.getOffset(), RHS.getOffset());
+}
+
+NDRange jit_compiler::combineNDRanges(ArrayRef<NDRange> NDRanges) {
+  assert(isValidCombination(NDRanges) && "Invalid ND-ranges combination");
+  const auto Dimensions =
+      std::max_element(NDRanges.begin(), NDRanges.end(),
+                       [](const auto &LHS, const auto &RHS) {
+                         return LHS.getDimensions() < RHS.getDimensions();
+                       })
+          ->getDimensions();
+  const auto GlobalSize = getMaximalGlobalSize(NDRanges);
+  const auto *End = NDRanges.end();
+  const auto *LocalSizeIter = findSpecifiedLocalSize(NDRanges);
+  const auto &LocalSize =
+      LocalSizeIter == End ? NDRange::AllZeros : LocalSizeIter->getLocalSize();
+  const auto &Front = NDRanges.front();
+  return {Dimensions, GlobalSize, LocalSize, Front.getOffset()};
+}
+
+bool jit_compiler::isHeterogeneousList(ArrayRef<NDRange> NDRanges) {
+  const auto *FirstSpecifiedLocalSize = findSpecifiedLocalSize(NDRanges);
+  const auto &ND = FirstSpecifiedLocalSize == NDRanges.end()
+                       ? NDRanges.front()
+                       : *FirstSpecifiedLocalSize;
+  return any_of(NDRanges, [&ND](const auto &Other) { return ND != Other; });
+}
+
+bool jit_compiler::isValidCombination(llvm::ArrayRef<NDRange> NDRanges) {
+  if (NDRanges.empty()) {
+    return false;
+  }
+  const auto *FirstSpecifiedLocalSize = findSpecifiedLocalSize(NDRanges);
+  const auto &ND = FirstSpecifiedLocalSize == NDRanges.end()
+                       ? NDRanges.front()
+                       : *FirstSpecifiedLocalSize;
+  return llvm::all_of(NDRanges, [&ND](const auto &Other) {
+    return compatibleRanges(ND, Other);
+  });
+}

sycl-fusion/common/lib/NDRangesHelper.cpp

@@ -0,0 +1,37 @@
+//==--------- NDRangesHelper.h - Helpers to handle ND-ranges ---------------==//
+//
+// 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_NDRANGESHELPER_H
+#define SYCL_FUSION_COMMON_NDRANGESHELPER_H
+
+#include "Kernel.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+
+namespace jit_compiler {
+///
+/// Combine a list of ND-ranges, obtaining the resulting "fused" ND-range.
+NDRange combineNDRanges(llvm::ArrayRef<NDRange> NDRanges);
+
+inline llvm::ArrayRef<NDRange>::const_iterator
+findSpecifiedLocalSize(llvm::ArrayRef<NDRange> NDRanges) {
+  return llvm::find_if(
+      NDRanges, [](const auto &ND) { return ND.hasSpecificLocalSize(); });
+}
+
+///
+/// Returns whether the input list of ND-ranges is heterogeneous or not.
+bool isHeterogeneousList(llvm::ArrayRef<NDRange> NDRanges);
+
+///
+/// Return whether a combination of ND-ranges is valid for fusion.
+bool isValidCombination(llvm::ArrayRef<NDRange> NDRanges);
+} // namespace jit_compiler
+
+#endif // SYCL_FUSION_COMMON_NDRANGESHELPER_H

sycl-fusion/common/lib/NDRangesHelper.h

@@ -12,6 +12,7 @@ add_llvm_library(sycl-fusion
    Core
    Support
    Analysis
+   IPO
    TransformUtils
    Passes
    Linker
@@ -33,6 +34,7 @@ find_package(Threads REQUIRED)
 
 target_link_libraries(sycl-fusion
   PRIVATE
+  sycl-fusion-common
   LLVMSPIRVLib
   SYCLKernelFusionPasses
   ${CMAKE_THREAD_LIBS_INIT}

sycl-fusion/jit-compiler/CMakeLists.txt

@@ -9,6 +9,7 @@
 #ifndef SYCL_FUSION_JIT_COMPILER_HASHING_H
 #define SYCL_FUSION_JIT_COMPILER_HASHING_H
 
+#include "Kernel.h"
 #include "Parameter.h"
 
 #include "llvm/ADT/Hashing.h"
@@ -32,13 +33,23 @@ inline llvm::hash_code hash_value(const JITConstant &C) {
 inline llvm::hash_code hash_value(const ParameterIdentity &IP) {
   return llvm::hash_combine(IP.LHS, IP.RHS);
 }
+
+inline llvm::hash_code hash_value(const NDRange &ND) {
+  return llvm::hash_combine(ND.getDimensions(), ND.getGlobalSize(),
+                            ND.getLocalSize(), ND.getOffset());
+}
 } // namespace jit_compiler
 
 namespace std {
 template <typename T> inline llvm::hash_code hash_value(const vector<T> &V) {
   return llvm::hash_combine_range(V.begin(), V.end());
 }
 
+template <typename T, std::size_t N>
+inline llvm::hash_code hash_value(const array<T, N> &A) {
+  return llvm::hash_combine_range(A.begin(), A.end());
+}
+
 template <typename... T> struct hash<tuple<T...>> {
   size_t operator()(const tuple<T...> &Tuple) const noexcept {
     return llvm::hash_value(Tuple);

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

@@ -29,7 +29,11 @@ namespace jit_compiler {
 
 using CacheKeyT =
     std::tuple<std::vector<std::string>, ParamIdentList, int,
-               std::vector<ParameterInternalization>, std::vector<JITConstant>>;
+               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.
+               std::optional<std::vector<NDRange>>>;
 
 ///
 /// Wrapper around a SPIR-V binary.