Fix tests

Author: gandhi56

llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

@@ -1319,26 +1319,35 @@ 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 if all instructions are of type int, FP or ptr.
   // This information will help us determine the type instructions should be
   // casted into.
   MapVector<EqClassKey, Bitset<3>> ClassAllTy;
-  for (auto C : EQClasses) {
-    if (all_of(EQClasses[C.first], 
-      [](Instruction *I) {
-        return I->getType()->isIntOrIntVectorTy();
-      }))
-      ClassAllTy[C.first].set(0);
-    else if (all_of(EQClasses[C.first],
-      [](Instruction *I) {
-        return I->getType()->isFPOrFPVectorTy();
-      }))
-      ClassAllTy[C.first].set(1);
-    else if (all_of(EQClasses[C.first],
-      [](Instruction *I) {
-        return I->getType()->isPtrOrPtrVectorTy();
-      }))
-      ClassAllTy[C.first].set(2);
+  for (const auto &C : EQClasses) {
+    auto CommonTypeKind = [](Instruction *I) {
+      if (I->getType()->isIntOrIntVectorTy())
+        return 0;
+      if (I->getType()->isFPOrFPVectorTy())
+        return 1;
+      if (I->getType()->isPtrOrPtrVectorTy())
+        return 2;
+      return -1; // Invalid type kind
+    };
+
+    int FirstTypeKind = CommonTypeKind(EQClasses[C.first][0]);
+    if (FirstTypeKind != -1 && all_of(EQClasses[C.first], [&](Instruction *I) {
+          return CommonTypeKind(I) == FirstTypeKind;
+        })) {
+      ClassAllTy[C.first].set(FirstTypeKind);
+    }
   }
 
   // Loop over all equivalence classes and try to merge them. Keep track of
@@ -1362,6 +1371,11 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
       if (Ptr1 != Ptr2 || AS1 != AS2 || IsLoad1 != IsLoad2 || TySize1 < TySize2)
         continue;
 
+      // An All-FP class should only be merged into another All-FP class.
+      if ((ClassAllTy[EC1.first].test(1) && !ClassAllTy[EC2.first].test(1)) ||
+          (!ClassAllTy[EC1.first].test(2) && ClassAllTy[EC2.first].test(2)))
+        continue;
+
       // Ensure all instructions in EC2 can be bitcasted into NewTy.
       /// TODO: NewTyBits is needed as stuctured binded variables cannot be
       /// captured by a lambda until C++20.
@@ -1381,13 +1395,14 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
           NewTy = Type::getFloatTy(Ctx);
         else if (NewTyBits == 64)
           NewTy = Type::getDoubleTy(Ctx);
-      } else if (ClassAllTy[EC1.first].test(2) && ClassAllTy[EC2.first].test(2)) {
+      } else if (ClassAllTy[EC1.first].test(2) &&
+                 ClassAllTy[EC2.first].test(2)) {
         NewTy = PointerType::get(Ctx, AS2);
       }
 
       for (auto *Inst : EC2.second) {
-        auto *Ptr = getLoadStorePointerOperand(Inst);
-        auto *OrigTy = Inst->getType();
+        Value *Ptr = getLoadStorePointerOperand(Inst);
+        Type *OrigTy = Inst->getType();
         if (OrigTy == NewTy)
           continue;
         if (auto *LI = dyn_cast<LoadInst>(Inst)) {
@@ -1406,6 +1421,7 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
               SI->getValueOperand()->getName() + ".cast");
           auto *NewStore = Builder.CreateStore(
               Cast, getLoadStorePointerOperand(SI), SI->isVolatile());
+          CopyMetaDataFromTo(SI, NewStore);
           SI->eraseFromParent();
           EQClasses[EC1.first].emplace_back(NewStore);
         }
@@ -1415,7 +1431,7 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
       // basic block. This is important to ensure that the instructions are
       // vectorized in the correct order.
       std::sort(EQClasses[EC1.first].begin(), EQClasses[EC1.first].end(),
-                [](Instruction *A, Instruction *B) {
+                [](const Instruction *A, const Instruction *B) {
                   return A && B && A->comesBefore(B);
                 });
       ClassesToErase.insert(EC2.first);