[SYCL] Implement device image properties for virtual functions

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 &&