[SYCL] Implement device image properties for virtual functions (#14875)

Implementation design explaining those changes in a bigger picture can
be found in #10540

Key things implemented here:
- device code split to outline virtual functions into separate device
images
- emission of new properties for virtual functions
- generation of `calls-indirectly` LLVM IR attribute for kernels that
construct objects with virtual functions, but don't do calls
- device image manipulations to cleanup or preserve virtual functions
depending on a device image

Even though those pieces are technically independent from each other, it
is hard to split them apart into separate PRs, because they all have to
be either present or absent for existing E2E tests for virtual functions
to work.

Author: AlexeySachkov

llvm/include/llvm/SYCLLowerIR/ModuleSplitter.h

@@ -82,6 +82,11 @@ struct EntryPointGroup {
       // Scope remains global
       return Res;
     }
+
+    // Indicates that this group holds definitions of virtual functions - they
+    // are outlined into separate device images and should be removed from all
+    // other modules. The flag is used in ModuleDesc::cleanup
+    bool HasVirtualFunctionDefinitions = false;
   };
 
   std::string GroupId;

llvm/include/llvm/Support/PropertySetIO.h

@@ -209,6 +209,7 @@ class PropertySetRegistry {
   static constexpr char SYCL_DEVICE_GLOBALS[] = "SYCL/device globals";
   static constexpr char SYCL_DEVICE_REQUIREMENTS[] = "SYCL/device requirements";
   static constexpr char SYCL_HOST_PIPES[] = "SYCL/host pipes";
+  static constexpr char SYCL_VIRTUAL_FUNCTIONS[] = "SYCL/virtual functions";
 
   /// Function for bulk addition of an entire property set in the given
   /// \p Category .

llvm/lib/SYCLLowerIR/ComputeModuleRuntimeInfo.cpp

@@ -8,6 +8,9 @@
 // See comments in the header.
 //===----------------------------------------------------------------------===//
 #include "llvm/SYCLLowerIR/ComputeModuleRuntimeInfo.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Demangle/Demangle.h"
 #include "llvm/IR/PassInstrumentation.h"
 #include "llvm/SYCLLowerIR/CompileTimePropertiesPass.h"
@@ -188,6 +191,10 @@ PropSetRegTy computeModuleProperties(const Module &M,
   if (GlobProps.EmitExportedSymbols) {
     // extract exported functions if any and save them into property set
     for (const auto *F : EntryPoints) {
+      // Virtual functions use a different mechanism of dynamic linking, they
+      // should not be registered here.
+      if (F->hasFnAttribute("indirectly-callable"))
+        continue;
       // TODO FIXME some of SYCL/ESIMD functions maybe marked with __regcall CC,
       // so they won't make it into the export list. Should the check be
       // F->getCallingConv() != CallingConv::SPIR_KERNEL?
@@ -201,11 +208,19 @@ PropSetRegTy computeModuleProperties(const Module &M,
   if (GlobProps.EmitImportedSymbols) {
     // record imported functions in the property set
     for (const auto &F : M) {
-      if ( // A function that can be imported may still be defined in one split
-           // image. Only add import property if this is not the image where the
-           // function is defined.
-          F.isDeclaration() && module_split::canBeImportedFunction(F)) {
+      // A function that can be imported may still be defined in one split
+      // image. Only add import property if this is not the image where the
+      // function is defined.
+      if (!F.isDeclaration())
+        continue;
 
+      // Even though virtual functions are considered to be imported by the
+      // function below, we shouldn't list them in the property because they
+      // use different mechanism for dynamic linking.
+      if (F.hasFnAttribute("indirectly-callable"))
+        continue;
+
+      if (module_split::canBeImportedFunction(F)) {
         // StripDeadPrototypes is called during module splitting
         // cleanup.  At this point all function decls should have uses.
         assert(!F.use_empty() && "Function F has no uses");
@@ -354,6 +369,50 @@ PropSetRegTy computeModuleProperties(const Module &M,
     PropSet.add(PropSetRegTy::SYCL_MISC_PROP, "specConstsReplacedWithDefault",
                 1);
 
+  { // Properties related to virtual functions
+    StringSet<> UsedVFSets;
+    bool AddedVFSetProperty = false;
+    for (const Function &F : M) {
+      if (F.isDeclaration())
+        continue;
+
+      if (F.hasFnAttribute("indirectly-callable")) {
+        PropSet.add(PropSetRegTy::SYCL_VIRTUAL_FUNCTIONS,
+                    "virtual-functions-set",
+                    F.getFnAttribute("indirectly-callable").getValueAsString());
+        AddedVFSetProperty = true;
+        // Device code split should ensure that virtual functions that belong
+        // to different sets are split into separate device images and hence
+        // there is no need to scan other functions.
+        break;
+      }
+
+      if (F.hasFnAttribute("calls-indirectly")) {
+        SmallVector<StringRef, 4> Sets;
+        F.getFnAttribute("calls-indirectly")
+            .getValueAsString()
+            .split(Sets, ',', /* MaxSplits */ -1, /* KeepEmpty */ false);
+        for (auto Set : Sets)
+          UsedVFSets.insert(Set);
+      }
+    }
+
+    if (!UsedVFSets.empty()) {
+      assert(!AddedVFSetProperty &&
+             "device image cannot have both virtual-functions-set and "
+             "uses-virtual-functions-set property");
+      SmallString<128> AllSets;
+      for (auto &It : UsedVFSets) {
+        if (!AllSets.empty())
+          AllSets += ',';
+        AllSets += It.getKey();
+      }
+
+      PropSet.add(PropSetRegTy::SYCL_VIRTUAL_FUNCTIONS,
+                   "uses-virtual-functions-set", AllSets);
+    }
+  }
+
   return PropSet;
 }
 std::string computeModuleSymbolTable(const Module &M,

llvm/lib/SYCLLowerIR/ModuleSplitter.cpp

@@ -136,7 +136,10 @@ bool isEntryPoint(const Function &F, bool EmitOnlyKernelsAsEntryPoints) {
            !isGenericBuiltin(F.getName());
   }
 
-  return false;
+  // Even if we are emitting only kernels as entry points, virtual functions
+  // should still be treated as entry points, because they are going to be
+  // outlined into separate device images and linked in later.
+  return F.hasFnAttribute("indirectly-callable");
 }
 
 // Represents "dependency" or "use" graph of global objects (functions and
@@ -668,6 +671,22 @@ bool mustPreserveGV(const GlobalValue &GV) {
 // TODO: try to move all passes (cleanup, spec consts, compile time properties)
 // in one place and execute MPM.run() only once.
 void ModuleDesc::cleanup() {
+  // Any definitions of virtual functions should be removed and turned into
+  // declarations, they are supposed to be provided by a different module.
+  if (!EntryPoints.Props.HasVirtualFunctionDefinitions) {
+    for (Function &F : *M)
+      if (F.hasFnAttribute("indirectly-callable")) {
+        F.deleteBody();
+        if (F.hasComdat())
+          F.setComdat(nullptr);
+      }
+  } else {
+    // Otherwise externalize them so they are not dropped by GlobalDCE
+    for (Function &F : *M)
+      if (F.hasFnAttribute("indirectly-callable"))
+        F.setLinkage(GlobalValue::LinkageTypes::ExternalLinkage);
+  }
+
   ModuleAnalysisManager MAM;
   MAM.registerPass([&] { return PassInstrumentationAnalysis(); });
   ModulePassManager MPM;
@@ -1057,6 +1076,17 @@ getDeviceCodeSplitter(ModuleDesc &&MD, IRSplitMode Mode, bool IROutputOnly,
     Categorizer.registerSimpleStringAttributeRule(
         sycl::utils::ATTR_SYCL_MODULE_ID);
 
+    // This attribute marks virtual functions and effectively dictates how they
+    // should be groupped together. By design we won't split those groups of
+    // virtual functions further even if functions from the same group use
+    // different optional features and therefore this rule is put here.
+    // Strictly speaking, we don't even care about module-id splitting for
+    // those, but to avoid that we need to refactor the whole categorizer.
+    // However, this is good enough as it is for an initial version.
+    // TODO: for AOT use case we shouldn't be outlining those and instead should
+    // only select those functions which are compatible with the target device
+    Categorizer.registerSimpleStringAttributeRule("indirectly-callable");
+
     // Optional features
     // Note: Add more rules at the end of the list to avoid chaning orders of
     // output files in existing tests.
@@ -1096,8 +1126,19 @@ getDeviceCodeSplitter(ModuleDesc &&MD, IRSplitMode Mode, bool IROutputOnly,
     Groups.reserve(EntryPointsMap.size());
     // Start with properties of a source module
     EntryPointGroup::Properties MDProps = MD.getEntryPointGroup().Props;
-    for (auto &[Key, EntryPoints] : EntryPointsMap)
-      Groups.emplace_back(Key, std::move(EntryPoints), MDProps);
+    for (auto &[Key, EntryPoints] : EntryPointsMap) {
+      bool HasVirtualFunctions = false;
+      for (auto *F : EntryPoints) {
+        if (F->hasFnAttribute("indirectly-callable")) {
+          HasVirtualFunctions = true;
+          break;
+        }
+      }
+
+      auto PropsCopy = MDProps;
+      PropsCopy.HasVirtualFunctionDefinitions = HasVirtualFunctions;
+      Groups.emplace_back(Key, std::move(EntryPoints), PropsCopy);
+    }
   }
 
   bool DoSplit = (Mode != SPLIT_NONE &&