PGO fixes

Author: ivanradanov

CMakeLists.txt

@@ -26,7 +26,9 @@ set(LLVM_LIBRARY_OUTPUT_INTDIR ${CMAKE_BINARY_DIR}/lib)
 find_program(XXD_BIN xxd)
 
 # TODO should depend on OS
-set(POLYGEIST_PGO_DATA_DIR "/var/tmp/polygeist/pgo/" CACHE STRING "Directory for PGO data")
+set(POLYGEIST_PGO_DEFAULT_DATA_DIR "/var/tmp/polygeist/pgo/" CACHE STRING "Directory for PGO data")
+set(POLYGEIST_PGO_ALTERNATIVE_ENV_VAR "POLYGEIST_PGO_ALTERNATIVE" CACHE STRING "Env var name to specify alternative to profile")
+set(POLYGEIST_PGO_DATA_DIR_ENV_VAR "POLYGEIST_PGO_DATA_DIR" CACHE STRING "Env var name to specify PGO data dir")
 
 if (CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
   project(polygeist LANGUAGES CXX C)

lib/polygeist/ExecutionEngine/CMakeLists.txt

@@ -52,7 +52,9 @@ if(POLYGEIST_ENABLE_CUDA)
       ${bc_flags}
       ${infile} -o ${bc_outfile}
       -I${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
-      -DPOLYGEIST_PGO_DATA_DIR="${POLYGEIST_PGO_DATA_DIR}"
+      -DPOLYGEIST_PGO_DEFAULT_DATA_DIR="${POLYGEIST_PGO_DEFAULT_DATA_DIR}"
+      -DPOLYGEIST_PGO_ALTERNATIVE_ENV_VAR="${POLYGEIST_PGO_ALTERNATIVE_ENV_VAR}"
+      -DPOLYGEIST_PGO_DATA_DIR_ENV_VAR="${POLYGEIST_PGO_DATA_DIR_ENV_VAR}"
       DEPENDS ${infile}
       COMMENT "Building LLVM bitcode ${bc_outfile}"
       VERBATIM

lib/polygeist/ExecutionEngine/CudaRuntimeWrappers.cpp

@@ -86,21 +86,21 @@ class PGOState {
     struct timespec start_clock;
   };
 
-  // TODO define these in cmake files (depending on target OS and used in the
-  // compiler too)
-  static constexpr const char *dirname = POLYGEIST_PGO_DATA_DIR;
-  static constexpr const char *alternativeEnvVar =
-      "POLYGEIST_PGO_KERNEL_ALTERNATIVE";
-
   inline static int alternative;
+  inline static std::string dirname;
   inline thread_local static std::mutex mutex;
   inline thread_local static std::map<std::string, State *> states;
 
   std::string kernelId;
   int totalAlternatives;
 
-  PGOState(const char *kernelId, int totalAlternatives)
-      : kernelId(kernelId), totalAlternatives(totalAlternatives) {}
+  PGOState(const char *kernelId_c, int totalAlternatives)
+      : totalAlternatives(totalAlternatives) {
+    kernelId = kernelId_c;
+    for (char &c : kernelId)
+      if (c == '/')
+        c = '+';
+  }
   void end() {
     struct timespec end_clock;
     cudaDeviceSynchronize();
@@ -123,7 +123,7 @@ class PGOState {
     if (0 <= alternative && alternative < totalAlternatives) {
       // TODO error handling
       std::ofstream ofile;
-      ofile.open(std::string(dirname) + kernelId,
+      ofile.open(std::string(dirname) + "/" + kernelId,
                  std::ios::out | std::ios::app);
       ofile << alternative << " " << elapsed << std::endl;
       ofile.close();
@@ -149,12 +149,18 @@ class PGOState {
 
   int getAlternative() {
     static int init = [&] {
-      if (char *i = getenv(alternativeEnvVar)) {
+      if (char *i = getenv(POLYGEIST_PGO_ALTERNATIVE_ENV_VAR)) {
         this->alternative = atoi(i);
       } else {
-        std::cerr << alternativeEnvVar << " not defined" << std::endl;
+        std::cerr << POLYGEIST_PGO_ALTERNATIVE_ENV_VAR << " not defined"
+                  << std::endl;
         exit(1);
       }
+      if (char *d = getenv(POLYGEIST_PGO_DATA_DIR_ENV_VAR)) {
+        this->dirname = d;
+      } else {
+        this->dirname = POLYGEIST_PGO_DEFAULT_DATA_DIR;
+      }
       std::filesystem::create_directories(dirname);
       return 0;
     }();

lib/polygeist/Passes/CMakeLists.txt

@@ -53,7 +53,9 @@ add_mlir_dialect_library(MLIRPolygeistTransforms
 
 target_compile_definitions(obj.MLIRPolygeistTransforms
   PRIVATE
-  POLYGEIST_PGO_DATA_DIR="${POLYGEIST_PGO_DATA_DIR}"
+  POLYGEIST_PGO_DEFAULT_DATA_DIR="${POLYGEIST_PGO_DEFAULT_DATA_DIR}"
+  POLYGEIST_PGO_ALTERNATIVE_ENV_VAR="${POLYGEIST_PGO_ALTERNATIVE_ENV_VAR}"
+  POLYGEIST_PGO_DATA_DIR_ENV_VAR="${POLYGEIST_PGO_DATA_DIR_ENV_VAR}"
 )
 
 if(POLYGEIST_ENABLE_CUDA)

lib/polygeist/Passes/ConvertPolygeistToLLVM.cpp

@@ -1692,6 +1692,9 @@ struct LowerGPUAlternativesOp
         return stream.str();
       }();
       locStr += gao->getAttrOfType<StringAttr>("polygeist.altop.id").data();
+      for (char &c : locStr)
+        if (c == '/')
+          c = '+';
 
       if (PolygeistAlternativesMode == PAM_PGO_Profile) {
         rewriter.setInsertionPoint(gao);
@@ -1734,15 +1737,21 @@ struct LowerGPUAlternativesOp
         rewriter.eraseOp(gao);
         return success();
       } else if (PolygeistAlternativesMode == PAM_PGO_Opt) {
-        static constexpr const char *dirname = POLYGEIST_PGO_DATA_DIR;
+        std::string dirname = []() {
+          if (char *d = getenv(POLYGEIST_PGO_DATA_DIR_ENV_VAR)) {
+            return std::string(d);
+          } else {
+            return std::string(POLYGEIST_PGO_DEFAULT_DATA_DIR);
+          }
+        }();
         // TODO error handling
         std::ifstream ifile;
         int numAlternatives = gao->getNumRegions();
         std::vector<std::vector<double>> timings;
         for (int i = 0; i < numAlternatives; i++) {
           timings.push_back({});
         }
-        ifile.open(std::string(dirname) + locStr, std::ios::in);
+        ifile.open(std::string(dirname) + "/" + locStr, std::ios::in);
         while (ifile) {
           int alt;
           double time;
@@ -1763,6 +1772,7 @@ struct LowerGPUAlternativesOp
             avgs.push_back(std::numeric_limits<double>::infinity());
           } else {
             // TODO might get some round off errors here, maybe use a better alg
+            // or median
             avgs.push_back(
                 std::accumulate(timings[i].begin(), timings[i].end(), 0.0f) /
                 timings[i].size());