From cbccccf67edba7f683339f7fe0e0493235f028db Mon Sep 17 00:00:00 2001
From: Brant Qian <scharfrichterq@gmail.com>
Date: Thu, 2 May 2024 05:20:01 +0800
Subject: [PATCH] More fprintf handling (#1844)

Co-authored-by: William S. Moses <gh@wsmoses.com>
---
 enzyme/Enzyme/Enzyme.cpp                  | 11 ++++++++-
 enzyme/Enzyme/EnzymeLogic.cpp             |  1 +
 enzyme/test/Enzyme/ReverseMode/fprintf.ll | 28 +++++++++++++++++++++++
 3 files changed, 39 insertions(+), 1 deletion(-)
 create mode 100644 enzyme/test/Enzyme/ReverseMode/fprintf.ll

diff --git a/enzyme/Enzyme/Enzyme.cpp b/enzyme/Enzyme/Enzyme.cpp
index bae293c04506..10780233cd49 100644
--- a/enzyme/Enzyme/Enzyme.cpp
+++ b/enzyme/Enzyme/Enzyme.cpp
@@ -125,6 +125,14 @@ llvm::cl::opt<std::string> EnzymeTruncateAll(
 #endif
 bool attributeKnownFunctions(llvm::Function &F) {
   bool changed = false;
+  if (F.getName() == "fprintf") {
+    for (auto &arg : F.args()) {
+      if (arg.getType()->isPointerTy()) {
+        arg.addAttr(Attribute::NoCapture);
+        changed = true;
+      }
+    }
+  }
   if (F.getName().contains("__enzyme_float") ||
       F.getName().contains("__enzyme_double") ||
       F.getName().contains("__enzyme_integer") ||
@@ -334,7 +342,8 @@ bool attributeKnownFunctions(llvm::Function &F) {
       F.getName() == "_ZNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEE9_M_"
                      "createERmm" ||
       F.getName() ==
-          "_ZNKSt8__detail20_Prime_rehash_policy14_M_need_rehashEmmm") {
+          "_ZNKSt8__detail20_Prime_rehash_policy14_M_need_rehashEmmm" ||
+      F.getName() == "fprintf") {
     changed = true;
     F.addAttribute(
         AttributeList::FunctionIndex,
diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp
index 96f5bc0e917d..b46ffa25dfbc 100644
--- a/enzyme/Enzyme/EnzymeLogic.cpp
+++ b/enzyme/Enzyme/EnzymeLogic.cpp
@@ -6364,6 +6364,7 @@ llvm::Function *EnzymeLogic::CreateNoFree(RequestContext context, Function *F) {
   };
 
   StringSet<> NoFrees = {"mpfr_greater_p",
+                        "fprintf",
                          "memchr",
                          "time",
                          "strlen",
diff --git a/enzyme/test/Enzyme/ReverseMode/fprintf.ll b/enzyme/test/Enzyme/ReverseMode/fprintf.ll
new file mode 100644
index 000000000000..a152ad2faeb3
--- /dev/null
+++ b/enzyme/test/Enzyme/ReverseMode/fprintf.ll
@@ -0,0 +1,28 @@
+; RUN: if [ %llvmver -ge 15 ]; then %opt < %s %OPnewLoadEnzyme -enzyme-preopt=false -passes="enzyme,function(mem2reg,instsimplify,%simplifycfg)" -S | FileCheck %s; fi
+
+@stdout = external local_unnamed_addr global ptr, align 8
+@.str = private unnamed_addr constant [18 x i8] c" whatever ( %g )\0A\00", align 1
+
+; Function Attrs: noinline nounwind optnone uwtable
+define internal double @dotabs(double %ld) {
+  %r = call i32 (ptr, ptr, ...) @fprintf(ptr @stdout, ptr noundef @.str, double noundef %ld)
+  ret double %ld
+}
+
+; Function Attrs: nofree nounwind
+declare noundef i32 @fprintf(ptr nocapture noundef, ptr nocapture noundef readonly, ...)
+
+; Function Attrs: noinline nounwind optnone uwtable
+define void @f(ptr %x, ptr %dx) {
+  %r = call double (...) @__enzyme_autodiff(ptr @dotabs, double 1.0)
+  ret void
+}
+
+declare double @__enzyme_autodiff(...)
+
+; CHECK: define internal { double } @diffedotabs(double %ld, double %differeturn) 
+; CHECK-NEXT: invert:
+; CHECK-NEXT:   %r = call i32 (ptr, ptr, ...) @fprintf(ptr @stdout, ptr noundef @.str, double noundef %ld)
+; CHECK-NEXT:   %0 = insertvalue { double } {{(undef|poison)}}, double %differeturn, 0
+; CHECK-NEXT:   ret { double } %0
+; CHECK-NEXT: }