[SYCL] Implement per-aspect device code split

llvm/test/tools/sycl-post-link/device-code-split/per-aspect-split-1.ll

@@ -0,0 +1,133 @@
+; This test emulates two translation units with 3 kernels:
+; TU0_kernel0 - 1st translation unit, no aspects used
+; TU0_kernel1 - 1st translation unit, aspect 1 is used
+; TU1_kernel2 - 2nd translation unit, no aspects used
+
+; The test is intended to check that sycl-post-link correctly separates kernels
+; that use aspects from kernels which doesn't use aspects regardless of device
+; code split mode
+
+; RUN: sycl-post-link -split=auto -symbols -S %s -o %t.table
+; RUN: FileCheck %s -input-file=%t_0.ll --check-prefixes CHECK-M0-IR \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_1.ll --check-prefixes CHECK-M1-IR \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_2.ll --check-prefixes CHECK-M2-IR \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_0.sym --check-prefixes CHECK-M0-SYMS \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_1.sym --check-prefixes CHECK-M1-SYMS \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_2.sym --check-prefixes CHECK-M2-SYMS \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+
+; RUN: sycl-post-link -split=source -symbols -S %s -o %t.table
+; RUN: FileCheck %s -input-file=%t_0.ll --check-prefixes CHECK-M0-IR \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_1.ll --check-prefixes CHECK-M1-IR \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_2.ll --check-prefixes CHECK-M2-IR \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_0.sym --check-prefixes CHECK-M0-SYMS \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_1.sym --check-prefixes CHECK-M1-SYMS \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_2.sym --check-prefixes CHECK-M2-SYMS \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+
+; RUN: sycl-post-link -split=kernel -symbols -S %s -o %t.table
+; RUN: FileCheck %s -input-file=%t_0.ll --check-prefixes CHECK-M0-IR \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_1.ll --check-prefixes CHECK-M1-IR \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_2.ll --check-prefixes CHECK-M2-IR \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_0.sym --check-prefixes CHECK-M0-SYMS \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_1.sym --check-prefixes CHECK-M1-SYMS \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+; RUN: FileCheck %s -input-file=%t_2.sym --check-prefixes CHECK-M2-SYMS \
+; RUN: --implicit-check-not kernel0 --implicit-check-not kernel1
+
+; Regardless of device code split mode, each kernel should go into a separate
+; device image
+
+; CHECK-M2-IR: define {{.*}} @TU0_kernel0
+; CHECK-M2-SYMS: TU0_kernel0
+
+; CHECK-M1-IR: define {{.*}} @TU0_kernel1
+; CHECK-M1-SYMS: TU0_kernel1
+
+; CHECK-M0-IR: define {{.*}} @TU1_kernel2
+; CHECK-M0-SYMS: TU1_kernel2
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
+target triple = "spir64-unknown-linux"
+
+; FIXME: device globals should also be properly distributed across device images
+; if they are of optional type
+@_ZL2GV = internal addrspace(1) constant [1 x i32] [i32 42], align 4
+
+define dso_local spir_kernel void @TU0_kernel0() #0 {
+entry:
+  call spir_func void @foo()
+  ret void
+}
+
+define dso_local spir_func void @foo() {
+entry:
+  %a = alloca i32, align 4
+  %call = call spir_func i32 @bar(i32 1)
+  %add = add nsw i32 2, %call
+  store i32 %add, i32* %a, align 4
+  ret void
+}
+
+; Function Attrs: nounwind
+define linkonce_odr dso_local spir_func i32 @bar(i32 %arg) {
+entry:
+  %arg.addr = alloca i32, align 4
+  store i32 %arg, i32* %arg.addr, align 4
+  %0 = load i32, i32* %arg.addr, align 4
+  ret i32 %0
+}
+
+define dso_local spir_kernel void @TU0_kernel1() #0 !sycl_used_aspects !2 {
+entry:
+  call spir_func void @foo1()
+  ret void
+}
+
+; Function Attrs: nounwind
+define dso_local spir_func void @foo1() {
+entry:
+  %a = alloca i32, align 4
+  store i32 2, i32* %a, align 4
+  ret void
+}
+
+define dso_local spir_kernel void @TU1_kernel2() #1 {
+entry:
+  call spir_func void @foo2()
+  ret void
+}
+
+; Function Attrs: nounwind
+define dso_local spir_func void @foo2() {
+entry:
+  %a = alloca i32, align 4
+  %0 = load i32, i32 addrspace(4)* getelementptr inbounds ([1 x i32], [1 x i32] addrspace(4)* addrspacecast ([1 x i32] addrspace(1)* @_ZL2GV to [1 x i32] addrspace(4)*), i64 0, i64 0), align 4
+  %add = add nsw i32 4, %0
+  store i32 %add, i32* %a, align 4
+  ret void
+}
+
+attributes #0 = { "sycl-module-id"="TU1.cpp" }
+attributes #1 = { "sycl-module-id"="TU2.cpp" }
+
+!opencl.spir.version = !{!0, !0}
+!spirv.Source = !{!1, !1}
+
+!0 = !{i32 1, i32 2}
+!1 = !{i32 4, i32 100000}
+!2 = !{i32 1}

