Preserve stack alignment (rust-lang#487)

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -8167,14 +8167,26 @@ class AdjointGenerator
         if (Mode == DerivativeMode::ReverseModeGradient) {
           eraseIfUnused(*orig, /*erase*/ true, /*check*/ false);
         } else {
-          if (hasMetadata(orig, "enzyme_fromstack")) {
+          if (auto MD = hasMetadata(orig, "enzyme_fromstack")) {
             IRBuilder<> B(newCall);
             if (auto CI = dyn_cast<ConstantInt>(orig->getArgOperand(0))) {
               B.SetInsertPoint(gutils->inversionAllocs);
             }
             auto replacement = B.CreateAlloca(
                 Type::getInt8Ty(orig->getContext()),
                 gutils->getNewFromOriginal(orig->getArgOperand(0)));
+            auto Alignment =
+                cast<ConstantInt>(
+                    cast<ConstantAsMetadata>(MD->getOperand(0))->getValue())
+                    ->getLimitedValue();
+            // Don't set zero alignment
+            if (Alignment) {
+#if LLVM_VERSION_MAJOR >= 10
+              replacement->setAlignment(Align(Alignment));
+#else
+              replacement->setAlignment(Alignment);
+#endif
+            }
             gutils->replaceAWithB(newCall, replacement);
             gutils->erase(newCall);
           }

enzyme/Enzyme/FunctionUtils.cpp

@@ -312,7 +312,11 @@ static inline void UpgradeAllocasToMallocs(Function *NewF,
     CallInst *CI = dyn_cast<CallInst>(rep);
     if (auto C = dyn_cast<CastInst>(rep))
       CI = cast<CallInst>(C->getOperand(0));
-    CI->setMetadata("enzyme_fromstack", MDNode::get(CI->getContext(), {}));
+    CI->setMetadata("enzyme_fromstack",
+                    MDNode::get(CI->getContext(),
+                                {ConstantAsMetadata::get(ConstantInt::get(
+                                    IntegerType::get(AI->getContext(), 64),
+                                    AI->getAlignment()))}));
 #if LLVM_VERSION_MAJOR >= 14
     CI->addRetAttr(Attribute::NoAlias);
 #else

enzyme/Enzyme/Utils.h

@@ -203,9 +203,9 @@ static inline bool hasMetadata(const llvm::GlobalObject *O,
 }
 
 /// Check if an instruction has metadata
-static inline bool hasMetadata(const llvm::Instruction *O,
-                               llvm::StringRef kind) {
-  return O->getMetadata(kind) != nullptr;
+static inline llvm::MDNode *hasMetadata(const llvm::Instruction *O,
+                                        llvm::StringRef kind) {
+  return O->getMetadata(kind);
 }
 
 /// Potential return type of generated functions

enzyme/test/Enzyme/ReverseMode/stackalign.ll

@@ -0,0 +1,37 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -S | FileCheck %s
+
+define float @f(float %this) {
+entry:
+  %call = tail call float @sub(float %this)
+  %res = fmul float %call, %call
+  ret float %res
+}
+
+declare void @julia.write_barrier(float* readnone nocapture)
+
+define float @sub(float %this)  {
+entry:
+  %alloc = alloca float, align 256
+  store float %this, float* %alloc, align 8
+  call void @julia.write_barrier(float* %alloc)
+  ret float %this
+}
+
+define float @g(float %t) {
+entry:
+  %0 = tail call float (float (float)*, ...) @__enzyme_autodiff(float (float)* @f, float %t)
+  ret float %0
+}
+
+declare float @__enzyme_autodiff(float (float)*, ...)
+
+; ensure both alignment is maintained and that the alloca
+; is not preserved for the reverse
+; CHECK: define internal float @augmented_sub(float %this)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = alloca i8, i64 4, align 256
+; CHECK-NEXT:   %alloc = bitcast i8* %0 to float*
+; CHECK-NEXT:   store float %this, float* %alloc, align 8
+; CHECK-NEXT:   call void @julia.write_barrier(float* %alloc)
+; CHECK-NEXT:   ret float %this
+; CHECK-NEXT: }