[SYCL] Add option to sycl-post-link to embed module properties/symbols as metadata

Signed-off-by: Sarnie, Nick <[email protected]>

Author: sarnex

clang/lib/Driver/ToolChains/Clang.cpp

@@ -10717,6 +10717,10 @@ getTripleBasedSYCLPostLinkOpts(const ToolChain &TC, const JobAction &JA,
       (IsAOT || NewOffloadDriver))
     addArgs(PostLinkArgs, TCArgs,
             {"-generate-device-image-default-spec-consts"});
+  bool IsUsingLTO = TC.getDriver().isUsingLTO(/*IsDeviceOffloadAction=*/true);
+  auto LTOMode = TC.getDriver().getLTOMode(/*IsDeviceOffloadAction=*/true);
+  if (IsUsingLTO && LTOMode == LTOK_Thin)
+    addArgs(PostLinkArgs, TCArgs, {"-embed-aux-info-as-metadata"});
 }
 
 // sycl-post-link tool normally outputs a file table (see the tool sources for

clang/test/Driver/sycl-lto.cpp

@@ -7,3 +7,4 @@
 // Verify there's no error and we see the expected cc1 flags with the new offload driver.
 // RUN: %clangxx -fsycl --offload-new-driver -foffload-lto=thin %s -### 2>&1 | FileCheck -check-prefix=CHECK_SUPPORTED %s
 // CHECK_SUPPORTED: clang{{.*}} "-cc1" "-triple" "spir64-unknown-unknown" {{.*}} "-flto=thin" "-flto-unit"
+// CHECK_SUPPORTED: sycl-post-link{{.*}} -embed-aux-info-as-metadata

llvm/test/tools/sycl-post-link/embed_aux_info_as_metadata.ll

