[SROA] Fix alignment for uses of PHI nodes.

[efriedma]

Summary

Splitting an alloca can decrease the alignment of GEPs into the partition. Normally, rewriting accounts for this, but the code was missing for uses of PHI nodes and select instructions.

Fixes https://bugs.llvm.org/show_bug.cgi?id=38707 .

[glaubitz]

This patch fixes a problem with Rust emitting code with causing unaligned access: https://github.com/rust-lang/rust/issues/53181

[davide]

LGTM.

lib/Transforms/Scalar/SROA.cpp

else
       New = IRB.CreateLifetimeEnd(Ptr, Size);
 
     (void)New;
     LLVM_DEBUG(dbgs() << "          to: " << *New << "\n");
 
     return true;
   }
 
+  void fixLoadStoreAlign(Instruction &Root) {
+    // This algorithm implements the same visitor loop as
+    // hasUnsafePHIOrSelectUse, and fixes the alignment of each load
+    // or store found.
+    SmallPtrSet<Instruction *, 4> Visited;
+    SmallVector<Instruction *, 4> Uses;
+    Visited.insert(&Root);
+    Uses.push_back(&Root);
+    do {
+      Instruction *I = Uses.pop_back_val();
+
+      if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
+        unsigned LoadAlign = LI->getAlignment();
+        if (!LoadAlign)
+          LoadAlign = DL.getABITypeAlignment(LI->getType());
+        LI->setAlignment(std::min(LoadAlign, getSliceAlign()));
+        continue;
+      }
+      if (StoreInst *SI = dyn_cast<StoreInst>(I)) {
+        unsigned StoreAlign = SI->getAlignment();
+        if (!StoreAlign) {
+          Value *Op = SI->getOperand(0);
+          StoreAlign = DL.getABITypeAlignment(Op->getType());
+        }
+        SI->setAlignment(std::min(StoreAlign, getSliceAlign()));
+        continue;
+      }
+
+      assert(isa<BitCastInst>(I) || isa<PHINode>(I) ||
+             isa<SelectInst>(I) || isa<GetElementPtrInst>(I));
+      for (User *U : I->users())
+        if (Visited.insert(cast<Instruction>(U)).second)
+          Uses.push_back(cast<Instruction>(U));
+    } while (!Uses.empty());
+  }
+
   bool visitPHINode(PHINode &PN) {
     LLVM_DEBUG(dbgs() << "    original: " << PN << "\n");
     assert(BeginOffset >= NewAllocaBeginOffset && "PHIs are unsplittable");
     assert(EndOffset <= NewAllocaEndOffset && "PHIs are unsplittable");
 
     // We would like to compute a new pointer in only one place, but have it be
     // as local as possible to the PHI. To do that, we re-use the location of
     // the old pointer, which necessarily must be in the right position to
     // dominate the PHI.
     IRBuilderTy PtrBuilder(IRB);
     if (isa<PHINode>(OldPtr))
       PtrBuilder.SetInsertPoint(&*OldPtr->getParent()->getFirstInsertionPt());
     else
       PtrBuilder.SetInsertPoint(OldPtr);
     PtrBuilder.SetCurrentDebugLocation(OldPtr->getDebugLoc());
 
     Value *NewPtr = getNewAllocaSlicePtr(PtrBuilder, OldPtr->getType());
     // Replace the operands which were using the old pointer.
     std::replace(PN.op_begin(), PN.op_end(), cast<Value>(OldPtr), NewPtr);
 
     LLVM_DEBUG(dbgs() << "          to: " << PN << "\n");
     deleteIfTriviallyDead(OldPtr);
 
+    // Fix the alignment of any loads or stores using this PHI node.
+    fixLoadStoreAlign(PN);
+
     // PHIs can't be promoted on their own, but often can be speculated. We
     // check the speculation outside of the rewriter so that we see the
     // fully-rewritten alloca.
     PHIUsers.insert(&PN);
     return true;
   }
 
   bool visitSelectInst(SelectInst &SI) {
     LLVM_DEBUG(dbgs() << "    original: " << SI << "\n");
     assert((SI.getTrueValue() == OldPtr || SI.getFalseValue() == OldPtr) &&
            "Pointer isn't an operand!");
     assert(BeginOffset >= NewAllocaBeginOffset && "Selects are unsplittable");
     assert(EndOffset <= NewAllocaEndOffset && "Selects are unsplittable");
 
     Value *NewPtr = getNewAllocaSlicePtr(IRB, OldPtr->getType());
     // Replace the operands which were using the old pointer.
     if (SI.getOperand(1) == OldPtr)
       SI.setOperand(1, NewPtr);
     if (SI.getOperand(2) == OldPtr)
       SI.setOperand(2, NewPtr);
 
     LLVM_DEBUG(dbgs() << "          to: " << SI << "\n");
     deleteIfTriviallyDead(OldPtr);
 
+    // Fix the alignment of any loads or stores using this select.
+    fixLoadStoreAlign(SI);
+
     // Selects can't be promoted on their own, but often can be speculated. We
     // check the speculation outside of the rewriter so that we see the
     // fully-rewritten alloca.
     SelectUsers.insert(&SI);
     return true;
   }
 };
 

test/Transforms/SROA/phi-and-select.ll

if.then2:                                         ; preds = %if.end
   br label %if.then5
 
 if.then5:                                         ; preds = %if.then2, %if.end
   %f1 = phi %struct.S* [ undef, %if.then2 ], [ %f2, %if.end ]
 ; CHECK: phi {{.*}} %[[cast]]
   store %struct.S undef, %struct.S* %f1, align 4
   ret void
 }
+
+define i32 @phi_align(i32* %z) {
+; CHECK-LABEL: @phi_align(
+entry:
+  %a = alloca [8 x i8], align 8
+; CHECK: alloca [7 x i8]
+
+  %a0x = getelementptr [8 x i8], [8 x i8]* %a, i64 0, i32 1
+  %a0 = bitcast i8* %a0x to i32*
+  %a1x = getelementptr [8 x i8], [8 x i8]* %a, i64 0, i32 4
+  %a1 = bitcast i8* %a1x to i32*
+; CHECK: store i32 0, {{.*}}, align 1
+  store i32 0, i32* %a0, align 1
+; CHECK: store i32 1, {{.*}}, align 1
+  store i32 1, i32* %a1, align 4
+; CHECK: load {{.*}}, align 1
+  %v0 = load i32, i32* %a0, align 1
+; CHECK: load {{.*}}, align 1
+  %v1 = load i32, i32* %a1, align 4
+  %cond = icmp sle i32 %v0, %v1
+  br i1 %cond, label %then, label %exit
+
+then:
+  br label %exit
+
+exit:
+; CHECK: %phi = phi i32* [ {{.*}}, %then ], [ %z, %entry ]
+; CHECK-NEXT: %result = load i32, i32* %phi, align 1
+  %phi = phi i32* [ %a1, %then ], [ %z, %entry ]
+  %result = load i32, i32* %phi, align 4
+  ret i32 %result
+}