fix cout

EnzymeAD · Sep 10, 2024 · db54e73 · db54e73
1 parent 70f9241
commit db54e73
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 22 deletions.
diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp
@@ -5614,11 +5614,6 @@ llvm::Function *EnzymeLogic::CreateBatch(RequestContext context,
       BasicBlock::Create(NewF->getContext(), "placeholders", NewF);
 
   IRBuilder<> PlaceholderBuilder(placeholderBB);
-#if LLVM_VERSION_MAJOR >= 18
-  auto It = PlaceholderBuilder.GetInsertPoint();
-  It.setHeadBit(true);
-  PlaceholderBuilder.SetInsertPoint(It);
-#endif
   PlaceholderBuilder.SetCurrentDebugLocation(DebugLoc());
   ValueToValueMapTy vmap;
   auto DestArg = NewF->arg_begin();
@@ -5990,22 +5985,23 @@ llvm::Value *EnzymeLogic::CreateNoFree(RequestContext context,
     }
   }
 
+  // clang-format off
+  const char* NoFreeDemanglesStartsWith[] = {
+    "std::basic_ostream<char, std::char_traits<char>>& std::__ostream_insert<char, std::char_traits<char>>",
+    "std::basic_ostream<char, std::char_traits<char>>::operator<<",
+    "std::ostream::operator<<",
+    "std::ostream& std::ostream::_M_insert",
+    "std::basic_ostream<char, std::char_traits<char>>& std::__ostream_insert",
+  };
+  // clang-format on
+
   if (auto CI = dyn_cast<CallInst>(todiff)) {
     TargetLibraryInfo &TLI =
         PPC.FAM.getResult<TargetLibraryAnalysis>(*CI->getParent()->getParent());
     if (isAllocationFunction(getFuncNameFromCall(CI), TLI))
       return CI;
     if (auto F = CI->getCalledFunction()) {
 
-      // clang-format off
-      const char* NoFreeDemanglesStartsWith[] = {
-          "std::basic_ostream<char, std::char_traits<char>>& std::__ostream_insert<char, std::char_traits<char>>",
-          "std::basic_ostream<char, std::char_traits<char>>::operator<<",
-          "std::ostream::operator<<",
-          "std::ostream& std::ostream::_M_insert",
-      };
-      // clang-format on
-
       demangledCall = llvm::demangle(F->getName().str());
       // replace all '> >' with '>>'
       size_t start = 0;
@@ -6018,6 +6014,44 @@ llvm::Value *EnzymeLogic::CreateNoFree(RequestContext context,
           return CI;
     }
   }
+  if (auto PN = dyn_cast<PHINode>(todiff)) {
+    Value *illegal = nullptr;
+    for (auto &op : PN->incoming_values()) {
+
+      if (auto CI = dyn_cast<CallInst>(op)) {
+        TargetLibraryInfo &TLI = PPC.FAM.getResult<TargetLibraryAnalysis>(
+            *CI->getParent()->getParent());
+        if (isAllocationFunction(getFuncNameFromCall(CI), TLI))
+          continue;
+        if (auto F = CI->getCalledFunction()) {
+
+          demangledCall = llvm::demangle(F->getName().str());
+          // replace all '> >' with '>>'
+          size_t start = 0;
+          while ((start = demangledCall.find("> >", start)) !=
+                 std::string::npos) {
+            demangledCall.replace(start, 3, ">>");
+          }
+
+          bool legal = false;
+          for (auto Name : NoFreeDemanglesStartsWith)
+            if (startsWith(demangledCall, Name)) {
+              legal = true;
+              break;
+            }
+          illegal = op;
+          break;
+        }
+        continue;
+      }
+    }
+    demangledCall = "";
+    illegal = op;
+    break;
+  }
+  if (!illegal)
+    return PN;
+}
 
   if (auto GV = dyn_cast<GlobalVariable>(todiff)) {
     if (GV->getName() == "_ZSt4cerr")
@@ -6073,6 +6107,11 @@ llvm::Value *EnzymeLogic::CreateNoFree(RequestContext context,
   llvm::raw_string_ostream ss(s);
   ss << "No create nofree of unknown value\n";
   ss << *todiff << "\n";
+  if (auto PN = dyn_cast<PHINode>(todiff)) {
+    for (auto &op : PN->incoming_values()) {
+      ss << " - " << *op << "\n";
+    }
+  }
   if (demangledCall.size()) {
     ss << " demangled (" << demangledCall << ")\n";
   }

diff --git a/enzyme/test/Enzyme/ReverseMode/freefn.ll b/enzyme/test/Enzyme/ReverseMode/freefn.ll
@@ -1,26 +1,26 @@
 ; RUN: if [ %llvmver -ge 15 ]; then %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="enzyme,function(mem2reg,sroa,instsimplify)" -S | FileCheck %s; fi
 
-define double @_f2(ptr %0, i64 %1) personality ptr null {
-  call void @__rust_dealloc(ptr %0, i64 %1)
+define double @_f2(i8* %0, i64 %1) {
+  call void @__rust_dealloc(i8* %0, i64 %1)
   ret double 0.000000e+00
 }
 
-declare void @__rust_dealloc(ptr, i64)
+declare void @__rust_dealloc(i8*, i64)
 
 declare double @__enzyme_fwddiff(...)
 
 define double @enzyme_opt_helper_0() {
-  %1 = call double (...) @__enzyme_fwddiff(ptr @_f2, metadata !"enzyme_dup", ptr null, ptr null, metadata !"enzyme_const", i64 0)
+  %1 = call double (...) @__enzyme_fwddiff(double (i8*, i64)* @_f2, metadata !"enzyme_dup", i8* null, i8* null, metadata !"enzyme_const", i64 0)
   ret double 0.000000e+00
 }
 
-; CHECK: define internal double @fwddiffe_f2(ptr %0, ptr %"'", i64 %1) 
-; CHECK-NEXT:   call void @__rust_dealloc(ptr %0, i64 %1)
-; CHECK-NEXT:   %3 = icmp ne ptr %0, %"'"
+; CHECK: define internal double @fwddiffe_f2(i8* %0, i8* %"'", i64 %1) 
+; CHECK-NEXT:   call void @__rust_dealloc(i8* %0, i64 %1)
+; CHECK-NEXT:   %3 = icmp ne i8* %0, %"'"
 ; CHECK-NEXT:   br i1 %3, label %free0.i, label %__enzyme_checked_free_1___rust_dealloc.exit
 
 ; CHECK: free0.i:
-; CHECK-NEXT:   call void @__rust_dealloc(ptr %"'", i64 %1) 
+; CHECK-NEXT:   call void @__rust_dealloc(i8* %"'", i64 %1) 
 ; CHECK-NEXT:   br label %__enzyme_checked_free_1___rust_dealloc.exit
 
 ; CHECK: __enzyme_checked_free_1___rust_dealloc.exit: