build under older LLVM's, handle rust mangling

Author: ZuseZ4

enzyme/Enzyme/GradientUtils.cpp

@@ -127,6 +127,10 @@ llvm::cl::opt<bool>
 llvm::cl::opt<bool>
     EnzymePrintDiffUse("enzyme-print-diffuse", cl::init(false), cl::Hidden,
                        cl::desc("Print differential use analysis"));
+
+llvm::cl::opt<std::string>
+    EnzymeRustDeallocName("rust-dealloc-name", cl::init(""), cl::Hidden,
+                            cl::desc("Name of Rust deallocation function"));
 }
 
 SmallVector<unsigned int, 9> MD_ToCopy = {
@@ -9474,6 +9478,7 @@ llvm::CallInst *freeKnownAllocation(llvm::IRBuilder<> &builder,
                                     GradientUtils *gutils) {
   assert(isAllocationFunction(allocationfn, TLI));
 
+#if LLVM_VERSION_MAJOR >= 17
   std::string demangledName = llvm::demangle(allocationfn);
   if (demangledName == "__rustc::__rust_alloc" || demangledName == "__rustc::__rust_alloc_zeroed" ) {
     Type *VoidTy = Type::getVoidTy(tofree->getContext());
@@ -9482,10 +9487,26 @@ llvm::CallInst *freeKnownAllocation(llvm::IRBuilder<> &builder,
     Type *inTys[3] = {IntPtrTy, RustSz, RustSz};
 
     auto FT = FunctionType::get(VoidTy, inTys, false);
+    if (EnzymeRustDeallocName == "") {
+      // Rust's (de)alloc names aren't stable. We expect rustc to set them 
+      // for us, but if it fails to do so we instead search for it here.
+      for (auto &F : *builder.GetInsertBlock()->getParent()->getParent()) {
+        auto demangledName = llvm::demangle(F.getName());
+        if (demangledName == "__rustc::__rust_dealloc") {
+          EnzymeRustDeallocName = F.getName();
+          break;
+        }
+      }
+      if (EnzymeRustDeallocName == "") {
+        // If we can't find it, use the raw __rust_dealloc as a fallback.
+        // FIXME: Make this a hard error once we pass the right name from rustc.
+        EnzymeRustDeallocName = "__rust_dealloc";
+      }
+    }
     Value *freevalue = builder.GetInsertBlock()
                            ->getParent()
                            ->getParent()
-                           ->getOrInsertFunction("__rust_dealloc", FT)
+                           ->getOrInsertFunction(EnzymeRustDeallocName, FT)
                            .getCallee();
     Value *vals[3];
     vals[0] = builder.CreatePointerCast(tofree, IntPtrTy);
@@ -9509,6 +9530,7 @@ llvm::CallInst *freeKnownAllocation(llvm::IRBuilder<> &builder,
       builder.Insert(freecall);
     return freecall;
   }
+#endif
   if (allocationfn == "julia.gc_alloc_obj" ||
       allocationfn == "jl_gc_alloc_typed" ||
       allocationfn == "ijl_gc_alloc_typed" ||

enzyme/Enzyme/LibraryFuncs.h

@@ -56,9 +56,11 @@ static inline bool isAllocationFunction(const llvm::StringRef name,
     return true;
   if (name == "__size_returning_new_experiment")
     return true;
+#if LLVM_VERSION_MAJOR >= 17
   std::string demangledName = llvm::demangle(name.str());
   if (demangledName == "__rustc::__rust_alloc" || demangledName == "__rustc::__rust_alloc_zeroed")
     return true;
+#endif
   if (name == "julia.gc_alloc_obj" || name == "jl_gc_alloc_typed" ||
       name == "ijl_gc_alloc_typed")
     return true;
@@ -210,8 +212,12 @@ static inline void zeroKnownAllocation(llvm::IRBuilder<> &bb,
   assert(isAllocationFunction(funcName, TLI));
 
   // Don't re-zero an already-zero buffer
+#if LLVM_VERSION_MAJOR >= 17
     std::string demangledName = llvm::demangle(funcName.str());
-  if (funcName == "calloc" || demangledName == "__rustc::__rust_alloc_zeroed")
+  if (demangledName == "__rustc::__rust_alloc_zeroed")
+    return;
+#endif
+  if (funcName == "calloc")
     return;
 
   Value *allocSize = argValues[0];