Frontend fixes

include/polygeist/Passes/Passes.h

@@ -37,6 +37,7 @@ std::unique_ptr<Pass> createParallelLowerPass(
 std::unique_ptr<Pass> createConvertCudaRTtoCPUPass();
 std::unique_ptr<Pass> createConvertCudaRTtoGPUPass();
 std::unique_ptr<Pass> createConvertCudaRTtoHipRTPass();
+std::unique_ptr<Pass> createFixGPUFuncPass();
 std::unique_ptr<Pass> createSCFParallelLoopUnrollPass(int unrollFactor = 2);
 std::unique_ptr<Pass>
 createConvertPolygeistToLLVMPass(const LowerToLLVMOptions &options,
@@ -49,6 +50,8 @@ createConvertParallelToGPUPass1(std::string arch = "sm_60");
 std::unique_ptr<Pass>
 createConvertParallelToGPUPass2(bool emitGPUKernelLaunchBounds = true);
 std::unique_ptr<Pass> createMergeGPUModulesPass();
+std::unique_ptr<Pass> createLowerAlternativesPass();
+std::unique_ptr<Pass> createCollectKernelStatisticsPass();
 std::unique_ptr<Pass> createGpuSerializeToCubinPass(
     StringRef arch, StringRef features, int llvmOptLevel, int ptxasOptLevel,
     std::string ptxasPath, std::string libDevicePath, bool outputIntermediate);

include/polygeist/Passes/Passes.td

@@ -23,13 +23,31 @@ def SCFParallelLoopUnroll : Pass<"scf-parallel-loop-unroll"> {
   ];
 }
 
+def CollectKernelStatistics : Pass<"collect-kernel-statistics", "mlir::ModuleOp"> {
+  let summary = "Lower cudart functions to cpu versions";
+  let dependentDialects = [];
+  let constructor = "mlir::polygeist::createCollectKernelStatisticsPass()";
+}
+
+def LowerAlternatives : Pass<"lower-alternatives", "mlir::ModuleOp"> {
+  let summary = "Lower alternatives if in opt mode";
+  let dependentDialects = [];
+  let constructor = "mlir::polygeist::createLowerAlternativesPass()";
+}
+
 def ConvertCudaRTtoCPU : Pass<"convert-cudart-to-cpu", "mlir::ModuleOp"> {
   let summary = "Lower cudart functions to cpu versions";
   let dependentDialects =
       ["memref::MemRefDialect", "func::FuncDialect", "LLVM::LLVMDialect"];
   let constructor = "mlir::polygeist::createConvertCudaRTtoCPUPass()";
 }
 
+def FixGPUFunc : Pass<"fix-gpu-func", "mlir::gpu::GPUModuleOp"> {
+  let summary = "Fix nested calls to gpu functions we generate in the frontend";
+  let dependentDialects = ["func::FuncDialect", "LLVM::LLVMDialect", "gpu::GPUDialect"];
+  let constructor = "mlir::polygeist::createFixGPUFuncPass()";
+}
+
 def ConvertCudaRTtoGPU : Pass<"convert-cudart-to-gpu", "mlir::ModuleOp"> {
   let summary = "Lower cudart functions to generic gpu versions";
   let dependentDialects =

include/polygeist/PolygeistOps.td

@@ -18,6 +18,18 @@ include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Dialect/LLVMIR/LLVMOpsInterfaces.td"
 
+def UndefOp
+    : Polygeist_Op<"undef", [Pure]> {
+  let summary = "More flexible undef op";
+  let skipDefaultBuilders = 1;
+  let results = (outs AnyType:$result);
+  let builders = [
+    OpBuilder<(ins "Type":$type), [{
+      $_state.types.push_back(type);
+    }]>];
+  let hasCanonicalizer = true;
+}
+
 def NoopOp
     : Polygeist_Op<"noop",
                    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
@@ -29,6 +41,16 @@ def NoopOp
   let description = [{}];
 }
 
+def GetDeviceGlobalOp
+    : Polygeist_Op<"get_device_global",
+                   [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+                    DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+  let summary = "";
+  let arguments = (ins FlatSymbolRefAttr:$name);
+  let results = (outs AnyStaticShapeMemRef:$result);
+  let description = [{}];
+}
+
 def CacheLoad
     : Polygeist_Op<"cacheload"> {
 

lib/polygeist/ExecutionEngine/PGORuntime.h

@@ -4,12 +4,15 @@
 // PGO functions which should know whether the code in the alternatives op is
 // GPU code - we can add an attrib to the alternatives op for that
 
+#include <cstdlib>
 #include <ctime>
 #include <filesystem>
 #include <fstream>
 #include <iostream>
 #include <map>
 #include <mutex>
+#include <numeric>
+#include <vector>
 
 extern "C" int32_t mgpurtDeviceSynchronizeErr(void);
 
@@ -26,26 +29,38 @@ class PGOState {
     struct timespec start_clock;
   };
 
+  struct Logger {
+    std::map<std::string, std::vector<double>> timings;
+    ~Logger() { PGOState::writeResults(); }
+  };
+
   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;
+  inline static Logger logger;
 
-  std::string kernelId;
+  const char *kernelId_c;
   int totalAlternatives;
 
-  PGOState(const char *kernelId_c, int totalAlternatives)
-      : totalAlternatives(totalAlternatives) {
-    kernelId = kernelId_c;
+  std::string getKernelId() {
+    std::string kernelId = kernelId_c;
     for (char &c : kernelId)
       if (c == '/')
         c = '+';
+    return kernelId;
+  }
+
+  PGOState(const char *kernelId_c, int totalAlternatives)
+      : totalAlternatives(totalAlternatives) {
+    this->kernelId_c = kernelId_c;
   }
   void end() {
     struct timespec end_clock;
     mgpurtDeviceSynchronizeErr();
     clock_gettime(CLOCK_MONOTONIC, &end_clock);
 
+    auto kernelId = getKernelId();
     std::unique_lock<std::mutex> lock(mutex);
     if (states.count(kernelId) == 0) {
       std::cerr << "No kernel with id " << kernelId << "running" << std::endl;
@@ -59,21 +74,16 @@ class PGOState {
     double elapsed =
         (tmp_clock.tv_sec + ((double)tmp_clock.tv_nsec) * .000000001);
 
-    // Only write to file if we are profiling a valid alternative
-    if (0 <= alternative && alternative < totalAlternatives) {
-      // TODO error handling
-      std::ofstream ofile;
-      ofile.open(std::string(dirname) + "/" + kernelId,
-                 std::ios::out | std::ios::app);
-      ofile << alternative << " " << elapsed << std::endl;
-      ofile.close();
-    }
+    if (states.count(kernelId) == 0)
+      logger.timings[kernelId] = {};
+    logger.timings[kernelId].push_back(elapsed);
 
     delete state;
     states.erase(states.find(kernelId));
   }
 
   void start() {
+    auto kernelId = getKernelId();
     std::unique_lock<std::mutex> lock(mutex);
     State *state = new State();
     if (states.count(kernelId) == 1) {
@@ -87,6 +97,21 @@ class PGOState {
     clock_gettime(CLOCK_MONOTONIC, &state->start_clock);
   }
 
+  static void writeResults() {
+    // Only write to file if we are profiling a valid alternative
+    for (auto &pair : logger.timings) {
+      auto &kernelId = std::get<0>(pair);
+      auto &timings = std::get<1>(pair);
+      auto elapsed = std::accumulate(timings.begin(), timings.end(), 0.0f);
+      // TODO error handling
+      std::ofstream ofile;
+      ofile.open(std::string(dirname) + "/" + kernelId,
+                 std::ios::out | std::ios::app);
+      ofile << alternative << " " << elapsed << std::endl;
+      ofile.close();
+    }
+  }
+
   int getAlternative() {
     static int init = [&] {
       if (char *i = getenv(POLYGEIST_PGO_ALTERNATIVE_ENV_VAR)) {
@@ -102,12 +127,10 @@ class PGOState {
         this->dirname = POLYGEIST_PGO_DEFAULT_DATA_DIR;
       }
       std::filesystem::create_directories(dirname);
+
       return 0;
     }();
-    if (0 <= alternative && alternative < totalAlternatives)
-      return alternative;
-    else
-      return 0;
+    return alternative % totalAlternatives;
   }
 
   ~PGOState() {}