llvm/test/tools/sycl-post-link/device-code-split/per-aspect-split-2.ll

@@ -0,0 +1,55 @@
+; The test is intended to check that sycl-post-link correctly groups kernels
+; by unique sets of aspects used in them
+
+; RUN: sycl-post-link -split=auto -symbols -S %s -o %t.table
+; RUN: FileCheck %s -input-file=%t.table --check-prefix CHECK-TABLE
+; RUN: FileCheck %s -input-file=%t_0.sym --check-prefix CHECK-M0-SYMS
+; RUN: FileCheck %s -input-file=%t_1.sym --check-prefix CHECK-M1-SYMS
+; RUN: FileCheck %s -input-file=%t_2.sym --check-prefix CHECK-M2-SYMS
+
+; FIXME: We expect to see exactly three modules
+; CHECK-TABLE: Code
+; CHECK-TABLE-NEXT: _0.sym
+; CHECK-TABLE-NEXT: _1.sym
+; CHECK-TABLE-NEXT: _2.sym
+; CHECK-TABLE-EMPTY:
+
+; CHECK-M0-SYMS: kernel3
+; CHECK-M0-SYMS-EMPTY:
+
+; CHECK-M1-SYMS: kernel0
+; CHECK-M1-SYMS-EMPTY:
+
+; CHECK-M2-SYMS: kernel1
+; CHECK-M2-SYMS: kernel2
+; CHECK-M2-SYMS-EMPTY:
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
+target triple = "spir64-unknown-linux"
+
+define dso_local spir_kernel void @kernel0() #0 !sycl_used_aspects !1 {
+entry:
+  ret void
+}
+
+define dso_local spir_kernel void @kernel1() #0 !sycl_used_aspects !2 {
+entry:
+  ret void
+}
+
+define dso_local spir_kernel void @kernel2() #0 !sycl_used_aspects !3 {
+entry:
+  ret void
+}
+
+define dso_local spir_kernel void @kernel3() #0 !sycl_used_aspects !4 {
+entry:
+  ret void
+}
+
+attributes #0 = { "sycl-module-id"="TU1.cpp" }
+
+!1 = !{i32 1}
+!2 = !{i32 1, i32 2}
+!3 = !{i32 2, i32 1}
+!4 = !{i32 2, i32 3, i32 4}

llvm/tools/sycl-post-link/ModuleSplitter.cpp

@@ -761,5 +761,146 @@ getDoubleGRFSplitter(ModuleDesc &&MD, bool EmitOnlyKernelsAsEntryPoints) {
     return std::make_unique<ModuleCopier>(std::move(MD), std::move(Groups));
 }
 
