[SYCL][Fusion] Enable fusion of kernels with different ND-ranges

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/common/include/Kernel.h

@@ -10,6 +10,8 @@
 #define SYCL_FUSION_COMMON_KERNEL_H
 
 #include <algorithm>
+#include <cassert>
+#include <map>
 #include <string>
 #include <vector>
 
@@ -97,6 +99,112 @@ 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};
+
+  ///
+  /// Returns whether or not two ranges can be fused.
+  ///
+  /// In order to be fusable, two ranges must:
+  /// 1. Both have the local size specified or not;
+  /// 2. If the local size is specified, it must be equal.
+  /// 3. Have the same offset;
+  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 (EqualIndices(LHS.getLocalSize(), RHS.getLocalSize()) ||
+            LHS.getLocalSize() == AllZeros || RHS.getLocalSize() == AllZeros) &&
+           EqualIndices(LHS.getOffset(), RHS.getOffset());
+  }
+
+  ///
+  /// Return whether a combination of ND-ranges is valid for fusion.
+  template <typename InputIt>
+  static bool isValidCombination(InputIt Begin, InputIt End) {
+    if (Begin == End) {
+      return false;
+    }
+    const auto &AllZeros = NDRange::AllZeros;
+    const auto FirstSpecLocal =
+        std::find_if(Begin, End, [&AllZeros](const auto &ND) {
+          return ND.getLocalSize() != AllZeros;
+        });
+    return std::all_of(Begin, End,
+                       [&ND = FirstSpecLocal == End ? *Begin : *FirstSpecLocal](
+                           const auto &Other) {
+                         return NDRange::compatibleRanges(ND, Other);
+                       });
+  }
+
+  ///
+  /// 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; }
+
+  friend constexpr bool operator==(const NDRange &LHS, const NDRange &RHS) {
+    return LHS.Dimensions == RHS.Dimensions &&
+           LHS.GlobalSize == RHS.GlobalSize &&
+           (LHS.LocalSize == AllZeros || RHS.LocalSize == AllZeros ||
+            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 +214,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/jit-compiler/CMakeLists.txt

@@ -5,6 +5,7 @@ add_llvm_library(sycl-fusion
    lib/translation/SPIRVLLVMTranslation.cpp
    lib/fusion/FusionPipeline.cpp
    lib/fusion/FusionHelper.cpp
+   lib/fusion/NDRangesHelper.cpp
    lib/fusion/ModuleHelper.cpp
    lib/helper/ConfigHelper.cpp
 

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

@@ -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/JITContext.h

@@ -29,7 +29,8 @@ 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>,
+               std::optional<std::vector<NDRange>>>;
 
 ///
 /// Wrapper around a SPIR-V binary.

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

@@ -0,0 +1,38 @@
+//==--------- 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_JIT_COMPILER_NDRANGESHELPER_H
+#define SYCL_FUSION_JIT_COMPILER_NDRANGESHELPER_H
+
+#include "Kernel.h"
+
+#include <algorithm>
+
+#include "llvm/ADT/ArrayRef.h"
+
+namespace jit_compiler {
+///
+/// Combine a list of ND-ranges, obtaining the resulting "fused" ND-range.
+NDRange combineNDRanges(llvm::ArrayRef<NDRange> NDRanges);
+
+///
+/// Returns whether the input list of ND-ranges is heterogeneous or not.
+inline bool isHeterogeneousList(llvm::ArrayRef<NDRange> NDRanges) {
+  const auto *Begin = NDRanges.begin();
+  const auto *End = NDRanges.end();
+  const auto *FirstSpecLocal =
+      std::find_if(Begin, End, [&AllZeros = NDRange::AllZeros](const auto &ND) {
+        return ND.getLocalSize() != AllZeros;
+      });
+  return std::any_of(Begin, End,
+                     [&ND = FirstSpecLocal == End ? *Begin : *FirstSpecLocal](
+                         const auto &Other) { return ND != Other; });
+}
+} // namespace jit_compiler
+
+#endif // SYCL_FUSION_JIT_COMPILER_NDRANGESHELPER_H

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

@@ -36,6 +36,16 @@ static FusionResult errorToFusionResult(llvm::Error &&Err,
   return FusionResult{ErrMsg.str()};
 }
 
+static std::vector<jit_compiler::NDRange>
+gatherNDRanges(llvm::ArrayRef<SYCLKernelInfo> KernelInformation) {
+  std::vector<jit_compiler::NDRange> NDRanges;
+  NDRanges.reserve(KernelInformation.size());
+  std::transform(KernelInformation.begin(), KernelInformation.end(),
+                 std::back_inserter(NDRanges),
+                 [](const auto &I) { return I.NDR; });
+  return NDRanges;
+}
+
 FusionResult KernelFusion::fuseKernels(
     JITContext &JITCtx, Config &&JITConfig,
     const std::vector<SYCLKernelInfo> &KernelInformation,
@@ -44,14 +54,21 @@ FusionResult KernelFusion::fuseKernels(
     int BarriersFlags,
     const std::vector<jit_compiler::ParameterInternalization> &Internalization,
     const std::vector<jit_compiler::JITConstant> &Constants) {
+  const auto NDRanges = gatherNDRanges(KernelInformation);
 
   // Initialize the configuration helper to make the options for this invocation
   // available (on a per-thread basis).
   ConfigHelper::setConfig(std::move(JITConfig));
 
   bool CachingEnabled = ConfigHelper::get<option::JITEnableCaching>();
-  CacheKeyT CacheKey{KernelsToFuse, Identities, BarriersFlags, Internalization,
-                     Constants};
+  CacheKeyT CacheKey{KernelsToFuse,
+                     Identities,
+                     BarriersFlags,
+                     Internalization,
+                     Constants,
+                     jit_compiler::isHeterogeneousList(NDRanges)
+                         ? std::optional<std::vector<NDRange>>{NDRanges}
+                         : std::optional<std::vector<NDRange>>{}};
   if (CachingEnabled) {
     std::optional<SYCLKernelInfo> CachedKernel = JITCtx.getCacheEntry(CacheKey);
     if (CachedKernel) {
@@ -82,8 +99,9 @@ FusionResult KernelFusion::fuseKernels(
 
   // Add information about the kernel that should be fused as metadata into the
   // LLVM module.
-  FusedFunction FusedKernel{FusedKernelName, KernelsToFuse,
-                            std::move(Identities), Internalization, Constants};
+  FusedFunction FusedKernel{
+      FusedKernelName, KernelsToFuse, std::move(Identities),
+      Internalization, Constants,     NDRanges};
   FusedFunctionList FusedKernelList;
   FusedKernelList.push_back(FusedKernel);
   llvm::Expected<std::unique_ptr<llvm::Module>> NewModOrError =
@@ -116,6 +134,8 @@ FusionResult KernelFusion::fuseKernels(
   }
   jit_compiler::SPIRVBinary *SPIRVBin = *BinaryOrError;
 
+  FusedKernelInfo.NDR = FusedKernel.FusedNDRange;
+
   // Update the KernelInfo for the fused kernel with the address and size of the
   // SPIR-V binary resulting from translation.
   SYCLKernelBinaryInfo &FusedBinaryInfo = FusedKernelInfo.BinaryInfo;

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

@@ -10,6 +10,7 @@
 
 #include "ModuleHelper.h"
 #include "helper/ErrorHandling.h"
+#include "kernel-fusion/SYCLKernelFusion.h"
 #include "llvm/IR/Constants.h"
 
 using namespace llvm;
@@ -81,6 +82,48 @@ Expected<std::unique_ptr<Module>> helper::FusionHelper::addFusedKernel(
     // The metadata can be identified by this fixed string providing a kind.
     F->setMetadata(MetadataKind, MDList);
 
+    // Attach ND-ranges related information. User of this API must pass the
+    // following information for each kernel, as well as for the fused kernel:
+    // 1. Number of dimensions;
+    // 2. Global size;
+    // 3. Local size;
+    // 4. Offset
+    {
+      const auto MDFromND = [&LLVMCtx](const auto &ND) {
+        auto MDFromIndices = [&LLVMCtx](const auto &Ind) -> Metadata * {
+          std::array<Metadata *, Ind.size()> MD{nullptr};
+          std::transform(
+              Ind.begin(), Ind.end(), MD.begin(),
+              [&LLVMCtx](auto I) { return getConstantIntMD(LLVMCtx, I); });
+          return MDNode::get(LLVMCtx, MD);
+        };
+        std::array<Metadata *, 4> MD;
+        MD[0] = getConstantIntMD(LLVMCtx, ND.getDimensions());
+        MD[1] = MDFromIndices(ND.getGlobalSize());
+        MD[2] = MDFromIndices(ND.getLocalSize());
+        MD[3] = MDFromIndices(ND.getOffset());
+        return MDNode::get(LLVMCtx, MD);
+      };
+
+      // Attach ND-range of the fused kernel
+      {
+        assert(!F->hasMetadata(SYCLKernelFusion::NDRangeMDKey));
+        F->setMetadata(SYCLKernelFusion::NDRangeMDKey,
+                       MDFromND(FF.FusedNDRange));
+      }
+
+      // Attach ND-ranges of each kernel to be fused
+      {
+        const auto SrcNDRanges = FF.NDRanges;
+        SmallVector<Metadata *> Nodes;
+        std::transform(SrcNDRanges.begin(), SrcNDRanges.end(),
+                       std::back_inserter(Nodes), MDFromND);
+        assert(!F->hasMetadata(SYCLKernelFusion::NDRangesMDKey));
+        F->setMetadata(SYCLKernelFusion::NDRangesMDKey,
+                       MDNode::get(LLVMCtx, Nodes));
+      }
+    }
+
     // The user of this API may be able to determine that
     // the same value is used for multiple input functions in the fused kernel,
     // e.g. when using the output of one kernel as the input to another kernel.

sycl-fusion/jit-compiler/lib/fusion/FusionHelper.h

@@ -9,6 +9,8 @@
 #ifndef SYCL_FUSION_JIT_COMPILER_FUSION_FUSIONHELPER_H
 #define SYCL_FUSION_JIT_COMPILER_FUSION_FUSIONHELPER_H
 
+#include "Kernel.h"
+#include "NDRangesHelper.h"
 #include "Parameter.h"
 #include "llvm/IR/Module.h"
 #include <llvm/ADT/ArrayRef.h>
@@ -25,12 +27,27 @@ class FusionHelper {
   /// Representation of a fused kernel named FusedName and fusing
   /// all the kernels listed in FusedKernels.
   struct FusedFunction {
+    FusedFunction(const std::string &FusedName,
+                  const std::vector<std::string> &FusedKernels,
+                  const jit_compiler::ParamIdentList &ParameterIdentities,
+                  llvm::ArrayRef<jit_compiler::ParameterInternalization>
+                      ParameterInternalization,
+                  llvm::ArrayRef<jit_compiler::JITConstant> Constants,
+                  llvm::ArrayRef<jit_compiler::NDRange> NDRanges)
+        : FusedName{FusedName}, FusedKernels{FusedKernels},
+          ParameterIdentities{ParameterIdentities},
+          ParameterInternalization{ParameterInternalization},
+          NDRanges{NDRanges}, FusedNDRange{
+                                  jit_compiler::combineNDRanges(NDRanges)} {}
+
     std::string FusedName;
     std::vector<std::string> FusedKernels;
     jit_compiler::ParamIdentList ParameterIdentities;
     llvm::ArrayRef<jit_compiler::ParameterInternalization>
         ParameterInternalization;
     llvm::ArrayRef<jit_compiler::JITConstant> Constants;
+    llvm::ArrayRef<jit_compiler::NDRange> NDRanges;
+    jit_compiler::NDRange FusedNDRange;
   };
 
   ///