4 files changed, 24 insertions, 1 deletions
diff --git a/llvm/lib/CodeGen/SafeStack.cpp b/llvm/lib/CodeGen/SafeStack.cpp
index 841a5ec000f..4762fbfd35e 100644
--- a/llvm/lib/CodeGen/SafeStack.cpp
+++ b/llvm/lib/CodeGen/SafeStack.cpp
@@ -557,6 +557,7 @@ Value *SafeStack::moveStaticAllocasToUnsafeStack(
 
   for (Argument *Arg : ByValArguments) {
     unsigned Offset = SSL.getObjectOffset(Arg);
+    unsigned Align = SSL.getObjectAlignment(Arg);
     Type *Ty = Arg->getType()->getPointerElementType();
 
     uint64_t Size = DL.getTypeStoreSize(Ty);
@@ -573,7 +574,7 @@ Value *SafeStack::moveStaticAllocasToUnsafeStack(
                       DIExpression::NoDeref, -Offset, DIExpression::NoDeref);
     Arg->replaceAllUsesWith(NewArg);
     IRB.SetInsertPoint(cast<Instruction>(NewArg)->getNextNode());
-    IRB.CreateMemCpy(Off, Arg, Size, Arg->getParamAlignment());
+    IRB.CreateMemCpy(Off, Align, Arg, Arg->getParamAlignment(), Size);
   }
 
   // Allocate space for every unsafe static AllocaInst on the unsafe stack.
diff --git a/llvm/lib/CodeGen/SafeStackLayout.cpp b/llvm/lib/CodeGen/SafeStackLayout.cpp
index b1759359e46..68ccdc27766 100644
--- a/llvm/lib/CodeGen/SafeStackLayout.cpp
+++ b/llvm/lib/CodeGen/SafeStackLayout.cpp
@@ -42,6 +42,7 @@ LLVM_DUMP_METHOD void StackLayout::print(raw_ostream &OS) {
 void StackLayout::addObject(const Value *V, unsigned Size, unsigned Alignment,
                             const StackColoring::LiveRange &Range) {
   StackObjects.push_back({V, Size, Alignment, Range});
+  ObjectAlignments[V] = Alignment;
   MaxAlignment = std::max(MaxAlignment, Alignment);
 }
 
diff --git a/llvm/lib/CodeGen/SafeStackLayout.h b/llvm/lib/CodeGen/SafeStackLayout.h
index 7c1292f251f..ac531d800f6 100644
--- a/llvm/lib/CodeGen/SafeStackLayout.h
+++ b/llvm/lib/CodeGen/SafeStackLayout.h
@@ -47,6 +47,7 @@ class StackLayout {
   SmallVector<StackObject, 8> StackObjects;
 
   DenseMap<const Value *, unsigned> ObjectOffsets;
+  DenseMap<const Value *, unsigned> ObjectAlignments;
 
   void layoutObject(StackObject &Obj);
 
@@ -64,6 +65,9 @@ public:
   /// Returns the offset to the object start in the stack frame.
   unsigned getObjectOffset(const Value *V) { return ObjectOffsets[V]; }
 
+  /// Returns the alignment of the object
+  unsigned getObjectAlignment(const Value *V) { return ObjectAlignments[V]; }
+
   /// Returns the size of the entire frame.
   unsigned getFrameSize() { return Regions.empty() ? 0 : Regions.back().End; }
 
diff --git a/llvm/test/Transforms/SafeStack/X86/byval.ll b/llvm/test/Transforms/SafeStack/X86/byval.ll
index 1bc5a298e78..1b131aadef0 100644
--- a/llvm/test/Transforms/SafeStack/X86/byval.ll
+++ b/llvm/test/Transforms/SafeStack/X86/byval.ll
@@ -33,6 +33,23 @@ entry:
   ret i32 %0
 }
 
+; Unsafe access to a byval argument.
+; Argument is copied to the unsafe stack.
+; Check that dest align of memcpy is set according to datalayout prefered alignment
+define i32 @ByValUnsafe2(%struct.S* byval nocapture readonly %zzz, i64 %idx) norecurse nounwind readonly safestack uwtable {
+entry:
+  ; CHECK-LABEL: @ByValUnsafe
+  ; CHECK: %[[A:.*]] = load {{.*}} @__safestack_unsafe_stack_ptr
+  ; CHECK: store {{.*}} @__safestack_unsafe_stack_ptr
+  ; CHECK: %[[B:.*]] = getelementptr i8, i8* %[[A]], i32 -400
+  ; CHECK: %[[C:.*]] = bitcast %struct.S* %zzz to i8*
+  ; CHECK: call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 %[[B]], i8* %[[C]], i64 400, i1 false)
+  ; CHECK: ret i32
+  %arrayidx = getelementptr inbounds %struct.S, %struct.S* %zzz, i64 0, i32 0, i64 %idx
+  %0 = load i32, i32* %arrayidx, align 4
+  ret i32 %0
+}
+
 ; Highly aligned byval argument.
 define i32 @ByValUnsafeAligned(%struct.S* byval nocapture readonly align 64 %zzz, i64 %idx) norecurse nounwind readonly safestack uwtable {
 entry: