diff --git a/llvm/include/llvm/IR/Instructions.h b/llvm/include/llvm/IR/Instructions.h
index 88c8c709c306..b6575d4c8572 100644
--- a/llvm/include/llvm/IR/Instructions.h
+++ b/llvm/include/llvm/IR/Instructions.h
@@ -4960,6 +4960,16 @@ inline Align getLoadStoreAlignment(const Value *I) {
   return cast<StoreInst>(I)->getAlign();
 }
 
+/// A helper function that set the alignment of load or store instruction.
+inline void setLoadStoreAlignment(Value *I, Align NewAlign) {
+  assert((isa<LoadInst>(I) || isa<StoreInst>(I)) &&
+         "Expected Load or Store instruction");
+  if (auto *LI = dyn_cast<LoadInst>(I))
+    LI->setAlignment(NewAlign);
+  else
+    cast<StoreInst>(I)->setAlignment(NewAlign);
+}
+
 /// A helper function that returns the address space of the pointer operand of
 /// load or store instruction.
 inline unsigned getLoadStoreAddressSpace(const Value *I) {
diff --git a/llvm/lib/Transforms/Scalar/InferAlignment.cpp b/llvm/lib/Transforms/Scalar/InferAlignment.cpp
index 6e0c206bd198..21d373790ac5 100644
--- a/llvm/lib/Transforms/Scalar/InferAlignment.cpp
+++ b/llvm/lib/Transforms/Scalar/InferAlignment.cpp
@@ -25,21 +25,14 @@ using namespace llvm;
 static bool tryToImproveAlign(
     const DataLayout &DL, Instruction *I,
     function_ref<Align(Value *PtrOp, Align OldAlign, Align PrefAlign)> Fn) {
-  if (auto *LI = dyn_cast<LoadInst>(I)) {
-    Value *PtrOp = LI->getPointerOperand();
-    Align OldAlign = LI->getAlign();
-    Align NewAlign = Fn(PtrOp, OldAlign, DL.getPrefTypeAlign(LI->getType()));
+
+  if (auto *PtrOp = getLoadStorePointerOperand(I)) {
+    Align OldAlign = getLoadStoreAlignment(I);
+    Align PrefAlign = DL.getPrefTypeAlign(getLoadStoreType(I));
+
+    Align NewAlign = Fn(PtrOp, OldAlign, PrefAlign);
     if (NewAlign > OldAlign) {
-      LI->setAlignment(NewAlign);
-      return true;
-    }
-  } else if (auto *SI = dyn_cast<StoreInst>(I)) {
-    Value *PtrOp = SI->getPointerOperand();
-    Value *ValOp = SI->getValueOperand();
-    Align OldAlign = SI->getAlign();
-    Align NewAlign = Fn(PtrOp, OldAlign, DL.getPrefTypeAlign(ValOp->getType()));
-    if (NewAlign > OldAlign) {
-      SI->setAlignment(NewAlign);
+      setLoadStoreAlignment(I, NewAlign);
       return true;
     }
   }