Handle ostream (rust-lang#242)

* Handle ostream

* add more printFnc to utils.h

* small refactor / formating

* addressing feedback

Co-authored-by: Manuel Drehwald <[email protected]>

Author: wsmoses

enzyme/Enzyme/ActivityAnalysis.cpp

@@ -80,7 +80,18 @@ cl::opt<bool>
 
 const char *KnownInactiveFunctionsStartingWith[] = {
     "_ZN4core3fmt", "_ZN3std2io5stdio6_print", "f90io", "$ss5print",
-    "_ZNSt7__cxx1112basic_string"};
+    "_ZNSt7__cxx1112basic_string",
+    // ostream generic <<
+    "_ZStlsISt11char_traitsIcEERSt13basic_ostreamIcT_ES5_",
+    "_ZSt16__ostream_insert", "_ZNSo9_M_insert",
+    // ostream put
+    "_ZNSo3put",
+    // std::cout
+    "_ZSt4cout",
+    // generic <<
+    "_ZNSolsE",
+    // std::endl
+    "_ZNSo5flushEv", "_ZSt4endl"};
 
 const char *KnownInactiveFunctionsContains[] = {
     "__enzyme_float", "__enzyme_double", "__enzyme_integer",
@@ -178,7 +189,6 @@ const std::set<std::string> KnownInactiveFunctions = {
     "_msize",
     "ftnio_fmt_write64",
     "f90_strcmp_klen",
-    "vprintf",
     "__swift_instantiateConcreteTypeFromMangledName"};
 
 /// Is the use of value val as an argument of call CI known to be inactive

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp

@@ -4246,7 +4246,7 @@ void TypeAnalyzer::visitCallInst(CallInst &call) {
     }
 
     if (funcName == "__cxa_guard_acquire" || funcName == "printf" ||
-        funcName == "vprintf" || funcName == "puts") {
+        funcName == "vprintf" || funcName == "puts" || funcName == "fprintf") {
       updateAnalysis(&call, TypeTree(BaseType::Integer).Only(-1), &call);
     }
 

enzyme/Enzyme/Utils.h

@@ -532,8 +532,16 @@ static inline bool isCertainPrintOrFree(llvm::Function *called) {
     return false;
 
   if (called->getName() == "printf" || called->getName() == "puts" ||
-      called->getName() == "fprintf" ||
+      called->getName() == "fprintf" || called->getName() == "putchar" ||
+      called->getName().startswith(
+          "_ZStlsISt11char_traitsIcEERSt13basic_ostreamIcT_ES5_") ||
+      called->getName().startswith("_ZNSolsE") ||
+      called->getName().startswith("_ZNSo9_M_insert") ||
+      called->getName().startswith("_ZSt16__ostream_insert") ||
+      called->getName().startswith("_ZNSo3put") ||
+      called->getName().startswith("_ZSt4endl") ||
       called->getName().startswith("_ZN3std2io5stdio6_print") ||
+      called->getName().startswith("_ZNSo5flushEv") ||
       called->getName().startswith("_ZN4core3fmt") ||
       called->getName() == "vprintf" || called->getName() == "_ZdlPv" ||
       called->getName() == "_ZdlPvm" || called->getName() == "free" ||
@@ -562,30 +570,12 @@ static inline bool isCertainPrintMallocOrFree(llvm::Function *called) {
   if (called == nullptr)
     return false;
 
-  if (called->getName() == "printf" || called->getName() == "puts" ||
-      called->getName() == "fprintf" ||
-      called->getName().startswith("_ZN3std2io5stdio6_print") ||
-      called->getName().startswith("_ZN4core3fmt") ||
-      called->getName() == "vprintf" || called->getName() == "malloc" ||
-      called->getName() == "swift_allocObject" ||
-      called->getName() == "swift_release" || called->getName() == "_Znwm" ||
-      called->getName() == "_ZdlPv" || called->getName() == "_ZdlPvm" ||
-      called->getName() == "free" ||
-      shadowHandlers.find(called->getName().str()) != shadowHandlers.end())
+  if (isCertainPrintOrFree(called))
     return true;
-  switch (called->getIntrinsicID()) {
-  case llvm::Intrinsic::dbg_declare:
-  case llvm::Intrinsic::dbg_value:
-#if LLVM_VERSION_MAJOR > 6
-  case llvm::Intrinsic::dbg_label:
-#endif
-  case llvm::Intrinsic::dbg_addr:
-  case llvm::Intrinsic::lifetime_start:
-  case llvm::Intrinsic::lifetime_end:
+
+  if (isCertainMallocOrFree(called))
     return true;
-  default:
-    break;
-  }
+
   return false;
 }
 

enzyme/test/Integration/ReverseMode/cout.cpp

@@ -0,0 +1,26 @@
+// RUN: %clang++ -std=c++11 -fno-exceptions -ffast-math -O0 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S
+// RUN: %clang++ -std=c++11 -fno-exceptions -ffast-math -O1 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S
+// RUN: %clang++ -std=c++11 -fno-exceptions -ffast-math -O2 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S
+// RUN: %clang++ -std=c++11 -fno-exceptions -ffast-math -O3 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -S
+// RUN: %clang++ -std=c++11 -fno-exceptions -ffast-math -O0 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S
+// RUN: %clang++ -std=c++11 -fno-exceptions -ffast-math -O1 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S
+// RUN: %clang++ -std=c++11 -fno-exceptions -ffast-math -O2 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S
+// RUN: %clang++ -std=c++11 -fno-exceptions -ffast-math -O3 %s -S -emit-llvm -o - | %opt - %loadEnzyme -enzyme -enzyme-inline=1 -S
+
+#include <iostream>
+#include "test_utils.h"
+
+extern double __enzyme_autodiff(void*, double);
+
+double fn(double vec) {
+  std::cout << "hello" << 7 << '7' << std::endl;
+  std::cerr << vec << vec * vec << "\n";
+  return vec * vec;
+}
+
+int main() {
+    double x = 2.1;
+    double dsq = __enzyme_autodiff((void*)fn, x);
+
+    APPROX_EQ(dsq, 2 * x, 1e-7);
+}