[SYCL][Fusion] Rebase and address feedback

Signed-off-by: Lukas Sommer <[email protected]>

Author: sommerlukas

clang/include/clang/Driver/Action.h

@@ -665,7 +665,7 @@ class OffloadWrapperJobAction : public JobAction {
 public:
   OffloadWrapperJobAction(ActionList &Inputs, types::ID Type);
   OffloadWrapperJobAction(Action *Input, types::ID OutputType,
-                          bool IsEmbeddedIR = false);
+                          bool EmbedIR = false);
 
   bool isEmbeddedIR() const { return EmbedIR; }
 

clang/lib/Driver/Action.cpp

@@ -478,7 +478,7 @@ void OffloadWrapperJobAction::anchor() {}
 
 OffloadWrapperJobAction::OffloadWrapperJobAction(ActionList &Inputs,
                                                  types::ID Type)
-  : JobAction(OffloadWrapperJobClass, Inputs, Type) {}
+    : JobAction(OffloadWrapperJobClass, Inputs, Type), EmbedIR(false) {}
 
 OffloadWrapperJobAction::OffloadWrapperJobAction(Action *Input, types::ID Type,
                                                  bool IsEmbeddedIR)

sycl-fusion/jit-compiler/CMakeLists.txt

@@ -22,9 +22,8 @@ add_llvm_library(sycl-fusion
    InstCombine
    Target
    TargetParser
-   NVPTX
-   X86
    MC
+   ${LLVM_TARGETS_TO_BUILD}
 )
 
 target_include_directories(sycl-fusion
@@ -47,6 +46,10 @@ target_link_libraries(sycl-fusion
   ${CMAKE_THREAD_LIBS_INIT}
 )
 
+if("NVPTX" IN_LIST LLVM_TARGETS_TO_BUILD)
+  target_compile_definitions(sycl-fusion PRIVATE FUSION_JIT_SUPPORT_PTX)
+endif()
+
 if (BUILD_SHARED_LIBS)
   if(NOT MSVC AND NOT APPLE)
     # Manage symbol visibility through the linker to make sure no LLVM symbols

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

@@ -39,7 +39,7 @@ using CacheKeyT =
 /// Wrapper around a kernel binary.
 class KernelBinary {
 public:
-  explicit KernelBinary(std::string Binary, BinaryFormat Format);
+  explicit KernelBinary(std::string &&Binary, BinaryFormat Format);
 
   jit_compiler::BinaryAddress address() const;
 
@@ -65,7 +65,10 @@ class JITContext {
 
   llvm::LLVMContext *getLLVMContext();
 
-  KernelBinary &emplaceSPIRVBinary(std::string Binary, BinaryFormat Format);
+  template <typename... Ts> KernelBinary &emplaceKernelBinary(Ts &&...Args) {
+    WriteLockT WriteLock{BinariesMutex};
+    return Binaries.emplace_back(std::forward<Ts>(Args)...);
+  }
 
   std::optional<SYCLKernelInfo> getCacheEntry(CacheKeyT &Identifier) const;
 

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

@@ -10,6 +10,7 @@
 #define SYCL_FUSION_JIT_COMPILER_OPTIONS_H
 
 #include "Kernel.h"
+
 #include <memory>
 #include <unordered_map>
 

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

@@ -11,7 +11,7 @@
 
 using namespace jit_compiler;
 
-KernelBinary::KernelBinary(std::string Binary, BinaryFormat Fmt)
+KernelBinary::KernelBinary(std::string &&Binary, BinaryFormat Fmt)
     : Blob{std::move(Binary)}, Format{Fmt} {}
 
 jit_compiler::BinaryAddress KernelBinary::address() const {
@@ -29,15 +29,6 @@ JITContext::~JITContext() = default;
 
 llvm::LLVMContext *JITContext::getLLVMContext() { return LLVMCtx.get(); }
 
-KernelBinary &JITContext::emplaceSPIRVBinary(std::string Binary,
-                                             BinaryFormat Format) {
-  WriteLockT WriteLock{BinariesMutex};
-  // NOTE: With C++17, which returns a reference from emplace_back, the
-  // following code would be even simpler.
-  Binaries.emplace_back(std::move(Binary), Format);
-  return Binaries.back();
-}
-
 std::optional<SYCLKernelInfo>
 JITContext::getCacheEntry(CacheKeyT &Identifier) const {
   ReadLockT ReadLock{CacheMutex};

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

@@ -48,6 +48,22 @@ gatherNDRanges(llvm::ArrayRef<SYCLKernelInfo> KernelInformation) {
   return NDRanges;
 }
 
+static bool isTargetFormatSupported(BinaryFormat TargetFormat) {
+  switch (TargetFormat) {
+  case BinaryFormat::SPIRV:
+    return true;
+  case BinaryFormat::PTX: {
+#ifdef FUSION_JIT_SUPPORT_PTX
+    return true;
+#else  // FUSION_JIT_SUPPORT_PTX
+    return false;
+#endif // FUSION_JIT_SUPPORT_PTX
+  }
+  default:
+    return false;
+  }
+}
+
 FusionResult KernelFusion::fuseKernels(
     JITContext &JITCtx, Config &&JITConfig,
     const std::vector<SYCLKernelInfo> &KernelInformation,
@@ -71,6 +87,12 @@ FusionResult KernelFusion::fuseKernels(
   bool IsHeterogeneousList = jit_compiler::isHeterogeneousList(NDRanges);
 
   BinaryFormat TargetFormat = ConfigHelper::get<option::JITTargetFormat>();
+
+  if (!isTargetFormatSupported(TargetFormat)) {
+    return FusionResult(
+        "Fusion output target format not supported by this build");
+  }
+
   if (TargetFormat == BinaryFormat::PTX && IsHeterogeneousList) {
     return FusionResult{"Heterogeneous ND ranges not supported for CUDA"};
   }

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

@@ -23,8 +23,8 @@
 #ifndef NDEBUG
 #include "llvm/IR/Verifier.h"
 #endif // NDEBUG
-#include "llvm/ADT/Triple.h"
 #include "llvm/Passes/PassBuilder.h"
+#include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/InstCombine/InstCombine.h"
 #include "llvm/Transforms/Scalar/ADCE.h"
 #include "llvm/Transforms/Scalar/EarlyCSE.h"
@@ -103,7 +103,7 @@ FusionPipeline::runFusionPasses(Module &Mod, SYCLModuleInfo &InputInfo,
     // to/from generic address-space as possible, because these hinder
     // internalization.
     // Ideally, the static compiler should have performed that job.
-    unsigned FlatAddressSpace = getFlatAddressSpace(Mod);
+    const unsigned FlatAddressSpace = getFlatAddressSpace(Mod);
     FPM.addPass(InferAddressSpacesPass(FlatAddressSpace));
     MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
   }

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

@@ -24,7 +24,7 @@ helper::ModuleHelper::cloneAndPruneModule(Module *Mod,
   identifyUnusedFunctions(Mod, CGRoots, UnusedFunctions);
 
   {
-    auto TFI = llvm::TargetFusionInfo::getTargetFusionInfo(Mod);
+    TargetFusionInfo TFI{Mod};
     SmallVector<Function *> Unused{UnusedFunctions.begin(),
                                    UnusedFunctions.end()};
     TFI.notifyFunctionsDelete(Unused);

sycl-fusion/jit-compiler/lib/translation/KernelTranslation.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "KernelTranslation.h"
+
 #include "SPIRVLLVMTranslation.h"
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/IR/Constants.h"
@@ -182,7 +183,8 @@ llvm::Error KernelTranslator::translateKernel(SYCLKernelInfo &Kernel,
     break;
   }
   case BinaryFormat::PTX: {
-    llvm::Expected<KernelBinary *> BinaryOrError = translateToPTX(Mod, JITCtx);
+    llvm::Expected<KernelBinary *> BinaryOrError =
+        translateToPTX(Kernel, Mod, JITCtx);
     if (auto Error = BinaryOrError.takeError()) {
       return Error;
     }
@@ -215,12 +217,20 @@ KernelTranslator::translateToSPIRV(llvm::Module &Mod, JITContext &JITCtx) {
 }
 
 llvm::Expected<KernelBinary *>
-KernelTranslator::translateToPTX(llvm::Module &Mod, JITContext &JITCtx) {
-  // FIXME: Can we limit this to the NVPTX specific target?
-  llvm::InitializeAllTargets();
-  llvm::InitializeAllAsmParsers();
-  llvm::InitializeAllAsmPrinters();
-  llvm::InitializeAllTargetMCs();
+KernelTranslator::translateToPTX(SYCLKernelInfo &KernelInfo, llvm::Module &Mod,
+                                 JITContext &JITCtx) {
+#ifndef FUSION_JIT_SUPPORT_PTX
+  return createStringError(inconvertibleErrorCode(),
+                           "PTX translation not supported in this build");
+#else  // FUSION_JIT_SUPPORT_PTX
+  LLVMInitializeNVPTXTargetInfo();
+  LLVMInitializeNVPTXTarget();
+  LLVMInitializeNVPTXAsmPrinter();
+  LLVMInitializeNVPTXTargetMC();
+#endif // FUSION_JIT_SUPPORT_PTX
+
+  static const char *TARGET_CPU_ATTRIBUTE = "target-cpu";
+  static const char *TARGET_FEATURE_ATTRIBUTE = "target-features";
 
   std::string TargetTriple{"nvptx64-nvidia-cuda"};
 
@@ -231,13 +241,26 @@ KernelTranslator::translateToPTX(llvm::Module &Mod, JITContext &JITCtx) {
   if (!Target) {
     return createStringError(
         inconvertibleErrorCode(),
-        "Failed to load and translate SPIR-V module with error %s",
+        "Failed to load and translate PTX LLVM IR module with error %s",
         ErrorMessage.c_str());
   }
 
+  llvm::StringRef TargetCPU{"sm_50"};
+  llvm::StringRef TargetFeatures{"+sm_50,+ptx76"};
+  if (auto *KernelFunc = Mod.getFunction(KernelInfo.Name)) {
+    if (KernelFunc->hasFnAttribute(TARGET_CPU_ATTRIBUTE)) {
+      TargetCPU =
+          KernelFunc->getFnAttribute(TARGET_CPU_ATTRIBUTE).getValueAsString();
+    }
+    if (KernelFunc->hasFnAttribute(TARGET_FEATURE_ATTRIBUTE)) {
+      TargetFeatures = KernelFunc->getFnAttribute(TARGET_FEATURE_ATTRIBUTE)
+                           .getValueAsString();
+    }
+  }
+
   // FIXME: Check whether we can provide more accurate target information here
   auto *TargetMachine = Target->createTargetMachine(
-      TargetTriple, "sm_50", "+sm_50,+ptx76", {}, llvm::Reloc::PIC_,
+      TargetTriple, TargetCPU, TargetFeatures, {}, llvm::Reloc::PIC_,
       std::nullopt, llvm::CodeGenOpt::Default);
 
   llvm::legacy::PassManager PM;
@@ -259,5 +282,5 @@ KernelTranslator::translateToPTX(llvm::Module &Mod, JITContext &JITCtx) {
     ASMStream.flush();
   }
 
-  return &JITCtx.emplaceSPIRVBinary(PTXASM, BinaryFormat::PTX);
+  return &JITCtx.emplaceKernelBinary(std::move(PTXASM), BinaryFormat::PTX);
 }

sycl-fusion/jit-compiler/lib/translation/KernelTranslation.h

@@ -5,12 +5,14 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+#ifndef SYCL_FUSION_JIT_COMPILER_TRANSLATION_KERNELTRANSLATION_H
+#define SYCL_FUSION_JIT_COMPILER_TRANSLATION_KERNELTRANSLATION_H
 
 #include "JITContext.h"
 #include "Kernel.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
-#include <llvm/Support/Error.h>
+#include "llvm/Support/Error.h"
 #include <vector>
 
 namespace jit_compiler {
@@ -39,8 +41,10 @@ class KernelTranslator {
   static llvm::Expected<KernelBinary *> translateToSPIRV(llvm::Module &Mod,
                                                          JITContext &JITCtx);
 
-  static llvm::Expected<KernelBinary *> translateToPTX(llvm::Module &Mod,
-                                                       JITContext &JITCtx);
+  static llvm::Expected<KernelBinary *>
+  translateToPTX(SYCLKernelInfo &Kernel, llvm::Module &Mod, JITContext &JITCtx);
 };
 } // namespace translation
 } // namespace jit_compiler
+
+#endif // SYCL_FUSION_JIT_COMPILER_TRANSLATION_KERNELTRANSLATION_H

sycl-fusion/jit-compiler/lib/translation/SPIRVLLVMTranslation.cpp

@@ -94,5 +94,5 @@ SPIRVLLVMTranslator::translateLLVMtoSPIRV(Module &Mod, JITContext &JITCtx) {
         "Translation of LLVM IR to SPIR-V failed with error %s",
         ErrMsg.c_str());
   }
-  return &JITCtx.emplaceSPIRVBinary(BinaryStream.str(), BinaryFormat::SPIRV);
+  return &JITCtx.emplaceKernelBinary(BinaryStream.str(), BinaryFormat::SPIRV);
 }

sycl-fusion/passes/internalization/Internalization.cpp

@@ -631,7 +631,7 @@ static void moduleCleanup(Module &M, ModuleAnalysisManager &AM,
 
 PreservedAnalyses llvm::SYCLInternalizer::run(Module &M,
                                               ModuleAnalysisManager &AM) {
-  auto TFI = TargetFusionInfo::getTargetFusionInfo(&M);
+  TargetFusionInfo TFI{&M};
   // Private promotion
   const PreservedAnalyses Tmp = SYCLInternalizerImpl{
       TFI.getPrivateAddressSpace(), PrivatePromotion, true, TFI}(M, AM);

sycl-fusion/passes/kernel-fusion/SYCLKernelFusion.cpp

@@ -138,7 +138,7 @@ PreservedAnalyses SYCLKernelFusion::run(Module &M, ModuleAnalysisManager &AM) {
       AM.getResult<SYCLModuleInfoAnalysis>(M).ModuleInfo;
   assert(ModuleInfo && "No module information available");
 
-  auto TFI = TargetFusionInfo::getTargetFusionInfo(&M);
+  TargetFusionInfo TFI{&M};
 
   // Iterate over the functions in the module and locate all
   // stub functions identified by metadata.
@@ -456,11 +456,20 @@ Error SYCLKernelFusion::fuseKernel(
       FT, GlobalValue::LinkageTypes::ExternalLinkage,
       M.getDataLayout().getProgramAddressSpace(), KernelName->getString(), &M);
   {
+    auto DefaultAttr = FusedFunction->getAttributes();
+    // Add uniform function attributes, i.e., attributes with identical value on
+    // each input function, to the fused function.
+    auto *FirstFunction = InputFunctions.front().F;
+    for (const auto &UniformKey : TargetInfo.getUniformKernelAttributes()) {
+      if (FirstFunction->hasFnAttribute(UniformKey)) {
+        DefaultAttr = DefaultAttr.addFnAttribute(
+            LLVMCtx, FirstFunction->getFnAttribute(UniformKey));
+      }
+    }
     // Add the collected parameter attributes to the fused function.
     // Copying the parameter attributes from their original definition in the
     // input kernels should be safe and they most likely can't be deducted later
     // on, as no caller is present in the module.
-    auto DefaultAttr = FusedFunction->getAttributes();
     auto FusedFnAttrs =
         AttributeList::get(LLVMCtx, DefaultAttr.getFnAttrs(),
                            DefaultAttr.getRetAttrs(), FusedParamAttributes);

sycl-fusion/passes/syclcp/SYCLCP.cpp

@@ -249,7 +249,7 @@ PreservedAnalyses SYCLCP::run(Module &M, ModuleAnalysisManager &AM) {
     Changed = propagateConstants(F, *ConstantsOrErr) || Changed;
   }
 
-  auto TFI = TargetFusionInfo::getTargetFusionInfo(&M);
+  TargetFusionInfo TFI{&M};
 
   if (Changed) {
     moduleCleanup(M, AM, TFI);