Better errors (rust-lang#588)

* More advanced error handler

* Improved keep minus one

Co-authored-by: William S. Moses <[email protected]>

Author: vchuravy

enzyme/Enzyme/CApi.cpp

@@ -221,6 +221,12 @@ void FreeTypeAnalysis(EnzymeTypeAnalysisRef TAR) {
   delete TA;
 }
 
+void *EnzymeAnalyzeTypes(EnzymeTypeAnalysisRef TAR, CFnTypeInfo CTI,
+                         LLVMValueRef F) {
+  FnTypeInfo FTI(eunwrap(CTI, cast<Function>(unwrap(F))));
+  return (void *)&((TypeAnalysis *)TAR)->analyzeFunction(FTI).analyzer;
+}
+
 void EnzymeRegisterAllocationHandler(char *Name, CustomShadowAlloc AHandle,
                                      CustomShadowFree FHandle) {
   shadowHandlers[std::string(Name)] =
@@ -539,8 +545,36 @@ const char *EnzymeTypeTreeToString(CTypeTreeRef src) {
 
   return cstr;
 }
+
+// TODO deprecated
 void EnzymeTypeTreeToStringFree(const char *cstr) { delete[] cstr; }
 
+const char *EnzymeTypeAnalyzerToString(void *src) {
+  auto TA = (TypeAnalyzer *)src;
+  std::string str;
+  raw_string_ostream ss(str);
+  TA->dump(ss);
+  ss.str();
+  char *cstr = new char[str.length() + 1];
+  std::strcpy(cstr, str.c_str());
+  return cstr;
+}
+
+const char *EnzymeGradientUtilsInvertedPointersToString(GradientUtils *gutils,
+                                                        void *src) {
+  std::string str;
+  raw_string_ostream ss(str);
+  for (auto z : gutils->invertedPointers) {
+    ss << "available inversion for " << *z.first << " of " << *z.second << "\n";
+  }
+  ss.str();
+  char *cstr = new char[str.length() + 1];
+  std::strcpy(cstr, str.c_str());
+  return cstr;
+}
+
+void EnzymeStringFree(const char *cstr) { delete[] cstr; }
+
 void EnzymeMoveBefore(LLVMValueRef inst1, LLVMValueRef inst2) {
   Instruction *I1 = cast<Instruction>(unwrap(inst1));
   Instruction *I2 = cast<Instruction>(unwrap(inst2));

enzyme/Enzyme/EnzymeLogic.cpp

@@ -89,8 +89,6 @@ cl::opt<bool> nonmarkedglobals_inactiveloads(
 cl::opt<bool> EnzymeJuliaAddrLoad(
     "enzyme-julia-addr-load", cl::init(false), cl::Hidden,
     cl::desc("Mark all loads resulting in an addr(13)* to be legal to redo"));
-
-void (*CustomErrorHandler)(const char *) = nullptr;
 }
 
 struct CacheAnalysis {
@@ -1826,7 +1824,8 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
       llvm::raw_string_ostream ss(s);
       ss << "No augmented forward pass found for " + todiff->getName() << "\n";
       ss << *todiff << "\n";
-      CustomErrorHandler(ss.str().c_str());
+      CustomErrorHandler(ss.str().c_str(), wrap(todiff),
+                         ErrorType::NoDerivative, nullptr);
     }
     llvm::errs() << "mod: " << *todiff->getParent() << "\n";
     llvm::errs() << *todiff << "\n";
@@ -3343,7 +3342,8 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
     ss << "No reverse pass found for " + key.todiff->getName() << "\n";
     ss << *key.todiff << "\n";
     if (CustomErrorHandler) {
-      CustomErrorHandler(ss.str().c_str());
+      CustomErrorHandler(ss.str().c_str(), wrap(key.todiff),
+                         ErrorType::NoDerivative, nullptr);
     } else {
       llvm_unreachable(ss.str().c_str());
     }
@@ -3958,7 +3958,8 @@ Function *EnzymeLogic::CreateForwardDiff(
     llvm::raw_string_ostream ss(s);
     ss << "No forward derivative found for " + todiff->getName() << "\n";
     ss << *todiff << "\n";
-    CustomErrorHandler(s.c_str());
+    CustomErrorHandler(s.c_str(), wrap(todiff), ErrorType::NoDerivative,
+                       nullptr);
   }
   if (todiff->empty())
     llvm::errs() << *todiff << "\n";

enzyme/Enzyme/EnzymeLogic.h

@@ -50,7 +50,6 @@
 
 extern "C" {
 extern llvm::cl::opt<bool> EnzymePrint;
-extern void (*CustomErrorHandler)(const char *);
 }
 
 enum class AugmentedStruct { Tape, Return, DifferentialReturn };

enzyme/Enzyme/GradientUtils.cpp

@@ -100,8 +100,6 @@ llvm::cl::opt<bool>
     EnzymeRematerialize("enzyme-rematerialize", cl::init(true), cl::Hidden,
                         cl::desc("Rematerialize allocations/shadows in the "
                                  "reverse rather than caching"));
-
-extern void (*CustomErrorHandler)(const char *);
 }
 
 unsigned int MD_ToCopy[5] = {LLVMContext::MD_dbg, LLVMContext::MD_tbaa,
@@ -4478,11 +4476,11 @@ end:;
   assert(BuilderM.GetInsertBlock()->getParent());
   assert(oval);
 
-  if (CustomErrorHandler && isa<Constant>(oval)) {
+  if (CustomErrorHandler) {
     std::string str;
     raw_string_ostream ss(str);
     ss << "cannot find shadow for " << *oval;
-    CustomErrorHandler(str.c_str());
+    CustomErrorHandler(str.c_str(), wrap(oval), ErrorType::NoShadow, this);
   }
 
   llvm::errs() << *newFunc->getParent() << "\n";

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp

@@ -631,6 +631,17 @@ void TypeAnalyzer::updateAnalysis(Value *Val, TypeTree Data, Value *Origin) {
       Invalid = true;
       return;
     }
+    if (CustomErrorHandler) {
+      std::string str;
+      raw_string_ostream ss(str);
+      ss << "Illegal updateAnalysis prev:" << prev.str()
+         << " new: " << Data.str() << "\n";
+      ss << "val: " << *Val;
+      if (Origin)
+        ss << " origin=" << *Origin;
+      CustomErrorHandler(str.c_str(), wrap(Val), ErrorType::IllegalTypeAnalysis,
+                         (void *)this);
+    }
     llvm::errs() << *fntypeinfo.Function->getParent() << "\n";
     llvm::errs() << *fntypeinfo.Function << "\n";
     dump();
@@ -1384,7 +1395,19 @@ void TypeAnalyzer::visitGetElementPtrInst(GetElementPtrInst &gep) {
     return;
 
   if (direction & DOWN) {
-    updateAnalysis(&gep, pointerAnalysis.KeepMinusOne(), &gep);
+    bool legal = true;
+    auto keepMinus = pointerAnalysis.KeepMinusOne(legal);
+    if (!legal) {
+      if (CustomErrorHandler)
+        CustomErrorHandler("Could not keep minus one", wrap(&gep),
+                           ErrorType::IllegalTypeAnalysis, this);
+      else {
+        dump();
+        llvm::errs() << " could not perform minus one for gep'd: " << gep
+                     << "\n";
+      }
+    }
+    updateAnalysis(&gep, keepMinus, &gep);
     updateAnalysis(&gep, TypeTree(pointerAnalysis.Inner0()).Only(-1), &gep);
   }
   if (direction & UP)
@@ -2114,14 +2137,13 @@ void TypeAnalyzer::visitInsertValueInst(InsertValueInst &I) {
     updateAnalysis(&I, new_res, &I);
 }
 
-void TypeAnalyzer::dump() {
-  llvm::errs() << "<analysis>\n";
+void TypeAnalyzer::dump(llvm::raw_ostream &ss) {
+  ss << "<analysis>\n";
   for (auto &pair : analysis) {
-    llvm::errs() << *pair.first << ": " << pair.second.str()
-                 << ", intvals: " << to_string(knownIntegralValues(pair.first))
-                 << "\n";
+    ss << *pair.first << ": " << pair.second.str()
+       << ", intvals: " << to_string(knownIntegralValues(pair.first)) << "\n";
   }
-  llvm::errs() << "</analysis>\n";
+  ss << "</analysis>\n";
 }
 
 void TypeAnalyzer::visitAtomicRMWInst(llvm::AtomicRMWInst &I) {

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h

@@ -347,7 +347,7 @@ class TypeAnalyzer : public llvm::InstVisitor<TypeAnalyzer> {
 
   TypeTree getReturnAnalysis();
 
-  void dump();
+  void dump(llvm::raw_ostream &ss = llvm::errs());
 
   std::set<int64_t> knownIntegralValues(llvm::Value *val);
 

enzyme/Enzyme/TypeAnalysis/TypeTree.h

@@ -657,7 +657,7 @@ class TypeTree : public std::enable_shared_from_this<TypeTree> {
   }
 
   /// Keep only pointers (or anything's) to a repeated value (represented by -1)
-  TypeTree KeepMinusOne() const {
+  TypeTree KeepMinusOne(bool &legal) const {
     TypeTree dat;
 
     for (const auto &pair : mapping) {
@@ -674,9 +674,8 @@ class TypeTree : public std::enable_shared_from_this<TypeTree> {
           dat.insert(pair.first, pair.second);
           continue;
         }
-        llvm::errs() << "could not merge test  " << str() << "\n";
-        assert(0 && "could not merge");
-        llvm_unreachable("could not merge");
+        legal = false;
+        break;
       }
 
       if (pair.first[1] == -1) {

enzyme/Enzyme/Utils.cpp

@@ -39,6 +39,11 @@
 
 using namespace llvm;
 
+extern "C" {
+void (*CustomErrorHandler)(const char *, LLVMValueRef, ErrorType,
+                           void *) = nullptr;
+}
+
 EnzymeFailure::EnzymeFailure(llvm::StringRef RemarkName,
                              const llvm::DiagnosticLocation &Loc,
                              const llvm::Instruction *CodeRegion)

enzyme/Enzyme/Utils.h

@@ -64,9 +64,17 @@ namespace llvm {
 class ScalarEvolution;
 }
 
+enum class ErrorType {
+  NoDerivative = 0,
+  NoShadow = 1,
+  IllegalTypeAnalysis = 2
+};
+
 extern "C" {
 /// Print additional debug info relevant to performance
 extern llvm::cl::opt<bool> EnzymePrintPerf;
+extern void (*CustomErrorHandler)(const char *, LLVMValueRef, ErrorType,
+                                  void *);
 }
 
 extern std::map<std::string, std::function<llvm::Value *(