[SYCL-JIT] Store env variables in configuration (#14860)

Jakub Chlanda · web-flow · commit 7c533bd8abdf · 2024-08-21T13:54:17.000+02:00
By storing the env variables in the config helper we eliminate the need
for extra function parameters while also making it accessible throughout
the codebase.
diff --git a/sycl-jit/common/include/DynArray.h b/sycl-jit/common/include/DynArray.h
@@ -10,6 +10,7 @@
 #define SYCL_FUSION_COMMON_DYNARRAY_H
 
 #include <algorithm>
+#include <cstring>
 
 namespace jit_compiler {
 
@@ -22,6 +23,11 @@ template <typename T> class DynArray {
 
   explicit DynArray(size_t Size) { init(Size); }
 
+  template <typename InputIt> DynArray(InputIt Begin, InputIt End) {
+    init(End - Begin);
+    std::copy(Begin, End, this->begin());
+  }
+
   ~DynArray() { deinit(); }
 
   DynArray(const DynArray &Other) {
diff --git a/sycl-jit/jit-compiler/include/KernelFusion.h b/sycl-jit/jit-compiler/include/KernelFusion.h
@@ -64,9 +64,7 @@ JITResult fuseKernels(View<SYCLKernelInfo> KernelInformation,
 
 JITResult materializeSpecConstants(const char *KernelName,
                                    jit_compiler::SYCLKernelBinaryInfo &BinInfo,
-                                   View<unsigned char> SpecConstBlob,
-                                   const char *TargetCPU,
-                                   const char *TargetFeatures);
+                                   View<unsigned char> SpecConstBlob);
 
 /// Clear all previously set options.
 void resetJITConfiguration();
diff --git a/sycl-jit/jit-compiler/include/Options.h b/sycl-jit/jit-compiler/include/Options.h
@@ -13,7 +13,13 @@
 
 namespace jit_compiler {
 
-enum OptionID { VerboseOutput, EnableCaching, TargetDeviceInfo };
+enum OptionID {
+  VerboseOutput,
+  EnableCaching,
+  TargetDeviceInfo,
+  TargetCPU,
+  TargetFeatures
+};
 
 class OptionPtrBase {
 protected:
@@ -90,6 +96,17 @@ struct JITTargetInfo
   using OptionBase::OptionBase;
 };
 
+struct JITTargetCPU
+    : public OptionBase<JITTargetCPU, OptionID::TargetCPU, DynArray<char>> {
+  using OptionBase::OptionBase;
+};
+
+struct JITTargetFeatures
+    : public OptionBase<JITTargetFeatures, OptionID::TargetFeatures,
+                        DynArray<char>> {
+  using OptionBase::OptionBase;
+};
+
 } // namespace option
 } // namespace jit_compiler
 
diff --git a/sycl-jit/jit-compiler/lib/KernelFusion.cpp b/sycl-jit/jit-compiler/lib/KernelFusion.cpp
@@ -73,8 +73,7 @@ static bool isTargetFormatSupported(BinaryFormat TargetFormat) {
 extern "C" JITResult
 materializeSpecConstants(const char *KernelName,
                          jit_compiler::SYCLKernelBinaryInfo &BinInfo,
-                         View<unsigned char> SpecConstBlob,
-                         const char *TargetCPU, const char *TargetFeatures) {
+                         View<unsigned char> SpecConstBlob) {
   auto &JITCtx = JITContext::getInstance();
 
   TargetInfo TargetInfo = ConfigHelper::get<option::JITTargetInfo>();
@@ -107,8 +106,7 @@ materializeSpecConstants(const char *KernelName,
 
   SYCLKernelInfo &MaterializerKernelInfo = *ModuleInfo.getKernelFor(KernelName);
   if (auto Error = translation::KernelTranslator::translateKernel(
-          MaterializerKernelInfo, *NewMod, JITCtx, TargetFormat, TargetCPU,
-          TargetFeatures)) {
+          MaterializerKernelInfo, *NewMod, JITCtx, TargetFormat)) {
     return errorToFusionResult(std::move(Error),
                                "Translation to output format failed");
   }
diff --git a/sycl-jit/jit-compiler/lib/translation/KernelTranslation.cpp b/sycl-jit/jit-compiler/lib/translation/KernelTranslation.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "KernelTranslation.h"
+#include "helper/ConfigHelper.h"
 
 #include "SPIRVLLVMTranslation.h"
 #include "llvm/Bitcode/BitcodeReader.h"
@@ -168,11 +169,10 @@ KernelTranslator::loadSPIRVKernel(llvm::LLVMContext &LLVMCtx,
   return SPIRVLLVMTranslator::loadSPIRVKernel(LLVMCtx, Kernel);
 }
 
-llvm::Error
-KernelTranslator::translateKernel(SYCLKernelInfo &Kernel, llvm::Module &Mod,
-                                  JITContext &JITCtx, BinaryFormat Format,
-                                  const std::string &TargetCPU,
-                                  const std::string &TargetFeatures) {
+llvm::Error KernelTranslator::translateKernel(SYCLKernelInfo &Kernel,
+                                              llvm::Module &Mod,
+                                              JITContext &JITCtx,
+                                              BinaryFormat Format) {
 
   KernelBinary *KernelBin = nullptr;
   switch (Format) {
@@ -187,7 +187,7 @@ KernelTranslator::translateKernel(SYCLKernelInfo &Kernel, llvm::Module &Mod,
   }
   case BinaryFormat::PTX: {
     llvm::Expected<KernelBinary *> BinaryOrError =
-        translateToPTX(Kernel, Mod, JITCtx, TargetCPU, TargetFeatures);
+        translateToPTX(Kernel, Mod, JITCtx);
     if (auto Error = BinaryOrError.takeError()) {
       return Error;
     }
@@ -196,7 +196,7 @@ KernelTranslator::translateKernel(SYCLKernelInfo &Kernel, llvm::Module &Mod,
   }
   case BinaryFormat::AMDGCN: {
     llvm::Expected<KernelBinary *> BinaryOrError =
-        translateToAMDGCN(Kernel, Mod, JITCtx, TargetCPU, TargetFeatures);
+        translateToAMDGCN(Kernel, Mod, JITCtx);
     if (auto Error = BinaryOrError.takeError())
       return Error;
     KernelBin = *BinaryOrError;
@@ -227,10 +227,9 @@ KernelTranslator::translateToSPIRV(llvm::Module &Mod, JITContext &JITCtx) {
   return SPIRVLLVMTranslator::translateLLVMtoSPIRV(Mod, JITCtx);
 }
 
-llvm::Expected<KernelBinary *> KernelTranslator::translateToPTX(
-    SYCLKernelInfo &KernelInfo, llvm::Module &Mod, JITContext &JITCtx,
-    [[maybe_unused]] const std::string &TargetCPU,
-    [[maybe_unused]] const std::string &TargetFeatures) {
+llvm::Expected<KernelBinary *>
+KernelTranslator::translateToPTX(SYCLKernelInfo &KernelInfo, llvm::Module &Mod,
+                                 JITContext &JITCtx) {
 #ifndef JIT_SUPPORT_PTX
   (void)KernelInfo;
   (void)Mod;
@@ -261,8 +260,10 @@ llvm::Expected<KernelBinary *> KernelTranslator::translateToPTX(
 
   // Give priority to user specified values (through environment variables:
   // SYCL_JIT_AMDGCN_PTX_TARGET_CPU and SYCL_JIT_AMDGCN_PTX_TARGET_FEATURES).
-  llvm::StringRef CPU{TargetCPU};
-  llvm::StringRef Features{TargetFeatures};
+  auto CPUVal = ConfigHelper::get<option::JITTargetCPU>();
+  auto FeaturesVal = ConfigHelper::get<option::JITTargetFeatures>();
+  llvm::StringRef CPU = CPUVal.begin();
+  llvm::StringRef Features = FeaturesVal.begin();
 
   auto *KernelFunc = Mod.getFunction(KernelInfo.Name.c_str());
   // If they were not set, use default and consult the module for alternatives
@@ -309,10 +310,9 @@ llvm::Expected<KernelBinary *> KernelTranslator::translateToPTX(
 #endif // JIT_SUPPORT_PTX
 }
 
-llvm::Expected<KernelBinary *> KernelTranslator::translateToAMDGCN(
-    SYCLKernelInfo &KernelInfo, llvm::Module &Mod, JITContext &JITCtx,
-    [[maybe_unused]] const std::string &TargetCPU,
-    [[maybe_unused]] const std::string &TargetFeatures) {
+llvm::Expected<KernelBinary *>
+KernelTranslator::translateToAMDGCN(SYCLKernelInfo &KernelInfo,
+                                    llvm::Module &Mod, JITContext &JITCtx) {
 #ifndef JIT_SUPPORT_AMDGCN
   (void)KernelInfo;
   (void)Mod;
@@ -341,8 +341,10 @@ llvm::Expected<KernelBinary *> KernelTranslator::translateToAMDGCN(
         "Failed to load and translate AMDGCN LLVM IR module with error %s",
         ErrorMessage.c_str());
 
-  llvm::StringRef CPU{TargetCPU};
-  llvm::StringRef Features{TargetFeatures};
+  auto CPUVal = ConfigHelper::get<option::JITTargetCPU>();
+  auto FeaturesVal = ConfigHelper::get<option::JITTargetFeatures>();
+  llvm::StringRef CPU = CPUVal.begin();
+  llvm::StringRef Features = FeaturesVal.begin();
 
   auto *KernelFunc = Mod.getFunction(KernelInfo.Name.c_str());
   if (CPU.empty()) {
diff --git a/sycl-jit/jit-compiler/lib/translation/KernelTranslation.h b/sycl-jit/jit-compiler/lib/translation/KernelTranslation.h
@@ -25,9 +25,7 @@ class KernelTranslator {
   loadKernels(llvm::LLVMContext &LLVMCtx, std::vector<SYCLKernelInfo> &Kernels);
 
   static llvm::Error translateKernel(SYCLKernelInfo &Kernel, llvm::Module &Mod,
-                                     JITContext &JITCtx, BinaryFormat Format,
-                                     const std::string &TargetCPU = {},
-                                     const std::string &TargetFeatures = {});
+                                     JITContext &JITCtx, BinaryFormat Format);
 
 private:
   ///
@@ -44,14 +42,11 @@ class KernelTranslator {
                                                          JITContext &JITCtx);
 
   static llvm::Expected<KernelBinary *>
-  translateToPTX(SYCLKernelInfo &Kernel, llvm::Module &Mod, JITContext &JITCtx,
-                 const std::string &TargetCPU = {},
-                 const std::string &TargetFeatures = {});
+  translateToPTX(SYCLKernelInfo &Kernel, llvm::Module &Mod, JITContext &JITCtx);
 
   static llvm::Expected<KernelBinary *>
   translateToAMDGCN(SYCLKernelInfo &KernelInfo, llvm::Module &Mod,
-                    JITContext &JITCtx, const std::string &TargetCPU = {},
-                    const std::string &TargetFeatures = {});
+                    JITContext &JITCtx);
 };
 } // namespace translation
 } // namespace jit_compiler
diff --git a/sycl/source/detail/jit_compiler.cpp b/sycl/source/detail/jit_compiler.cpp
@@ -663,15 +663,20 @@ ur_kernel_handle_t jit_compiler::materializeSpecConstants(
       detail::SYCLConfig<detail::SYCL_RT_WARNING_LEVEL>::get() > 0;
   AddToConfigHandle(
       ::jit_compiler::option::JITEnableVerbose::set(DebugEnabled));
-
-  std::string TargetCPU =
-      detail::SYCLConfig<detail::SYCL_JIT_AMDGCN_PTX_TARGET_CPU>::get();
-  std::string TargetFeatures =
-      detail::SYCLConfig<detail::SYCL_JIT_AMDGCN_PTX_TARGET_FEATURES>::get();
+  auto SetUpOption = [](const std::string &Value) {
+    ::jit_compiler::JITEnvVar Option(Value.begin(), Value.end());
+    return Option;
+  };
+  ::jit_compiler::JITEnvVar TargetCPUOpt = SetUpOption(
+      detail::SYCLConfig<detail::SYCL_JIT_AMDGCN_PTX_TARGET_CPU>::get());
+  AddToConfigHandle(::jit_compiler::option::JITTargetCPU::set(TargetCPUOpt));
+  ::jit_compiler::JITEnvVar TargetFeaturesOpt = SetUpOption(
+      detail::SYCLConfig<detail::SYCL_JIT_AMDGCN_PTX_TARGET_FEATURES>::get());
+  AddToConfigHandle(
+      ::jit_compiler::option::JITTargetFeatures::set(TargetFeaturesOpt));
 
   auto MaterializerResult =
-      MaterializeSpecConstHandle(KernelName.c_str(), BinInfo, SpecConstBlob,
-                                 TargetCPU.c_str(), TargetFeatures.c_str());
+      MaterializeSpecConstHandle(KernelName.c_str(), BinInfo, SpecConstBlob);
   if (MaterializerResult.failed()) {
     std::string Message{"Compilation for kernel failed with message:\n"};
     Message.append(MaterializerResult.getErrorMessage());
diff --git a/sycl/source/detail/jit_compiler.hpp b/sycl/source/detail/jit_compiler.hpp
@@ -25,6 +25,7 @@ struct SYCLKernelInfo;
 struct SYCLKernelAttribute;
 template <typename T> class DynArray;
 using ArgUsageMask = DynArray<uint8_t>;
+using JITEnvVar = DynArray<char>;
 } // namespace jit_compiler
 
 namespace sycl {
