From b031af8430196b21dbfceafca20ef4c49daf3dfc Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Thu, 23 Dec 2021 21:39:56 -0500
Subject: [PATCH] Mem2Reg on LLVM

---
 include/polygeist/Passes/Passes.td               |  2 +-
 lib/polygeist/Passes/Mem2Reg.cpp                 | 14 ++++++++------
 test/polygeist-opt/llvmmem2reg.mlir              | 16 ++++++++++++++++
 .../mlir-clang/Test/Verification/charswitch.cpp  |  4 +++-
 4 files changed, 28 insertions(+), 8 deletions(-)
 create mode 100644 test/polygeist-opt/llvmmem2reg.mlir

diff --git a/include/polygeist/Passes/Passes.td b/include/polygeist/Passes/Passes.td
index 0ecd1cf..a85213d 100644
--- a/include/polygeist/Passes/Passes.td
+++ b/include/polygeist/Passes/Passes.td
@@ -10,7 +10,7 @@ def AffineCFG : FunctionPass<"affine-cfg"> {
 
 def Mem2Reg : FunctionPass<"mem2reg"> {
   let summary = "Replace scf.if and similar with affine.if";
-  let constructor = "mlir::polygeist::replaceAffineCFGPass()";
+  let constructor = "mlir::polygeist::createMem2RegPass()";
 }
 
 def ParallelLower : Pass<"parallel-lower", "mlir::ModuleOp"> {
diff --git a/lib/polygeist/Passes/Mem2Reg.cpp b/lib/polygeist/Passes/Mem2Reg.cpp
index a8226aa..ea52569 100644
--- a/lib/polygeist/Passes/Mem2Reg.cpp
+++ b/lib/polygeist/Passes/Mem2Reg.cpp
@@ -1344,10 +1344,12 @@ StoreMap getLastStored(mlir::Value AI) {
           }
         }
         lastStored.insert(vec);
-      } else if (auto LO = dyn_cast<LLVM::LoadOp>(U)) {
-        lastStored.insert({});
-      } else if (auto SO = dyn_cast<LLVM::StoreOp>(U)) {
-        lastStored.insert({});
+      } else if (isa<LLVM::LoadOp>(U)) {
+        std::vector<ssize_t> vec;
+        lastStored.insert(vec);
+      } else if (isa<LLVM::StoreOp>(U)) {
+        std::vector<ssize_t> vec;
+        lastStored.insert(vec);
       } else if (auto SO = dyn_cast<memref::LoadOp>(U)) {
         std::vector<ssize_t> vec;
         for (auto idx : SO.getIndices()) {
@@ -1419,7 +1421,6 @@ void Mem2Reg::runOnFunction() {
       LLVM_DEBUG(llvm::dbgs() << " attempting to promote " << AI << "\n");
       auto lastStored = getLastStored(AI);
       for (auto &vec : lastStored) {
-
         LLVM_DEBUG(llvm::dbgs() << " + forwarding vec to promote {";
                    for (auto m
                         : vec) llvm::dbgs()
@@ -1462,7 +1463,8 @@ void Mem2Reg::runOnFunction() {
         list.pop_front();
 
         for (auto U : val.getUsers()) {
-          if (isa<memref::StoreOp, AffineStoreOp, memref::DeallocOp>(U)) {
+          if (isa<LLVM::StoreOp, memref::StoreOp, AffineStoreOp,
+                  memref::DeallocOp>(U)) {
             toErase.push_back(U);
           } else if (isa<CallOp>(U) && cast<CallOp>(U).callee() == "free") {
             toErase.push_back(U);
diff --git a/test/polygeist-opt/llvmmem2reg.mlir b/test/polygeist-opt/llvmmem2reg.mlir
new file mode 100644
index 0000000..6acda29
--- /dev/null
+++ b/test/polygeist-opt/llvmmem2reg.mlir
@@ -0,0 +1,16 @@
+// RUN: polygeist-opt --mem2reg --split-input-file %s | FileCheck %s
+
+module {
+  func @ll(%arg0: !llvm.ptr<i8>) -> !llvm.ptr<i8> {
+    %c1_i64 = arith.constant 1 : i64
+    %2 = llvm.alloca %c1_i64 x !llvm.ptr<i8> : (i64) -> !llvm.ptr<ptr<i8>>
+    llvm.store %arg0, %2 : !llvm.ptr<ptr<i8>>
+    %3 = llvm.load %2 : !llvm.ptr<ptr<i8>>
+	return %3 : !llvm.ptr<i8>
+  }
+}
+
+// CHECK:   func @ll(%arg0: !llvm.ptr<i8>) -> !llvm.ptr<i8> {
+// CHECK-NEXT:     %c1_i64 = arith.constant 1 : i64
+// CHECK-NEXT:     return %arg0 : !llvm.ptr<i8>
+// CHECK-NEXT:   }
diff --git a/tools/mlir-clang/Test/Verification/charswitch.cpp b/tools/mlir-clang/Test/Verification/charswitch.cpp
index ddc880a..e2bd369 100644
--- a/tools/mlir-clang/Test/Verification/charswitch.cpp
+++ b/tools/mlir-clang/Test/Verification/charswitch.cpp
@@ -1,5 +1,6 @@
 // RUN: mlir-clang %s --function=foo -S | FileCheck %s
 
+extern "C" {
 int foo(char t) {
   int n = 10;
   switch (t) {
@@ -14,9 +15,10 @@ int foo(char t) {
   }
   return n;
 }
+}
 
 // TODO the select should be canonicalized better
-// CHECK:   func @foo(%arg0: i32) -> i32 attributes {llvm.linkage = #llvm.linkage<external>} {
+// CHECK:   func @foo(%arg0: i8) -> i32 attributes {llvm.linkage = #llvm.linkage<external>} {
 // CHECK-NEXT:     %c-1_i32 = arith.constant -1 : i32
 // CHECK-NEXT:     %c30_i32 = arith.constant 30 : i32
 // CHECK-NEXT:     %false = arith.constant false