@@ -0,0 +1,37 @@
+; This test verifies the -embed-aux-info-as-metadata option.
+; In particular, we should see the properties and symbols file are not added to the output table
+; and are instead embedded in the module as metadata.
+;
+; RUN: sycl-post-link -split=source -embed-aux-info-as-metadata -symbols -S < %s -o %t.table
+; RUN: FileCheck %s -input-file=%t.table -check-prefix=CHECK-TABLE
+; RUN: FileCheck %s -input-file=%t_0.ll -check-prefix=CHECK-BAR
+; RUN: FileCheck %s -input-file=%t_1.ll -check-prefix=CHECK-FOO
+
+; CHECK-TABLE: [Code]
+; CHECK-TABLE: {{.*}}_0.ll
+; CHECK-TABLE: {{.*}}_1.ll
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
+target triple = "spir64-unknown-unknown"
+
+; CHECK-FOO: define dso_local spir_func noundef void @foo
+define dso_local spir_func noundef void @foo(i32 noundef %a, i32 noundef %b) local_unnamed_addr #0 {
+entry:
+ret void
+}
+
+define dso_local spir_func noundef void @bar(i32 noundef %a, i32 noundef %b) #1 {
+entry:
+ret void
+}
+attributes #0 = { mustprogress nofree norecurse nosync nounwind willreturn memory(none) "sycl-module-id"="test.cpp" "sycl-grf-size"="128" }
+attributes #1 = { convergent mustprogress noinline norecurse nounwind "sycl-module-id"="test.cpp" "sycl-grf-size"="256" }
+; CHECK-FOO: !sycl_properties = !{![[#FOO_PROP:]]}
+; CHECK-FOO: !sycl_symbol_table = !{![[#FOO_SYM:]]}
+; CHECK-FOO: ![[#FOO_PROP]] = !{!"[SYCL/devicelib req mask]\0ADeviceLibReqMask=1|0\0A[SYCL/device requirements]\0Aaspects=2|AAAAAAAAAAA\0A[SYCL/misc properties]\0Asycl-grf-size=1|128\0A"}
+; CHECK-FOO: ![[#FOO_SYM]] = !{!"foo\0A"}
+
+; CHECK-BAR: !sycl_properties = !{![[#BAR_PROP:]]}
+; CHECK-BAR: !sycl_symbol_table = !{![[#BAR_SYM:]]}
+; CHECK-BAR: ![[#BAR_PROP]] = !{!"[SYCL/devicelib req mask]\0ADeviceLibReqMask=1|0\0A[SYCL/device requirements]\0Aaspects=2|AAAAAAAAAAA\0A[SYCL/misc properties]\0Asycl-grf-size=1|256\0A"}
+; CHECK-BAR: ![[#BAR_SYM]] = !{!"bar\0A"}

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

@@ -75,6 +75,7 @@
 using namespace llvm;
 
 using string_vector = std::vector<std::string>;
+using PropSetRegTy = llvm::util::PropertySetRegistry;
 
 namespace {
 
@@ -251,6 +252,11 @@ cl::opt<bool> GenerateDeviceImageWithDefaultSpecConsts{
              "replaced with default values from specialization id(s)."),
     cl::cat(PostLinkCat)};
 
+cl::opt<bool> EmbedAuxillaryInfoAsMetadata{
+    "embed-aux-info-as-metadata",
+    cl::desc("Embed the module properties and symbols as metadata"),
+    cl::cat(PostLinkCat)};
+
 struct GlobalBinImageProps {
   bool EmitKernelParamInfo;
   bool EmitProgramMetadata;
@@ -393,10 +399,31 @@ std::string makeResultFileName(Twine Ext, int I, StringRef Suffix) {
       .str();
 }
 
-void saveModuleIR(Module &M, StringRef OutFilename) {
+// Add the module properties and symbol table as metadata. This is used when we
+// compile with thinLTO, as sycl-post-link is run early.
+void addAuxInfoAsMetadata(Module &M, const PropSetRegTy &PropSet,
+                          const std::string &SymT) {
+  auto AddMD = [&M](StringRef MDName, const std::string &MDVal) {
+    auto NamedMD = M.getOrInsertNamedMetadata(MDName);
+    auto MDStr = MDString::get(M.getContext(), MDVal);
+    SmallVector<Metadata *, 2> MD{MDStr};
+    auto MDOp = MDNode::get(M.getContext(), MD);
+    NamedMD->addOperand(MDOp);
+  };
+  std::string Buf;
+  raw_string_ostream OStr(Buf);
+  PropSet.write(OStr);
+  AddMD("sycl_properties", Buf);
+  AddMD("sycl_symbol_table", SymT);
+}
+
+void saveModuleIR(Module &M, StringRef OutFilename, const PropSetRegTy &PropSet,
+                  const std::string &SymT) {
   std::error_code EC;
   raw_fd_ostream Out{OutFilename, EC, sys::fs::OF_None};
   checkError(EC, "error opening the file '" + OutFilename + "'");
+  if (EmbedAuxillaryInfoAsMetadata)
+    addAuxInfoAsMetadata(M, PropSet, SymT);
 
   ModulePassManager MPM;
   ModuleAnalysisManager MAM;
@@ -409,11 +436,12 @@ void saveModuleIR(Module &M, StringRef OutFilename) {
   MPM.run(M, MAM);
 }
 
-std::string saveModuleIR(Module &M, int I, StringRef Suff) {
+std::string saveModuleIR(Module &M, int I, StringRef Suff,
+                         const PropSetRegTy &PropSet, const std::string &SymT) {
   DUMP_ENTRY_POINTS(M, EmitOnlyKernelsAsEntryPoints, "saving IR");
   StringRef FileExt = (OutputAssembly) ? ".ll" : ".bc";
   std::string OutFilename = makeResultFileName(FileExt, I, Suff);
-  saveModuleIR(M, OutFilename);
+  saveModuleIR(M, OutFilename, PropSet, SymT);
   return OutFilename;
 }
 
@@ -436,10 +464,8 @@ bool isImportedFunction(const Function &F) {
   return ReturnValue;
 }
 
-std::string saveModuleProperties(module_split::ModuleDesc &MD,
-                                 const GlobalBinImageProps &GlobProps, int I,
-                                 StringRef Suff) {
-  using PropSetRegTy = llvm::util::PropertySetRegistry;
+PropSetRegTy getModuleProperties(module_split::ModuleDesc &MD,
+                                 const GlobalBinImageProps &GlobProps) {
   PropSetRegTy PropSet;
   Module &M = MD.getModule();
   {
@@ -646,6 +672,14 @@ std::string saveModuleProperties(module_split::ModuleDesc &MD,
     PropSet.add(PropSetRegTy::SYCL_MISC_PROP, "specConstsReplacedWithDefault",
                 1);
 
+  return PropSet;
+}
+
+std::string saveModuleProperties(module_split::ModuleDesc &MD,
+                                 const GlobalBinImageProps &GlobProps,
+                                 const PropSetRegTy &PropSet, int I,
+                                 StringRef Suff) {
+
   std::error_code EC;
   std::string SCFile = makeResultFileName(".prop", I, Suff);
   raw_fd_ostream SCOut(SCFile, EC);
@@ -655,9 +689,7 @@ std::string saveModuleProperties(module_split::ModuleDesc &MD,
   return SCFile;
 }
 
-// Saves specified collection of symbols to a file.
-std::string saveModuleSymbolTable(const module_split::EntryPointSet &Es, int I,
-                                  StringRef Suffix) {
+std::string getModuleSymbolTable(const module_split::EntryPointSet &Es) {
 #ifndef NDEBUG
   if (DebugPostLink > 0) {
     llvm::errs() << "ENTRY POINTS saving Sym table {\n";
@@ -673,6 +705,12 @@ std::string saveModuleSymbolTable(const module_split::EntryPointSet &Es, int I,
   for (const auto *F : Es) {
     SymT = (Twine(SymT) + Twine(F->getName()) + Twine("\n")).str();
   }
+  return SymT;
+}
+
+// Saves specified collection of symbols to a file.
+std::string saveModuleSymbolTable(const std::string &SymT, int I,
+                                  StringRef Suffix) {
   // Save to file.
   std::string OutFileName = makeResultFileName(".sym", I, Suffix);
   writeToFile(OutFileName, SymT);
@@ -757,22 +795,28 @@ IrPropSymFilenameTriple saveModule(module_split::ModuleDesc &MD, int I,
                                    StringRef IRFilename = "") {
   IrPropSymFilenameTriple Res;
   StringRef Suffix = getModuleSuffix(MD);
-
+  GlobalBinImageProps Props = {EmitKernelParamInfo, EmitProgramMetadata,
+                               EmitExportedSymbols, EmitImportedSymbols,
+                               DeviceGlobals};
+  PropSetRegTy PropSet = getModuleProperties(MD, Props);
+  std::string SymT = getModuleSymbolTable(MD.entries());
   if (!IRFilename.empty()) {
     // don't save IR, just record the filename
     Res.Ir = IRFilename.str();
   } else {
     MD.cleanup();
-    Res.Ir = saveModuleIR(MD.getModule(), I, Suffix);
+    Res.Ir = saveModuleIR(MD.getModule(), I, Suffix, PropSet, SymT);
   }
-  GlobalBinImageProps Props = {EmitKernelParamInfo, EmitProgramMetadata,
-                               EmitExportedSymbols, EmitImportedSymbols,
-                               DeviceGlobals};
-  Res.Prop = saveModuleProperties(MD, Props, I, Suffix);
-
-  if (DoSymGen) {
+  // We don't need a seperate file with properties if we saved it as
+  // metadata inside the module itself.
+  if (!EmbedAuxillaryInfoAsMetadata)
+    Res.Prop = saveModuleProperties(MD, Props, PropSet, I, Suffix);
+
+  // Only save the symbol table to a seperate module if it
+  // was requested and was not placed inside the module itself.
+  if (DoSymGen && !EmbedAuxillaryInfoAsMetadata) {
     // save the names of the entry points - the symbol table
-    Res.Sym = saveModuleSymbolTable(MD.entries(), I, Suffix);
+    Res.Sym = saveModuleSymbolTable(SymT, I, Suffix);
   }
   return Res;
 }
@@ -1071,10 +1115,12 @@ std::vector<std::unique_ptr<util::SimpleTable>>
 processInputModule(std::unique_ptr<Module> M) {
   // Construct the resulting table which will accumulate all the outputs.
   SmallVector<StringRef, MAX_COLUMNS_IN_FILE_TABLE> ColumnTitles{
-      StringRef(COL_CODE), StringRef(COL_PROPS)};
-
-  if (DoSymGen) {
-    ColumnTitles.push_back(COL_SYM);
+      StringRef(COL_CODE)};
+  if (!EmbedAuxillaryInfoAsMetadata) {
+    ColumnTitles.push_back(COL_PROPS);
+    if (DoSymGen) {
+      ColumnTitles.push_back(COL_SYM);
+    }
   }
   Expected<std::unique_ptr<util::SimpleTable>> TableE =
       util::SimpleTable::create(ColumnTitles);
@@ -1185,7 +1231,13 @@ processInputModule(std::unique_ptr<Module> M) {
               "' can't be used");
       }
       MMs.front().cleanup();
-      saveModuleIR(MMs.front().getModule(), OutputFiles[0].Filename);
+      const auto &ModuleProps = getModuleProperties(
+          MMs.front(),
+          {EmitKernelParamInfo, EmitProgramMetadata, EmitExportedSymbols,
+           EmitImportedSymbols, DeviceGlobals});
+      std::string SymT = getModuleSymbolTable(MMs.front().entries());
+      saveModuleIR(MMs.front().getModule(), OutputFiles[0].Filename,
+                   ModuleProps, SymT);
       return Tables;
     }
     // Empty IR file name directs saveModule to generate one and save IR to
@@ -1338,6 +1390,10 @@ int main(int argc, char **argv) {
     return 1;
   }
 
+  if (!DoSymGen && EmbedAuxillaryInfoAsMetadata)
+    errs() << "error: -" << EmbedAuxillaryInfoAsMetadata.ArgStr
+           << " can't be used without -" << DoSymGen.ArgStr << "\n";
+
   SMDiagnostic Err;
   std::unique_ptr<Module> M = parseIRFile(InputFilename, Err, Context);
   // It is OK to use raw pointer here as we control that it does not outlive M