[LSV] Introduce copyMetadataForStore

Author: gandhi56

llvm/include/llvm/Transforms/Utils/Local.h

@@ -430,6 +430,10 @@ void combineAAMetadata(Instruction *K, const Instruction *J);
 /// replacement for the source instruction).
 void copyMetadataForLoad(LoadInst &Dest, const LoadInst &Source);
 
+/// Copy the metadata from the source instruction to the destination (the
+/// replace)
+void copyMetadataForStore(StoreInst &Dest, const StoreInst &Source);
+
 /// Patch the replacement so that it is not more restrictive than the value
 /// being replaced. It assumes that the replacement does not get moved from
 /// its original position.

llvm/lib/Transforms/Utils/Local.cpp

@@ -3510,6 +3510,46 @@ void llvm::copyMetadataForLoad(LoadInst &Dest, const LoadInst &Source) {
   }
 }
 
+void llvm::copyMetadataForStore(StoreInst &Dest, const StoreInst &Source) {
+  SmallVector<std::pair<unsigned, MDNode *>, 8> MD;
+  Source.getAllMetadata(MD);
+  MDBuilder MDB(Dest.getContext());
+  Type *NewType = Dest.getType();
+  for (const auto &MDPair : MD) {
+    unsigned ID = MDPair.first;
+    MDNode *N = MDPair.second;
+    switch (ID) {
+    case LLVMContext::MD_dbg:
+    case LLVMContext::MD_tbaa:
+    case LLVMContext::MD_prof:
+    case LLVMContext::MD_tbaa_struct:
+    case LLVMContext::MD_alias_scope:
+    case LLVMContext::MD_noalias:
+    case LLVMContext::MD_nontemporal:
+    case LLVMContext::MD_access_group:
+    case LLVMContext::MD_noundef:
+    case LLVMContext::MD_noalias_addrspace:
+    case LLVMContext::MD_mem_parallel_loop_access:
+      Dest.setMetadata(ID, N);
+      break;
+    
+    case LLVMContext::MD_nonnull:
+      break;
+
+    case LLVMContext::MD_align:
+    case LLVMContext::MD_dereferenceable:
+    case LLVMContext::MD_dereferenceable_or_null:
+      // These only directly apply if the new type is also a pointer.
+      if (NewType->isPointerTy())
+        Dest.setMetadata(ID, N);
+      break;
+
+    case LLVMContext::MD_range:
+      break;
+    }
+  }
+}
+
 void llvm::patchReplacementInstruction(Instruction *I, Value *Repl) {
   auto *ReplInst = dyn_cast<Instruction>(Repl);
   if (!ReplInst)

llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

@@ -1333,14 +1333,6 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
   if (EQClasses.size() < 2)
     return;
 
-  auto CopyMetaDataFromTo = [&](Instruction *Src, Instruction *Dst) {
-    SmallVector<std::pair<unsigned, MDNode *>, 4> MD;
-    Src->getAllMetadata(MD);
-    for (const auto [ID, Node] : MD) {
-      Dst->setMetadata(ID, Node);
-    }
-  };
-
   // For each class, determine the most defined type. This information will
   // help us determine the type instructions should be casted into.
   MapVector<EqClassKey, unsigned> ClassToNewTyID;
@@ -1405,7 +1397,7 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
           auto *Cast = Builder.CreateBitOrPointerCast(
               NewLoad, OrigTy, NewLoad->getName() + ".cast");
           LI->replaceAllUsesWith(Cast);
-          CopyMetaDataFromTo(LI, NewLoad);
+          copyMetadataForLoad(*NewLoad, *LI);
           LI->eraseFromParent();
           EQClasses[EC1.first].emplace_back(NewLoad);
         } else {
@@ -1416,7 +1408,7 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
               SI->getValueOperand()->getName() + ".cast");
           auto *NewStore = Builder.CreateStore(
               Cast, getLoadStorePointerOperand(SI), SI->isVolatile());
-          CopyMetaDataFromTo(SI, NewStore);
+          copyMetadataForStore(*NewStore, *SI);
           SI->eraseFromParent();
           EQClasses[EC1.first].emplace_back(NewStore);
         }