Embed required devicelib modules as device image

Signed-off-by: jinge90 <[email protected]>

Author: jinge90

llvm/include/llvm/SYCLLowerIR/ComputeModuleRuntimeInfo.h

@@ -12,6 +12,7 @@
 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/SYCLLowerIR/ModuleSplitter.h"
+#include "llvm/SYCLLowerIR/SYCLDeviceLibReqMask.h"
 #include "llvm/Support/PropertySetIO.h"
 #include <string>
 namespace llvm {

llvm/include/llvm/SYCLLowerIR/SYCLDeviceLibReqMask.h

@@ -17,6 +17,8 @@
 #pragma once
 
 #include <cstdint>
+#include <string>
+#include <vector>
 
 namespace llvm {
 
@@ -39,5 +41,7 @@ enum class DeviceLibExt : std::uint32_t {
 };
 
 uint32_t getSYCLDeviceLibReqMask(const Module &M);
-
+void getSYCLDeviceLibReqNames(unsigned int ReqMask,
+                              std::vector<std::string> &ReqNames);
+unsigned int getSYCLDeviceLibMeta(std::string &DeviceLibFn);
 } // namespace llvm

llvm/include/llvm/Support/PropertySetIO.h

@@ -200,6 +200,7 @@ class PropertySetRegistry {
   static constexpr char SYCL_SPEC_CONSTANTS_DEFAULT_VALUES[] =
       "SYCL/specialization constants default values";
   static constexpr char SYCL_DEVICELIB_REQ_MASK[] = "SYCL/devicelib req mask";
+  static constexpr char SYCL_DEVICELIB_METADATA[] = "SYCL/devicelib metadata";
   static constexpr char SYCL_KERNEL_PARAM_OPT_INFO[] = "SYCL/kernel param opt";
   static constexpr char SYCL_PROGRAM_METADATA[] = "SYCL/program metadata";
   static constexpr char SYCL_MISC_PROP[] = "SYCL/misc properties";

llvm/lib/SYCLLowerIR/ComputeModuleRuntimeInfo.cpp

@@ -160,6 +160,7 @@ PropSetRegTy computeModuleProperties(const Module &M,
     std::map<StringRef, uint32_t> RMEntry = {{"DeviceLibReqMask", MRMask}};
     PropSet.add(PropSetRegTy::SYCL_DEVICELIB_REQ_MASK, RMEntry);
   }
+
   {
     PropSet.add(PropSetRegTy::SYCL_DEVICE_REQUIREMENTS,
                 computeDeviceRequirements(M, EntryPoints).asMap());

llvm/lib/SYCLLowerIR/SYCLDeviceLibReqMask.cpp

@@ -757,6 +757,28 @@ uint32_t getDeviceLibBits(const std::string &FuncName) {
 
 } // namespace
 
+static std::unordered_map<DeviceLibExt, std::vector<const char *>>
+    SYCLDeviceLibFn = {
+        {DeviceLibExt::cl_intel_devicelib_assert,
+         {"libsycl-fallback-cassert.bc"}},
+        {DeviceLibExt::cl_intel_devicelib_math, {"libsycl-fallback-cmath.bc"}},
+        {DeviceLibExt::cl_intel_devicelib_math_fp64,
+         {"libsycl-fallback-cmath-fp64.bc"}},
+        {DeviceLibExt::cl_intel_devicelib_complex,
+         {"libsycl-fallback-complex.bc"}},
+        {DeviceLibExt::cl_intel_devicelib_complex_fp64,
+         {"libsycl-fallback-complex-fp64.bc"}},
+        {DeviceLibExt::cl_intel_devicelib_cstring,
+         {"libsycl-fallback-cstring.bc"}},
+        {DeviceLibExt::cl_intel_devicelib_imf, {"libsycl-fallback-imf.bc"}},
+        {DeviceLibExt::cl_intel_devicelib_imf_fp64,
+         {"libsycl-fallback-imf-fp64.bc"}},
+        {DeviceLibExt::cl_intel_devicelib_imf_bf16,
+         {"libsycl-fallback-imf-bf16.bc"}},
+        {DeviceLibExt::cl_intel_devicelib_bfloat16,
+         {"libsycl-fallback-bfloat16.bc", "libsycl-native-bfloat16.bc"}},
+};
+
 // For each device image module, we go through all functions which meets
 // 1. The function name has prefix "__devicelib_"
 // 2. The function is declaration which means it doesn't have function body
@@ -775,3 +797,55 @@ uint32_t llvm::getSYCLDeviceLibReqMask(const Module &M) {
   }
   return ReqMask;
 }
+
+void llvm::getSYCLDeviceLibReqNames(unsigned int ReqMask,
+                                    std::vector<std::string> &ReqNames) {
+  DeviceLibExt DeviceLibExts[] = {DeviceLibExt::cl_intel_devicelib_assert,
+                                  DeviceLibExt::cl_intel_devicelib_math,
+                                  DeviceLibExt::cl_intel_devicelib_math_fp64,
+                                  DeviceLibExt::cl_intel_devicelib_complex,
+                                  DeviceLibExt::cl_intel_devicelib_complex_fp64,
+                                  DeviceLibExt::cl_intel_devicelib_cstring,
+                                  DeviceLibExt::cl_intel_devicelib_imf,
+                                  DeviceLibExt::cl_intel_devicelib_imf_fp64,
+                                  DeviceLibExt::cl_intel_devicelib_imf_bf16,
+                                  DeviceLibExt::cl_intel_devicelib_bfloat16};
+
+  unsigned int Temp;
+  for (auto Ext : DeviceLibExts) {
+    Temp =
+        0x1 << (static_cast<size_t>(Ext) -
+                static_cast<size_t>(DeviceLibExt::cl_intel_devicelib_assert));
+    if (Temp & ReqMask) {
+      for (auto Fn : SYCLDeviceLibFn[Ext])
+        ReqNames.push_back(std::string(Fn));
+    }
+  }
+}
+
+unsigned int llvm::getSYCLDeviceLibMeta(std::string &DeviceLibFn) {
+  DeviceLibExt DeviceLibExts[] = {DeviceLibExt::cl_intel_devicelib_assert,
+                                  DeviceLibExt::cl_intel_devicelib_math,
+                                  DeviceLibExt::cl_intel_devicelib_math_fp64,
+                                  DeviceLibExt::cl_intel_devicelib_complex,
+                                  DeviceLibExt::cl_intel_devicelib_complex_fp64,
+                                  DeviceLibExt::cl_intel_devicelib_cstring,
+                                  DeviceLibExt::cl_intel_devicelib_imf,
+                                  DeviceLibExt::cl_intel_devicelib_imf_fp64,
+                                  DeviceLibExt::cl_intel_devicelib_imf_bf16,
+                                  DeviceLibExt::cl_intel_devicelib_bfloat16};
+
+  unsigned int DeviceLibMeta = 0;
+  for (auto Ext : DeviceLibExts) {
+    for (auto Fn : SYCLDeviceLibFn[Ext]) {
+      if (DeviceLibFn == Fn) {
+        DeviceLibMeta = static_cast<unsigned int>(Ext);
+        if (DeviceLibFn == "libsycl-native-bfloat16.bc")
+          DeviceLibMeta |= 0x80000000ULL;
+        break;
+      }
+    }
+  }
+
+  return DeviceLibMeta;
+}

llvm/lib/Support/PropertySetIO.cpp

@@ -196,6 +196,7 @@ void PropertyValue::copy(const PropertyValue &P) {
 
 constexpr char PropertySetRegistry::SYCL_SPECIALIZATION_CONSTANTS[];
 constexpr char PropertySetRegistry::SYCL_DEVICELIB_REQ_MASK[];
+constexpr char PropertySetRegistry::SYCL_DEVICELIB_METADATA[];
 constexpr char PropertySetRegistry::SYCL_SPEC_CONSTANTS_DEFAULT_VALUES[];
 constexpr char PropertySetRegistry::SYCL_KERNEL_PARAM_OPT_INFO[];
 constexpr char PropertySetRegistry::SYCL_PROGRAM_METADATA[];

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

@@ -12,7 +12,6 @@
 // - module splitter to split a big input module into smaller ones
 // - specialization constant intrinsic transformation
 //===----------------------------------------------------------------------===//
-
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/AssumptionCache.h"
@@ -311,9 +310,20 @@ std::string saveModuleIR(Module &M, int I, StringRef Suff) {
 
 std::string saveModuleProperties(module_split::ModuleDesc &MD,
                                  const GlobalBinImageProps &GlobProps, int I,
-                                 StringRef Suff, StringRef Target = "") {
-  auto PropSet =
-      computeModuleProperties(MD.getModule(), MD.entries(), GlobProps);
+                                 StringRef Suff, StringRef Target = "",
+                                 bool IsDeviceLib = false) {
+  PropSetRegTy PropSet;
+
+  // For fallback devicelib module, no kernel included and no specialization
+  // constant used, skip regular Prop emit.
+  if (!IsDeviceLib)
+    PropSet = computeModuleProperties(MD.getModule(), MD.entries(), GlobProps);
+  else {
+    auto SYCLDeviceLibMeta = getSYCLDeviceLibMeta(MD.Name);
+    std::map<StringRef, unsigned int> RMEntry = {
+        {"DeviceLibMetaData", SYCLDeviceLibMeta}};
+    PropSet.add(PropSetRegTy::SYCL_DEVICELIB_METADATA, RMEntry);
+  }
 
   std::string NewSuff = Suff.str();
   if (!Target.empty()) {
@@ -421,17 +431,24 @@ void addTableRow(util::SimpleTable &Table,
 //   IR component saving is skipped, and this file name is recorded as such in
 //   the result.
 void saveModule(std::vector<std::unique_ptr<util::SimpleTable>> &OutTables,
-                module_split::ModuleDesc &MD, int I, StringRef IRFilename) {
+                module_split::ModuleDesc &MD, int I, StringRef IRFilename,
+                bool IsDeviceLib = false) {
   IrPropSymFilenameTriple BaseTriple;
   StringRef Suffix = getModuleSuffix(MD);
   MD.saveSplitInformationAsMetadata();
-  if (!IRFilename.empty()) {
-    // don't save IR, just record the filename
-    BaseTriple.Ir = IRFilename.str();
+  if (!IsDeviceLib) {
+    if (!IRFilename.empty()) {
+      // don't save IR, just record the filename
+      BaseTriple.Ir = IRFilename.str();
+    } else {
+      MD.cleanup();
+      BaseTriple.Ir = saveModuleIR(MD.getModule(), I, Suffix);
+    }
   } else {
-    MD.cleanup();
+    // For DeviceLib Modules, don't need to do clean up.
     BaseTriple.Ir = saveModuleIR(MD.getModule(), I, Suffix);
   }
+
   if (DoSymGen) {
     // save the names of the entry points - the symbol table
     BaseTriple.Sym = saveModuleSymbolTable(MD, I, Suffix);
@@ -445,13 +462,24 @@ void saveModule(std::vector<std::unique_ptr<util::SimpleTable>> &OutTables,
       GlobalBinImageProps Props = {EmitKernelParamInfo, EmitProgramMetadata,
                                    EmitExportedSymbols, EmitImportedSymbols,
                                    DeviceGlobals};
-      CopyTriple.Prop =
-          saveModuleProperties(MD, Props, I, Suffix, OutputFile.Target);
+      CopyTriple.Prop = saveModuleProperties(MD, Props, I, Suffix,
+                                             OutputFile.Target, IsDeviceLib);
     }
     addTableRow(*Table, CopyTriple);
   }
 }
 
+void saveDeviceLibModule(
+    std::vector<std::unique_ptr<util::SimpleTable>> &OutTables,
+    const std::string &IRFile, int I, LLVMContext &Context) {
+  SMDiagnostic Err;
+  StringRef DeviceLibLoc = DeviceLibDir;
+  std::string IRPath = DeviceLibLoc.str() + "/" + IRFile;
+  std::unique_ptr<Module> IRModule = parseIRFile(IRPath, Err, Context);
+  llvm::module_split::ModuleDesc LibIRMD(std::move(IRModule), IRFile);
+  saveModule(OutTables, LibIRMD, I, IRFile, true);
+}
+
 module_split::ModuleDesc link(module_split::ModuleDesc &&MD1,
                               module_split::ModuleDesc &&MD2) {
   std::vector<std::string> Names;
@@ -747,7 +775,7 @@ bool isTargetCompatibleWithModule(const std::string &Target,
 }
 
 std::vector<std::unique_ptr<util::SimpleTable>>
-processInputModule(std::unique_ptr<Module> M) {
+processInputModule(std::unique_ptr<Module> M, LLVMContext &Context) {
   // Construct the resulting table which will accumulate all the outputs.
   SmallVector<StringRef, MAX_COLUMNS_IN_FILE_TABLE> ColumnTitles{
       StringRef(COL_CODE)};
@@ -776,6 +804,9 @@ processInputModule(std::unique_ptr<Module> M) {
   // if none were made.
   bool Modified = false;
 
+  // Keeps track of required device libraries by all device images.
+  unsigned int DeviceLibReqMask = 0;
+
   // Propagate ESIMD attribute to wrapper functions to prevent
   // spurious splits and kernel link errors.
   Modified |= runModulePass<SYCLFixupESIMDKernelWrapperMDPass>(*M);
@@ -887,6 +918,7 @@ processInputModule(std::unique_ptr<Module> M) {
                 "have been made\n";
     }
     for (module_split::ModuleDesc &IrMD : MMs) {
+      DeviceLibReqMask |= getSYCLDeviceLibReqMask(IrMD.getModule());
       saveModule(Tables, IrMD, ID, OutIRFileName);
     }
 
@@ -895,12 +927,22 @@ processInputModule(std::unique_ptr<Module> M) {
     if (!MMsWithDefaultSpecConsts.empty()) {
       for (size_t i = 0; i != MMsWithDefaultSpecConsts.size(); ++i) {
         module_split::ModuleDesc &IrMD = MMsWithDefaultSpecConsts[i];
+        DeviceLibReqMask |= getSYCLDeviceLibReqMask(IrMD.getModule());
         saveModule(Tables, IrMD, ID, OutIRFileName);
       }
 
       ++ID;
     }
   }
+
+  if ((DeviceLibReqMask > 0) && (DeviceLibDir.getNumOccurrences() > 0)) {
+    string_vector DeviceLibReqNames;
+    getSYCLDeviceLibReqNames(DeviceLibReqMask, DeviceLibReqNames);
+    for (auto Fn : DeviceLibReqNames) {
+      saveDeviceLibModule(Tables, Fn, ID, Context);
+      ++ID;
+    }
+  }
   return Tables;
 }
 
@@ -1044,7 +1086,7 @@ int main(int argc, char **argv) {
   }
 
   std::vector<std::unique_ptr<util::SimpleTable>> Tables =
-      processInputModule(std::move(M));
+      processInputModule(std::move(M), Context);
 
   // Input module was processed and a single output file was requested.
   if (IROutputOnly)