Add fromstack utility (rust-lang#776)

* Add fromstack utility

* Add addrspace fromstack

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -10379,8 +10379,7 @@ class AdjointGenerator
                                                      Attribute::NonNull);
 #endif
 
-                  if (called->getName() == "malloc" ||
-                      called->getName() == "_Znwm") {
+                  if (funcName == "malloc" || funcName == "_Znwm") {
                     if (auto ci = dyn_cast<ConstantInt>(args[0])) {
                       unsigned derefBytes = ci->getLimitedValue();
                       CallInst *cal =
@@ -10438,18 +10437,37 @@ class AdjointGenerator
                     bb, anti, getIndex(orig, CacheType::Shadow));
               else {
                 if (auto MD = hasMetadata(orig, "enzyme_fromstack")) {
-                  AllocaInst *replacement = bb.CreateAlloca(
-                      Type::getInt8Ty(orig->getContext()), args[0]);
+                  Value *Size;
+                  if (funcName == "malloc")
+                    Size = args[0];
+                  else if (funcName == "julia.gc_alloc_obj")
+                    Size = args[1];
+                  else
+                    llvm_unreachable("Unknown allocation to upgrade");
+                  Value *replacement = bb.CreateAlloca(
+                      Type::getInt8Ty(orig->getContext()), Size);
                   replacement->takeName(anti);
                   auto Alignment = cast<ConstantInt>(cast<ConstantAsMetadata>(
                                                          MD->getOperand(0))
                                                          ->getValue())
                                        ->getLimitedValue();
 #if LLVM_VERSION_MAJOR >= 10
-                  replacement->setAlignment(Align(Alignment));
+                  cast<AllocaInst>(replacement)->setAlignment(Align(Alignment));
 #else
-                  replacement->setAlignment(Alignment);
+                  cast<AllocaInst>(replacement)->setAlignment(Alignment);
 #endif
+                  if (!anti->getType()->getPointerElementType()->isIntegerTy(8))
+                    replacement = bb.CreatePointerCast(
+                        replacement,
+                        PointerType::getUnqual(
+                            anti->getType()->getPointerElementType()));
+
+                  if (int AS =
+                          cast<PointerType>(anti->getType())->getAddressSpace())
+                    replacement = bb.CreateAddrSpaceCast(
+                        replacement,
+                        PointerType::get(
+                            anti->getType()->getPointerElementType(), AS));
 
                   gutils->replaceAWithB(cast<Instruction>(anti), replacement);
                   gutils->erase(cast<Instruction>(anti));
@@ -10583,23 +10601,43 @@ class AdjointGenerator
             // allocation where possible.
             if (auto MD = hasMetadata(orig, "enzyme_fromstack")) {
               IRBuilder<> B(newCall);
-              if (auto CI = dyn_cast<ConstantInt>(orig->getArgOperand(0))) {
+
+              Value *Size;
+              if (funcName == "malloc")
+                Size = orig->getArgOperand(0);
+              else if (funcName == "julia.gc_alloc_obj")
+                Size = orig->getArgOperand(1);
+              else
+                llvm_unreachable("Unknown allocation to upgrade");
+              Size = gutils->getNewFromOriginal(Size);
+
+              if (auto CI = dyn_cast<ConstantInt>(Size)) {
                 B.SetInsertPoint(gutils->inversionAllocs);
               }
-
               auto rule = [&]() {
-                auto replacement = B.CreateAlloca(
-                    Type::getInt8Ty(orig->getContext()),
-                    gutils->getNewFromOriginal(orig->getArgOperand(0)));
+                Value *replacement =
+                    B.CreateAlloca(Type::getInt8Ty(orig->getContext()), Size);
                 auto Alignment =
                     cast<ConstantInt>(
                         cast<ConstantAsMetadata>(MD->getOperand(0))->getValue())
                         ->getLimitedValue();
 #if LLVM_VERSION_MAJOR >= 10
-                replacement->setAlignment(Align(Alignment));
-#else
-                replacement->setAlignment(Alignment);
-#endif
+                cast<AllocaInst>(replacement)->setAlignment(Align(Alignment));
+#else
+                cast<AllocaInst>(replacement)->setAlignment(Alignment);
+#endif
+                if (!orig->getType()->getPointerElementType()->isIntegerTy(8))
+                  replacement = B.CreatePointerCast(
+                      replacement,
+                      PointerType::getUnqual(
+                          orig->getType()->getPointerElementType()));
+
+                if (int AS =
+                        cast<PointerType>(orig->getType())->getAddressSpace())
+                  replacement = B.CreateAddrSpaceCast(
+                      replacement,
+                      PointerType::get(orig->getType()->getPointerElementType(),
+                                       AS));
                 return replacement;
               };
 

enzyme/Enzyme/CApi.cpp

@@ -600,6 +600,11 @@ void EnzymeSetMustCache(LLVMValueRef inst1) {
   I1->setMetadata("enzyme_mustcache", MDNode::get(I1->getContext(), {}));
 }
 
+uint8_t EnzymeHasFromStack(LLVMValueRef inst1) {
+  Instruction *I1 = cast<Instruction>(unwrap(inst1));
+  return hasMetadata(I1, "enzyme_fromstack") != 0;
+}
+
 void EnzymeReplaceFunctionImplementation(LLVMModuleRef M) {
   ReplaceFunctionImplementation(*unwrap(M));
 }

enzyme/Enzyme/FunctionUtils.cpp

@@ -265,6 +265,116 @@ static inline bool OnlyUsedInOMP(AllocaInst *AI) {
   return true;
 }
 
+void RecursivelyReplaceAddressSpace(Value *AI, Value *rep, bool legal) {
+  SmallVector<std::tuple<Value *, Value *, Instruction *>, 1> Todo;
+  for (auto U : AI->users()) {
+    Todo.push_back(
+        std::make_tuple((Value *)rep, (Value *)AI, cast<Instruction>(U)));
+  }
+  SmallVector<Instruction *, 1> toErase;
+  if (auto I = dyn_cast<Instruction>(AI))
+    toErase.push_back(I);
+  while (Todo.size()) {
+    auto cur = Todo.back();
+    Todo.pop_back();
+    Value *rep = std::get<0>(cur);
+    Value *prev = std::get<1>(cur);
+    Value *inst = std::get<2>(cur);
+    if (auto ASC = dyn_cast<AddrSpaceCastInst>(inst)) {
+      auto AS = cast<PointerType>(rep->getType())->getAddressSpace();
+      if (AS == ASC->getDestAddressSpace()) {
+        ASC->replaceAllUsesWith(rep);
+        continue;
+      }
+      ASC->setOperand(0, rep);
+      continue;
+    }
+    if (auto CI = dyn_cast<CastInst>(inst)) {
+      if (!CI->getType()->isPointerTy()) {
+        CI->setOperand(0, rep);
+        continue;
+      }
+      IRBuilder<> B(CI);
+      auto nCI = cast<CastInst>(B.CreateCast(
+          CI->getOpcode(), rep,
+          PointerType::get(
+              CI->getType()->getPointerElementType(),
+              cast<PointerType>(rep->getType())->getAddressSpace())));
+      nCI->takeName(CI);
+      for (auto U : CI->users()) {
+        Todo.push_back(
+            std::make_tuple((Value *)nCI, (Value *)CI, cast<Instruction>(U)));
+      }
+      continue;
+    }
+    if (auto GEP = dyn_cast<GetElementPtrInst>(inst)) {
+      IRBuilder<> B(GEP);
+      SmallVector<Value *, 1> ind(GEP->indices());
+#if LLVM_VERSION_MAJOR > 7
+      auto nGEP = cast<GetElementPtrInst>(
+          B.CreateGEP(GEP->getSourceElementType(), rep, ind));
+#else
+      auto nGEP = cast<GetElementPtrInst>(B.CreateGEP(rep, ind));
+#endif
+      nGEP->takeName(GEP);
+      for (auto U : GEP->users()) {
+        Todo.push_back(
+            std::make_tuple((Value *)nGEP, (Value *)GEP, cast<Instruction>(U)));
+      }
+      toErase.push_back(GEP);
+      continue;
+    }
+    if (auto LI = dyn_cast<LoadInst>(inst)) {
+      LI->setOperand(0, rep);
+      continue;
+    }
+    if (auto SI = dyn_cast<StoreInst>(inst)) {
+      if (SI->getPointerOperand() == prev) {
+        SI->setOperand(1, rep);
+        continue;
+      }
+    }
+    if (auto MS = dyn_cast<MemSetInst>(inst)) {
+      IRBuilder<> B(MS);
+
+      Value *nargs[] = {rep, MS->getArgOperand(1), MS->getArgOperand(2),
+                        MS->getArgOperand(3)};
+
+      Type *tys[] = {nargs[0]->getType(), nargs[2]->getType()};
+
+      auto nMS = cast<CallInst>(B.CreateCall(
+          Intrinsic::getDeclaration(MS->getParent()->getParent()->getParent(),
+                                    Intrinsic::memset, tys),
+          nargs));
+      nMS->copyIRFlags(MS);
+      toErase.push_back(MS);
+      continue;
+    }
+    if (auto CI = dyn_cast<CallInst>(inst)) {
+      if (auto F = CI->getCalledFunction()) {
+        if (F->getName() == "julia.write_barrier" && legal) {
+          toErase.push_back(CI);
+          continue;
+        }
+      }
+    }
+    if (legal) {
+      IRBuilder<> B(cast<Instruction>(rep)->getNextNode());
+      rep = B.CreateAddrSpaceCast(
+          rep, PointerType::get(
+                   rep->getType()->getPointerElementType(),
+                   cast<PointerType>(prev->getType())->getAddressSpace()));
+      prev->replaceAllUsesWith(rep);
+      continue;
+    }
+    llvm::errs() << " rep: " << *rep << " prev: " << *prev << " inst: " << *inst
+                 << "\n";
+    llvm_unreachable("Illegal address space propagation");
+  }
+  for (auto I : llvm::reverse(toErase))
+    I->eraseFromParent();
+}
+
 /// Convert necessary stack allocations into mallocs for use in the reverse
 /// pass. Specifically if we're not topLevel all allocations must be upgraded
 /// Even if topLevel any allocations that aren't in the entry block (and
@@ -312,13 +422,13 @@ static inline void UpgradeAllocasToMallocs(Function *NewF,
 
     auto PT0 = cast<PointerType>(rep->getType());
     auto PT1 = cast<PointerType>(AI->getType());
-    if (PT0->getAddressSpace() != PT1->getAddressSpace())
-      rep = B.CreateAddrSpaceCast(rep,
-                                  PointerType::get(PT0->getPointerElementType(),
-                                                   PT1->getAddressSpace()));
-    assert(rep->getType() == AI->getType());
-    AI->replaceAllUsesWith(rep);
-    AI->eraseFromParent();
+    if (PT0->getAddressSpace() != PT1->getAddressSpace()) {
+      RecursivelyReplaceAddressSpace(AI, rep, /*legal*/ false);
+    } else {
+      assert(rep->getType() == AI->getType());
+      AI->replaceAllUsesWith(rep);
+      AI->eraseFromParent();
+    }
   }
 }
 

enzyme/Enzyme/FunctionUtils.h

@@ -346,6 +346,9 @@ static inline void calculateUnusedStores(
   }
 }
 
+void RecursivelyReplaceAddressSpace(llvm::Value *AI, llvm::Value *rep,
+                                    bool legal);
+
 void ReplaceFunctionImplementation(llvm::Module &M);
 
 /// Is the use of value val as an argument of call CI potentially captured