Fix mem2reg issue (#292)

Author: wsmoses

lib/polygeist/Ops.cpp

@@ -104,6 +104,28 @@ bool collectEffects(Operation *op,
 
         return true;
       }
+      if (*callee == "fscanf" || *callee == "__isoc99_fscanf") {
+        // Global read
+        effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Read>());
+
+        for (auto argp : llvm::enumerate(cop.getArgOperands())) {
+          auto arg = argp.value();
+          auto idx = argp.index();
+          if (idx == 0) {
+            effects.emplace_back(::mlir::MemoryEffects::Read::get(), arg,
+                                 ::mlir::SideEffects::DefaultResource::get());
+            effects.emplace_back(::mlir::MemoryEffects::Write::get(), arg,
+                                 ::mlir::SideEffects::DefaultResource::get());
+          } else if (idx == 1) {
+            effects.emplace_back(::mlir::MemoryEffects::Read::get(), arg,
+                                 ::mlir::SideEffects::DefaultResource::get());
+          } else
+            effects.emplace_back(::mlir::MemoryEffects::Write::get(), arg,
+                                 ::mlir::SideEffects::DefaultResource::get());
+        }
+
+        return true;
+      }
       if (*callee == "printf") {
         // Global read
         effects.emplace_back(

lib/polygeist/Passes/Mem2Reg.cpp

@@ -1067,10 +1067,12 @@ void removeRedundantBlockArgs(
 }
 
 std::set<std::string> NonCapturingFunctions = {
-    "free",         "printf",        "fprintf", "scanf",      "fscanf",
-    "gettimeofday", "clock_gettime", "getenv",  "strrchr",    "strlen",
-    "sprintf",      "sscanf",        "mkdir",   "fwrite",     "fread",
-    "memcpy",       "cudaMemcpy",    "memset",  "cudaMemset", "__isoc99_scanf"};
+    "free",           "printf",       "fprintf",       "scanf",
+    "fscanf",         "gettimeofday", "clock_gettime", "getenv",
+    "strrchr",        "strlen",       "sprintf",       "sscanf",
+    "mkdir",          "fwrite",       "fread",         "memcpy",
+    "cudaMemcpy",     "memset",       "cudaMemset",    "__isoc99_scanf",
+    "__isoc99_fscanf"};
 // fopen, fclose
 std::set<std::string> NoWriteFunctions = {"exit", "__errno_location"};
 // This is a straightforward implementation not optimized for speed. Optimize
@@ -1816,15 +1818,19 @@ bool isPromotable(mlir::Value AI) {
         continue;
       } else if (isa<memref::DeallocOp>(U)) {
         continue;
-      } else if (isa<func::CallOp>(U) &&
-                 cast<func::CallOp>(U).getCallee() == "free") {
-        continue;
-      } else if (isa<func::CallOp>(U)) {
-        // TODO check "no capture", currently assume as a fallback always
-        // nocapture
-        continue;
+      } else if (auto callOp = dyn_cast<func::CallOp>(U)) {
+        if (NonCapturingFunctions.count(callOp.getCallee().str()))
+          continue;
+      } else if (auto callOp = dyn_cast<LLVM::CallOp>(U)) {
+        if (auto callee = callOp.getCallee())
+          if (NonCapturingFunctions.count(callee->str()))
+            continue;
       } else if (auto CO = dyn_cast<memref::CastOp>(U)) {
         list.push_back(CO);
+      } else if (auto CO = dyn_cast<polygeist::Memref2PointerOp>(U)) {
+        list.push_back(CO);
+      } else if (auto CO = dyn_cast<polygeist::Pointer2MemrefOp>(U)) {
+        list.push_back(CO);
       } else {
         LLVM_DEBUG(llvm::dbgs()
                    << "non promotable " << AI << " due to " << *U << "\n");

test/polygeist-opt/promoteonscan.mlir

@@ -0,0 +1,26 @@
+// RUN: polygeist-opt --mem2reg --split-input-file %s | FileCheck %s
+
+module {
+  llvm.func @scanf(!llvm.ptr<i8>, ...) -> i32
+  llvm.mlir.global internal constant @str4("%d\00") {addr_space = 0 : i32}
+  func.func @_Z8BFSGraphiPPc(%arg0: i32, %arg1: memref<?xmemref<?xi8>>) -> (i32, i32) {
+    %c0_i32 = arith.constant 0 : i32
+    %0 = llvm.mlir.undef : i32
+    %alloca = memref.alloca() : memref<1xi32>
+    affine.store %0, %alloca[0] : memref<1xi32>
+    affine.store %c0_i32, %alloca[0] : memref<1xi32>
+    %4 = affine.load %arg1[1] : memref<?xmemref<?xi8>>
+    %8 = llvm.mlir.addressof @str4 : !llvm.ptr<array<3 x i8>>
+    %9 = llvm.getelementptr %8[0, 0] : (!llvm.ptr<array<3 x i8>>) -> !llvm.ptr<i8>
+    %10 = "polygeist.memref2pointer"(%alloca) : (memref<1xi32>) -> !llvm.ptr<i32>
+    %11 = llvm.call @scanf(%9, %10) : (!llvm.ptr<i8>, !llvm.ptr<i32>) -> i32
+    %12 = affine.load %alloca[0] : memref<1xi32>
+    %13 = affine.load %alloca[0] : memref<1xi32>
+    return %13, %12 : i32, i32
+  }
+// CHECK:          %[[i4:.+]] = "polygeist.memref2pointer"(%[[alloca:.+]]) : (memref<1xi32>) -> !llvm.ptr<i32>
+// CHECK-NEXT:     %[[i5:.+]] = llvm.call @scanf(%[[i3:.+]], %[[i4]]) : (!llvm.ptr<i8>, !llvm.ptr<i32>) -> i32
+// CHECK-NEXT:     %[[i6:.+]] = affine.load %[[alloca]][0] : memref<1xi32>
+// CHECK-NEXT:     return %[[i6]], %[[i6]] : i32, i32
+}
+

tools/cgeist/Test/Verification/fscanf.c

@@ -40,10 +40,10 @@ int* alloc() {
 // CHECK-NEXT:    %[[V7:.+]] = arith.index_cast %[[V6]] : i64 to index
 // CHECK-NEXT:    %[[V8:.+]] = arith.divui %[[V7]], %[[c4]] : index
 // CHECK-NEXT:    %[[i8:.+]] = memref.alloc(%[[V8]]) : memref<?xi32>
+// CHECK-NEXT:    %[[n:.+]] = arith.index_cast %[[V4]] : i32 to index
 // CHECK-NEXT:    %[[i9:.+]] = llvm.mlir.addressof @str1 : !llvm.ptr<array<4 x i8>>
 // CHECK-NEXT:    %[[i10:.+]] = llvm.getelementptr %[[i9]][0, 0] : (!llvm.ptr<array<4 x i8>>) -> !llvm.ptr<i8>
 // CHECK-NEXT:    %[[V12:.+]] = "polygeist.memref2pointer"(%[[i8]]) : (memref<?xi32>) -> !llvm.ptr<i32>
-// CHECK-NEXT:    %[[n:.+]] = arith.index_cast %[[V4]] : i32 to index
 // CHECK-NEXT:    scf.for %[[arg0:.+]] = %[[c0]] to %[[n]] step %[[c1]] {
 // CHECK-NEXT:      %[[i14:.+]] = arith.index_cast %[[arg0]] : index to i64
 // CHECK-NEXT:      %[[i15:.+]] = llvm.getelementptr %[[V12]][%[[i14]]] : (!llvm.ptr<i32>, i64) -> !llvm.ptr<i32>