Support declaring a custom allocation / free function. (rust-lang#937)

* starting custom allocator

* Fix custom alloc

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -10738,7 +10738,7 @@ class AdjointGenerator
                       bb,
                       [&](Value *anti) {
                         zeroKnownAllocation(bb, anti, args, funcName,
-                                            gutils->TLI);
+                                            gutils->TLI, orig);
                       },
                       anti);
                 }
@@ -10760,7 +10760,7 @@ class AdjointGenerator
             assert(tofree->getType());
             auto rule = [&](Value *tofree) {
               auto CI = freeKnownAllocation(Builder2, tofree, funcName, dbgLoc,
-                                            gutils->TLI);
+                                            gutils->TLI, orig, gutils);
               if (CI)
 #if LLVM_VERSION_MAJOR >= 14
                 CI->addAttributeAtIndex(AttributeList::FirstArgIndex,
@@ -10978,7 +10978,7 @@ class AdjointGenerator
               getReverseBuilder(Builder2);
               auto dbgLoc = gutils->getNewFromOriginal(orig->getDebugLoc());
               freeKnownAllocation(Builder2, lookup(newCall, Builder2), funcName,
-                                  dbgLoc, gutils->TLI);
+                                  dbgLoc, gutils->TLI, orig, gutils);
               if (Mode == DerivativeMode::ReverseModeGradient &&
                   found->second.LI && found->second.LI->contains(orig))
                 gutils->rematerializedPrimalOrShadowAllocations.push_back(
@@ -11063,7 +11063,7 @@ class AdjointGenerator
           getReverseBuilder(Builder2);
           auto dbgLoc = gutils->getNewFromOriginal(orig->getDebugLoc());
           freeKnownAllocation(Builder2, lookup(nop, Builder2), funcName, dbgLoc,
-                              gutils->TLI);
+                              gutils->TLI, orig, gutils);
         }
       } else if (Mode == DerivativeMode::ReverseModeGradient ||
                  Mode == DerivativeMode::ReverseModeCombined ||

enzyme/Enzyme/Enzyme.cpp

@@ -391,15 +391,15 @@ handleFunctionLike(llvm::Module &M, llvm::GlobalVariable &g,
                         Attribute::get(g.getContext(), "enzyme_math", nameVal));
       } else {
         llvm::errs() << M << "\n";
-        llvm::errs() << "Param of __enzyme_inactivefn must be a "
+        llvm::errs() << "Param of __enzyme_function_like must be a "
                         "function"
                      << g << "\n"
                      << *V << "\n";
         llvm_unreachable("__enzyme_inactivefn");
       }
     } else {
       llvm::errs() << M << "\n";
-      llvm::errs() << "Use of __enzyme_inactivefn must be a "
+      llvm::errs() << "Use of __enzyme_function_like must be a "
                       "constant function "
                    << g << "\n";
       llvm_unreachable("__enzyme_register_gradient");
@@ -408,6 +408,108 @@ handleFunctionLike(llvm::Module &M, llvm::GlobalVariable &g,
   }
 }
 
