[SYCL] Add support for JIT-ing in AMD and NVIDIA backends (#14280)

This patch provides the following:
* support for JIT compilation of Nvidia and AMD kernels
This is guarded by `SYCL_JIT_KERNELS` environment variable. Target CPU
and features can be provided through environment variables
(`SYCL_JIT_TARGET_CPU` and `SYCL_JIT_TARGET_FEATURES` respectively),
otherwise default, per-backend, values will be chosen.

* caching of JIT-compiled kernels
The runtime maintains a map of available JIT-ed kernels, accessible
through a key, which is constructed from kernel's name and values of
specialization constant (if provided).

* materialization of specialization
Materialization is done through a `SYCLSpecConstMaterializer` pass that
receives the values of all specialization constants used by a kernel
(from `SYCLSpecConstDataInserter`) and then walks all the uses of
implicit kernel argument (`_arg__specialization_constants_buffer`),
representing emulated specialization constants, with concrete values,
turning them to de-facto compile time constants.

This PR extends the work done for kernel fusion and in a similar fashion
it requires embedding of IR (`-fsycl-embed-ir`) during the initial
compilation.

Author: jchlanda

sycl-fusion/jit-compiler/include/KernelFusion.h

@@ -19,19 +19,18 @@
 
 namespace jit_compiler {
 
-class FusionResult {
+class JITResult {
 public:
-  explicit FusionResult(const char *ErrorMessage)
-      : Type{FusionResultType::FAILED}, KernelInfo{},
-        ErrorMessage{ErrorMessage} {}
+  explicit JITResult(const char *ErrorMessage)
+      : Type{JITResultType::FAILED}, KernelInfo{}, ErrorMessage{ErrorMessage} {}
 
-  explicit FusionResult(const SYCLKernelInfo &KernelInfo, bool Cached = false)
-      : Type{(Cached) ? FusionResultType::CACHED : FusionResultType::NEW},
+  explicit JITResult(const SYCLKernelInfo &KernelInfo, bool Cached = false)
+      : Type{(Cached) ? JITResultType::CACHED : JITResultType::NEW},
         KernelInfo(KernelInfo), ErrorMessage{} {}
 
-  bool failed() const { return Type == FusionResultType::FAILED; }
+  bool failed() const { return Type == JITResultType::FAILED; }
 
-  bool cached() const { return Type == FusionResultType::CACHED; }
+  bool cached() const { return Type == JITResultType::CACHED; }
 
   const char *getErrorMessage() const {
     assert(failed() && "No error message present");
@@ -44,9 +43,9 @@ class FusionResult {
   }
 
 private:
-  enum class FusionResultType { FAILED, CACHED, NEW };
+  enum class JITResultType { FAILED, CACHED, NEW };
 
-  FusionResultType Type;
+  JITResultType Type;
   SYCLKernelInfo KernelInfo;
   sycl::detail::string ErrorMessage;
 };
@@ -56,12 +55,18 @@ extern "C" {
 #ifdef __clang__
 #pragma clang diagnostic ignored "-Wreturn-type-c-linkage"
 #endif // __clang__
-FusionResult fuseKernels(View<SYCLKernelInfo> KernelInformation,
-                         const char *FusedKernelName,
-                         View<ParameterIdentity> Identities,
-                         BarrierFlags BarriersFlags,
-                         View<ParameterInternalization> Internalization,
-                         View<jit_compiler::JITConstant> JITConstants);
+JITResult fuseKernels(View<SYCLKernelInfo> KernelInformation,
+                      const char *FusedKernelName,
+                      View<ParameterIdentity> Identities,
+                      BarrierFlags BarriersFlags,
+                      View<ParameterInternalization> Internalization,
+                      View<jit_compiler::JITConstant> JITConstants);
+
+JITResult materializeSpecConstants(const char *KernelName,
+                                   jit_compiler::SYCLKernelBinaryInfo &BinInfo,
+                                   View<unsigned char> SpecConstBlob,
+                                   const char *TargetCPU,
+                                   const char *TargetFeatures);
 
 /// Clear all previously set options.
 void resetJITConfiguration();

sycl-fusion/jit-compiler/ld-version-script.txt

@@ -2,6 +2,7 @@
   global:
     /* Export the library entry points */
     fuseKernels;
+    materializeSpecConstants;
     resetJITConfiguration;
     addToJITConfiguration;
 

sycl-fusion/jit-compiler/lib/KernelFusion.cpp

@@ -24,8 +24,8 @@ using namespace jit_compiler;
 using FusedFunction = helper::FusionHelper::FusedFunction;
 using FusedFunctionList = std::vector<FusedFunction>;
 
-static FusionResult errorToFusionResult(llvm::Error &&Err,
-                                        const std::string &Msg) {
+static JITResult errorToFusionResult(llvm::Error &&Err,
+                                     const std::string &Msg) {
   std::stringstream ErrMsg;
   ErrMsg << Msg << "\nDetailed information:\n";
   llvm::handleAllErrors(std::move(Err),
@@ -34,7 +34,7 @@ static FusionResult errorToFusionResult(llvm::Error &&Err,
                           // compiled without exception support.
                           ErrMsg << "\t" << StrErr.getMessage() << "\n";
                         });
-  return FusionResult{ErrMsg.str().c_str()};
+  return JITResult{ErrMsg.str().c_str()};
 }
 
 static std::vector<jit_compiler::NDRange>
@@ -70,11 +70,58 @@ static bool isTargetFormatSupported(BinaryFormat TargetFormat) {
   }
 }
 
-extern "C" FusionResult
-fuseKernels(View<SYCLKernelInfo> KernelInformation, const char *FusedKernelName,
-            View<ParameterIdentity> Identities, BarrierFlags BarriersFlags,
-            View<ParameterInternalization> Internalization,
-            View<jit_compiler::JITConstant> Constants) {
+extern "C" JITResult
+materializeSpecConstants(const char *KernelName,
+                         jit_compiler::SYCLKernelBinaryInfo &BinInfo,
+                         View<unsigned char> SpecConstBlob,
+                         const char *TargetCPU, const char *TargetFeatures) {
+  auto &JITCtx = JITContext::getInstance();
+
+  TargetInfo TargetInfo = ConfigHelper::get<option::JITTargetInfo>();
+  BinaryFormat TargetFormat = TargetInfo.getFormat();
+  if (TargetFormat != BinaryFormat::PTX &&
+      TargetFormat != BinaryFormat::AMDGCN) {
+    return JITResult("Output target format not supported by this build. "
+                     "Available targets are: PTX or AMDGCN.");
+  }
+
+  ::jit_compiler::SYCLKernelInfo KernelInfo{
+      KernelName, ::jit_compiler::SYCLArgumentDescriptor{},
+      ::jit_compiler::NDRange{}, BinInfo};
+  SYCLModuleInfo ModuleInfo;
+  ModuleInfo.kernels().insert(ModuleInfo.kernels().end(), KernelInfo);
+  // Load all input kernels from their respective modules into a single
+  // LLVM IR module.
+  llvm::Expected<std::unique_ptr<llvm::Module>> ModOrError =
+      translation::KernelTranslator::loadKernels(*JITCtx.getLLVMContext(),
+                                                 ModuleInfo.kernels());
+  if (auto Error = ModOrError.takeError()) {
+    return errorToFusionResult(std::move(Error), "Failed to load kernels");
+  }
+  std::unique_ptr<llvm::Module> NewMod = std::move(*ModOrError);
+  if (!fusion::FusionPipeline::runMaterializerPasses(
+          *NewMod, SpecConstBlob.to<llvm::ArrayRef>()) ||
+      !NewMod->getFunction(KernelName)) {
+    return JITResult{"Materializer passes should not fail"};
+  }
+
+  SYCLKernelInfo &MaterializerKernelInfo = *ModuleInfo.getKernelFor(KernelName);
+  if (auto Error = translation::KernelTranslator::translateKernel(
+          MaterializerKernelInfo, *NewMod, JITCtx, TargetFormat, TargetCPU,
+          TargetFeatures)) {
+    return errorToFusionResult(std::move(Error),
+                               "Translation to output format failed");
+  }
+
+  return JITResult{MaterializerKernelInfo};
+}
+
+extern "C" JITResult fuseKernels(View<SYCLKernelInfo> KernelInformation,
+                                 const char *FusedKernelName,
+                                 View<ParameterIdentity> Identities,
+                                 BarrierFlags BarriersFlags,
+                                 View<ParameterInternalization> Internalization,
+                                 View<jit_compiler::JITConstant> Constants) {
 
   std::vector<std::string> KernelsToFuse;
   llvm::transform(KernelInformation, std::back_inserter(KernelsToFuse),
@@ -93,8 +140,7 @@ fuseKernels(View<SYCLKernelInfo> KernelInformation, const char *FusedKernelName,
   }
 
   if (!isTargetFormatSupported(TargetFormat)) {
-    return FusionResult(
-        "Fusion output target format not supported by this build");
+    return JITResult("Fusion output target format not supported by this build");
   }
 
   auto &JITCtx = JITContext::getInstance();
@@ -117,7 +163,7 @@ fuseKernels(View<SYCLKernelInfo> KernelInformation, const char *FusedKernelName,
         // before returning the kernel info to the runtime.
         CachedKernel->NDR = FusedNDR->getNDR();
       }
-      return FusionResult{*CachedKernel, /*Cached*/ true};
+      return JITResult{*CachedKernel, /*Cached*/ true};
     }
     helper::printDebugMessage(
         "Compiling new kernel, no suitable cached kernel found");
@@ -165,13 +211,13 @@ fuseKernels(View<SYCLKernelInfo> KernelInformation, const char *FusedKernelName,
                                               BarriersFlags);
 
   if (!NewMod->getFunction(FusedKernelName)) {
-    return FusionResult{"Kernel fusion failed"};
+    return JITResult{"Kernel fusion failed"};
   }
 
   // Get the updated kernel info for the fused kernel and add the information to
   // the existing KernelInfo.
   if (!NewModInfo->hasKernelFor(FusedKernelName)) {
-    return FusionResult{"No KernelInfo for fused kernel"};
+    return JITResult{"No KernelInfo for fused kernel"};
   }
 
   SYCLKernelInfo &FusedKernelInfo = *NewModInfo->getKernelFor(FusedKernelName);
@@ -188,7 +234,7 @@ fuseKernels(View<SYCLKernelInfo> KernelInformation, const char *FusedKernelName,
     JITCtx.addCacheEntry(CacheKey, FusedKernelInfo);
   }
 
-  return FusionResult{FusedKernelInfo};
+  return JITResult{FusedKernelInfo};
 }
 
 extern "C" void resetJITConfiguration() { ConfigHelper::reset(); }

sycl-fusion/jit-compiler/lib/fusion/FusionPipeline.cpp

@@ -12,6 +12,7 @@
 #include "helper/ConfigHelper.h"
 #include "internalization/Internalization.h"
 #include "kernel-fusion/SYCLKernelFusion.h"
+#include "kernel-fusion/SYCLSpecConstMaterializer.h"
 #include "kernel-info/SYCLKernelInfo.h"
 #include "syclcp/SYCLCP.h"
 
@@ -141,3 +142,48 @@ FusionPipeline::runFusionPasses(Module &Mod, SYCLModuleInfo &InputInfo,
 
   return std::make_unique<SYCLModuleInfo>(std::move(*NewModInfo.ModuleInfo));
 }
+
+bool FusionPipeline::runMaterializerPasses(
+    llvm::Module &Mod, llvm::ArrayRef<unsigned char> SpecConstData) {
+  PassBuilder PB;
+  LoopAnalysisManager LAM;
+  FunctionAnalysisManager FAM;
+  CGSCCAnalysisManager CGAM;
+  ModuleAnalysisManager MAM;
+  PB.registerModuleAnalyses(MAM);
+  PB.registerCGSCCAnalyses(CGAM);
+  PB.registerFunctionAnalyses(FAM);
+  PB.registerLoopAnalyses(LAM);
+  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+
+  ModulePassManager MPM;
+  // Register inserter and materializer passes.
+  {
+    FunctionPassManager FPM;
+    MPM.addPass(SYCLSpecConstDataInserter{SpecConstData});
+    FPM.addPass(SYCLSpecConstMaterializer{});
+    MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+  }
+  // Add generic optimizations,
+  {
+    FunctionPassManager FPM;
+    MPM.addPass(AlwaysInlinerPass{});
+    FPM.addPass(SROAPass{SROAOptions::ModifyCFG});
+    FPM.addPass(SCCPPass{});
+    FPM.addPass(ADCEPass{});
+    FPM.addPass(EarlyCSEPass{/*UseMemorySSA*/ true});
+    MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+  }
+  // followed by unrolling.
+  {
+    FunctionPassManager FPM;
+    FPM.addPass(createFunctionToLoopPassAdaptor(IndVarSimplifyPass{}));
+    LoopUnrollOptions UnrollOptions;
+    FPM.addPass(LoopUnrollPass{UnrollOptions});
+    MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+  }
+
+  MPM.run(Mod, MAM);
+
+  return true;
+}

sycl-fusion/jit-compiler/lib/fusion/FusionPipeline.h

@@ -27,6 +27,13 @@ class FusionPipeline {
   static std::unique_ptr<SYCLModuleInfo>
   runFusionPasses(llvm::Module &Mod, SYCLModuleInfo &InputInfo,
                   BarrierFlags BarriersFlags);
+
+  ///
+  /// Run the necessary passes in a custom pass pipeline to perform
+  /// materialization of kernel specialization constants.
+  static bool
+  runMaterializerPasses(llvm::Module &Mod,
+                        llvm::ArrayRef<unsigned char> SpecConstData);
 };
 } // namespace fusion
 } // namespace jit_compiler