[offload][SYCL] Add Module splitting by categories.

llvm/include/llvm/Transforms/Utils/SplitModuleByCategory.h

@@ -8,8 +8,8 @@
 // Functionality to split a module by categories.
 //===----------------------------------------------------------------------===//
 
-#ifndef LLVM_FRONTEND_SYCL_SPLIT_MODULE_H
-#define LLVM_FRONTEND_SYCL_SPLIT_MODULE_H
+#ifndef LLVM_TRANSFORM_UTILS_SPLIT_MODULE_BY_CATEGORY_H
+#define LLVM_TRANSFORM_UTILS_SPLIT_MODULE_BY_CATEGORY_H
 
 #include "llvm/ADT/STLFunctionalExtras.h"
 
@@ -22,21 +22,23 @@ namespace llvm {
 class Module;
 class Function;
 
-namespace sycl {
-
-/// FunctionCategorizer returns integer category for the given Function.
-/// Otherwise, it returns std::nullopt if function doesn't have a category.
-using FunctionCategorizer = function_ref<std::optional<int>(const Function &F)>;
-
-using PostSplitCallbackType = function_ref<void(std::unique_ptr<Module> Part)>;
-
-/// Splits the given module \p M.
-/// Every split image is being passed to \p Callback for further possible
+/// Splits the given module \p M using the given \p FunctionCategorizer.
+/// \p FunctionCategorizer returns integer category for an input Function.
+/// It may return std::nullopt if a function doesn't have a category.
+/// Module's functions are being grouped by categories. Every such group
+/// populates a call graph containing group's functions themselves and all
+/// reachable functions and globals. Split outputs are populated from each call
+/// graph associated with some category.
+///
+/// Every split output is being passed to \p Callback for further possible
 /// processing.
-void splitModuleByCategory(std::unique_ptr<Module> M, FunctionCategorizer FC,
-                           PostSplitCallbackType Callback);
+///
+/// Currently, the supported targets are SPIRV, AMDGPU and NVPTX.
+void splitModuleByCategory(
+    std::unique_ptr<Module> M,
+    function_ref<std::optional<int>(const Function &F)> FunctionCategorizer,
+    function_ref<void(std::unique_ptr<Module> Part)> Callback);
 
-} // namespace sycl
 } // namespace llvm
 
-#endif // LLVM_FRONTEND_SYCL_SPLIT_MODULE_H
+#endif // LLVM_TRANSFORM_UTILS_SPLIT_MODULE_BY_CATEGORY_H

llvm/lib/Transforms/Utils/CMakeLists.txt

@@ -84,7 +84,6 @@ add_llvm_component_library(LLVMTransformUtils
   SplitModule.cpp
   SplitModuleByCategory.cpp
   StripNonLineTableDebugInfo.cpp
-  SYCLUtils.cpp
   SymbolRewriter.cpp
   UnifyFunctionExitNodes.cpp
   UnifyLoopExits.cpp

llvm/lib/Transforms/Utils/SplitModuleByCategory.cpp

@@ -25,18 +25,11 @@
 #include <utility>
 
 using namespace llvm;
-using namespace llvm::sycl;
 
-#define DEBUG_TYPE "sycl-split-module"
+#define DEBUG_TYPE "split-module-by-category"
 
 namespace {
 
-bool isKernel(const Function &F) {
-  return F.getCallingConv() == CallingConv::SPIR_KERNEL ||
-         F.getCallingConv() == CallingConv::AMDGPU_KERNEL ||
-         F.getCallingConv() == CallingConv::PTX_Kernel;
-}
-
 // A vector that contains a group of function with the same category.
 using EntryPointSet = SetVector<const Function *>;
 
@@ -106,6 +99,12 @@ class ModuleDesc {
 #endif
 };
 
+bool isKernel(const Function &F) {
+  return F.getCallingConv() == CallingConv::SPIR_KERNEL ||
+         F.getCallingConv() == CallingConv::AMDGPU_KERNEL ||
+         F.getCallingConv() == CallingConv::PTX_Kernel;
+}
+
 // Represents "dependency" or "use" graph of global objects (functions and
 // global variables) in a module. It is used during device code split to
 // understand which global variables and functions (other than entry points)
@@ -299,17 +298,18 @@ class ModuleSplitter {
   bool hasMoreSplits() const { return Groups.size() > 0; }
 };
 
-EntryPointGroupVec selectEntryPointGroups(const Module &M,
-                                          FunctionCategorizer FC) {
+EntryPointGroupVec
+selectEntryPointGroups(const Module &M,
+                       function_ref<std::optional<int>(const Function &F)> FC) {
   // std::map is used here to ensure stable ordering of entry point groups,
   // which is based on their contents, this greatly helps LIT tests
   std::map<int, EntryPointSet> EntryPointsMap;
 
   for (const auto &F : M.functions()) {
-    if (auto Key = FC(F); Key) {
-      auto It = EntryPointsMap.find(*Key);
+    if (auto Category = FC(F); Category) {
+      auto It = EntryPointsMap.find(*Category);
       if (It == EntryPointsMap.end())
-        It = EntryPointsMap.emplace(*Key, EntryPointSet()).first;
+        It = EntryPointsMap.emplace(*Category, EntryPointSet()).first;
 
       It->second.insert(&F);
     }
@@ -326,10 +326,11 @@ EntryPointGroupVec selectEntryPointGroups(const Module &M,
 
 } // namespace
 
-void llvm::sycl::splitModuleByCategory(std::unique_ptr<Module> M,
-                                       FunctionCategorizer FC,
-                                       PostSplitCallbackType Callback) {
-  EntryPointGroupVec Groups = selectEntryPointGroups(*M, FC);
+void llvm::splitModuleByCategory(
+    std::unique_ptr<Module> M,
+    function_ref<std::optional<int>(const Function &F)> FunctionCategorizer,
+    function_ref<void(std::unique_ptr<Module> Part)> Callback) {
+  EntryPointGroupVec Groups = selectEntryPointGroups(*M, FunctionCategorizer);
   ModuleDesc MD = std::move(M);
   ModuleSplitter Splitter(std::move(MD), std::move(Groups));
   while (Splitter.hasMoreSplits()) {