+static void
+handleAllocationLike(llvm::Module &M, llvm::GlobalVariable &g,
+                     SmallVectorImpl<GlobalVariable *> &globalsToErase) {
+  if (g.hasInitializer()) {
+    if (auto CA = dyn_cast<ConstantAggregate>(g.getInitializer())) {
+      if (CA->getNumOperands() != 4) {
+        llvm::errs() << M << "\n";
+        llvm::errs() << "Use of "
+                     << "enzyme_allocation_like"
+                     << " must be a "
+                        "constant of size at least "
+                     << 4 << " " << g << "\n";
+        llvm_unreachable("enzyme_allocation_like");
+      }
+      Value *V = CA->getOperand(0);
+      Value *name = CA->getOperand(1);
+      while (auto CE = dyn_cast<ConstantExpr>(V)) {
+        V = CE->getOperand(0);
+      }
+      while (auto CE = dyn_cast<ConstantExpr>(name)) {
+        name = CE->getOperand(0);
+      }
+      Value *deallocind = CA->getOperand(2);
+      while (auto CE = dyn_cast<ConstantExpr>(deallocind)) {
+        deallocind = CE->getOperand(0);
+      }
+      Value *deallocfn = CA->getOperand(3);
+      while (auto CE = dyn_cast<ConstantExpr>(deallocfn)) {
+        deallocfn = CE->getOperand(0);
+      }
+      size_t index = 0;
+      if (auto CI = dyn_cast<ConstantInt>(name)) {
+        index = CI->getZExtValue();
+      } else {
+        llvm::errs() << *name << "\n";
+        llvm::errs() << "Use of "
+                     << "enzyme_allocation_like"
+                     << "requires an integer index"
+                     << "\n";
+        llvm_unreachable("enzyme_allocation_like");
+      }
+
+      StringRef deallocIndStr;
+      bool foundInd = false;
+      if (auto GV = dyn_cast<GlobalVariable>(deallocind))
+        if (GV->isConstant())
+          if (auto C = GV->getInitializer())
+            if (auto CA = dyn_cast<ConstantDataArray>(C))
+              if (CA->getType()->getElementType()->isIntegerTy(8) &&
+                  CA->isCString()) {
+                deallocIndStr = CA->getAsCString();
+                foundInd = true;
+              }
+
+      if (!foundInd) {
+        llvm::errs() << *deallocind << "\n";
+        llvm::errs() << "Use of "
+                     << "enzyme_allocation_like"
+                     << "requires a deallocation index string"
+                     << "\n";
+        llvm_unreachable("enzyme_allocation_like");
+      }
+      if (auto F = dyn_cast<Function>(V)) {
+        F->addAttribute(AttributeList::FunctionIndex,
+                        Attribute::get(g.getContext(), "enzyme_allocator",
+                                       std::to_string(index)));
+      } else {
+        llvm::errs() << M << "\n";
+        llvm::errs() << "Param of __enzyme_allocation_like must be a "
+                        "function"
+                     << g << "\n"
+                     << *V << "\n";
+        llvm_unreachable("__enzyme_allocation_like");
+      }
+      cast<Function>(V)->addAttribute(
+          AttributeList::FunctionIndex,
+          Attribute::get(g.getContext(), "enzyme_deallocator", deallocIndStr));
+
+      if (auto F = dyn_cast<Function>(deallocfn)) {
+        cast<Function>(V)->setMetadata(
+            "enzyme_deallocator_fn",
+            llvm::MDTuple::get(F->getContext(),
+                               {llvm::ValueAsMetadata::get(F)}));
+      } else {
+        llvm::errs() << M << "\n";
+        llvm::errs() << "Free fn of __enzyme_allocation_like must be a "
+                        "function"
+                     << g << "\n"
+                     << *deallocfn << "\n";
+        llvm_unreachable("__enzyme_allocation_like");
+      }
+    } else {
+      llvm::errs() << M << "\n";
+      llvm::errs() << "Use of __enzyme_allocation_like must be a "
+                      "constant function "
+                   << g << "\n";
+      llvm_unreachable("__enzyme_allocation_like");
+    }
+    globalsToErase.push_back(&g);
+  }
+}
+
 static void handleKnownFunctions(llvm::Function &F) {
   if (F.getName() == "memcmp") {
     F.addFnAttr(Attribute::ReadOnly);
@@ -2452,6 +2554,8 @@ class Enzyme final : public ModulePass {
         handleInactiveFunction(M, g, globalsToErase);
       } else if (g.getName().contains("__enzyme_function_like")) {
         handleFunctionLike(M, g, globalsToErase);
+      } else if (g.getName().contains("__enzyme_allocation_like")) {
+        handleAllocationLike(M, g, globalsToErase);
       }
     }
     for (auto g : globalsToErase) {