+namespace {
+  // Data structure, which represent a combination of all possible optional
+  // features used in a function.
+  //
+  // It has extra methods to be useable as a key in llvm::DenseMap.
+  struct UsedOptionalFeatures {
+    SmallVector<int, 4> Aspects;
+    // TODO: extend this further with reqd-sub-group-size, reqd-work-group-size,
+    // double-grf and other properties
+
+    UsedOptionalFeatures() = default;
+
+    UsedOptionalFeatures(const Function *F) {
+      if (const MDNode *MDN = F->getMetadata("sycl_used_aspects")) {
+        auto ExtractIntegerFromMDNodeOperand = [=](const MDNode *N,
+                                                   unsigned OpNo) -> auto {
+          Constant *C =
+              cast<ConstantAsMetadata>(N->getOperand(OpNo).get())->getValue();
+          return C->getUniqueInteger().getSExtValue();
+        };
+
+        for (size_t I = 0, E = MDN->getNumOperands(); I < E; ++I) {
+          Aspects.push_back(ExtractIntegerFromMDNodeOperand(MDN, I));
+        }
+        llvm::sort(Aspects);
+      }
+
+      llvm::hash_code AspectsHash =
+          llvm::hash_combine_range(Aspects.begin(), Aspects.end());
+      Hash = static_cast<unsigned>(llvm::hash_combine(AspectsHash));
+    }
+
+    std::string getName(StringRef BaseName) const {
+      if (Aspects.empty())
+        return BaseName.str() + "-no-aspects";
+
+      std::string Ret = BaseName.str() + "-aspects";
+      for (int A : Aspects) {
+        Ret += "-" + std::to_string(A);
+      }
+      return Ret;
+    }
+
+    static UsedOptionalFeatures getTombstone() {
+      UsedOptionalFeatures Ret;
+      Ret.IsTombstoneKey = true;
+      return Ret;
+    }
+
+    static UsedOptionalFeatures getEmpty() {
+      UsedOptionalFeatures Ret;
+      Ret.IsEmpty = true;
+      return Ret;
+    }
+
+  private:
+    // For DenseMap:
+    llvm::hash_code Hash = {};
+    bool IsTombstoneKey = false;
+    bool IsEmpty = false;
+
+  public:
+    bool operator==(const UsedOptionalFeatures &Other) const {
+      // Tombstone does not compare equal to any other item
+      if (IsTombstoneKey || Other.IsTombstoneKey)
+        return false;
+
+      if (Aspects.size() != Other.Aspects.size())
+        return false;
+
+      for (size_t I = 0, E = Aspects.size(); I != E; ++I) {
+        if (Aspects[I] != Other.Aspects[I])
+          return false;
+      }
+
+      return IsEmpty == Other.IsEmpty;
+    }
+
+    unsigned hash() const {
+      return static_cast<unsigned>(Hash);
+    }
+  };
+
+  struct UsedOptionalFeaturesAsKeyInfo {
+    static inline UsedOptionalFeatures getEmptyKey() {
+      return UsedOptionalFeatures::getEmpty();
+    }
+
+    static inline UsedOptionalFeatures getTombstoneKey() {
+      return UsedOptionalFeatures::getTombstone();
+    }
+
+    static unsigned getHashValue(const UsedOptionalFeatures &Value) {
+      return Value.hash();
+    }
+
+    static bool isEqual(const UsedOptionalFeatures &LHS,
+                        const UsedOptionalFeatures &RHS) {
+      return LHS == RHS;
+    }
+  };
+}
+
+std::unique_ptr<ModuleSplitterBase>
+getPropertiesBasedSplitter(ModuleDesc &&MD, bool EmitOnlyKernelsAsEntryPoints) {
+  EntryPointGroupVec Groups;
+
+  DenseMap<UsedOptionalFeatures, EntryPointSet, UsedOptionalFeaturesAsKeyInfo>
+      PropertiesToFunctionsMap;
+
+  Module &M = MD.getModule();
+
+  // Only process module entry points:
+  for (auto &F : M.functions()) {
+    if (!isEntryPoint(F, EmitOnlyKernelsAsEntryPoints) ||
+        !MD.isEntryPointCandidate(F)) {
+      continue;
+    }
+
+    auto Key = UsedOptionalFeatures(&F);
+    PropertiesToFunctionsMap[Key].insert(&F);
+  }
+
+  if (!PropertiesToFunctionsMap.empty()) {
+    Groups.reserve(PropertiesToFunctionsMap.size());
+    for (auto &EPG : PropertiesToFunctionsMap) {
+      Groups.emplace_back(EntryPointGroup{
+          EPG.first.getName(MD.getEntryPointGroup().GroupId),
+          std::move(EPG.second), MD.getEntryPointGroup().Props});
+    }
+  } else {
+    // No entry points met, record this.
+    Groups.push_back({GLOBAL_SCOPE_NAME, {}});
+  }
+
+  if (Groups.size() > 1)
+    return std::make_unique<ModuleSplitter>(std::move(MD), std::move(Groups));
+  else
+    return std::make_unique<ModuleCopier>(std::move(MD), std::move(Groups));
+}
+
 } // namespace module_split
 } // namespace llvm

llvm/tools/sycl-post-link/ModuleSplitter.h

@@ -253,6 +253,9 @@ getSplitterByMode(ModuleDesc &&MD, IRSplitMode Mode,
 std::unique_ptr<ModuleSplitterBase>
 getDoubleGRFSplitter(ModuleDesc &&MD, bool EmitOnlyKernelsAsEntryPoints);
 
+std::unique_ptr<ModuleSplitterBase>
+getPropertiesBasedSplitter(ModuleDesc &&MD, bool EmitOnlyKernelsAsEntryPoints);
+
 #ifndef NDEBUG
 void dumpEntryPoints(const EntryPointSet &C, const char *msg = "", int Tab = 0);
 void dumpEntryPoints(const Module &M, bool OnlyKernelsAreEntryPoints = false,