@@ -1319,23 +1319,36 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
13191319 if (EQClasses.size () < 2 )
13201320 return ;
13211321
1322+ auto CopyMetaDataFromTo = [&](Instruction *Src, Instruction *Dst) {
1323+ SmallVector<std::pair<unsigned , MDNode *>, 4 > MD;
1324+ Src->getAllMetadata (MD);
1325+ for (const auto [ID, Node] : MD) {
1326+ Dst->setMetadata (ID, Node);
1327+ }
1328+ };
1329+
13221330 // For each class, determine if all instructions are of type int, FP or ptr.
13231331 // This information will help us determine the type instructions should be
13241332 // casted into.
13251333 MapVector<EqClassKey, Bitset<3 >> ClassAllTy;
1326- for (auto C : EQClasses) {
1327- if (all_of (EQClasses[C.first ], [](Instruction *I) {
1328- return I->getType ()->isIntOrIntVectorTy ();
1329- }))
1330- ClassAllTy[C.first ].set (0 );
1331- else if (all_of (EQClasses[C.first ], [](Instruction *I) {
1332- return I->getType ()->isFPOrFPVectorTy ();
1333- }))
1334- ClassAllTy[C.first ].set (1 );
1335- else if (all_of (EQClasses[C.first ], [](Instruction *I) {
1336- return I->getType ()->isPtrOrPtrVectorTy ();
1337- }))
1338- ClassAllTy[C.first ].set (2 );
1334+ for (const auto &C : EQClasses) {
1335+ auto CommonTypeKind = [](Instruction *I) {
1336+ if (I->getType ()->isIntOrIntVectorTy ())
1337+ return 0 ;
1338+ if (I->getType ()->isFPOrFPVectorTy ())
1339+ return 1 ;
1340+ if (I->getType ()->isPtrOrPtrVectorTy ())
1341+ return 2 ;
1342+ return -1 ; // Invalid type kind
1343+ };
1344+
1345+ int FirstTypeKind = CommonTypeKind (EQClasses[C.first ][0 ]);
1346+ if (FirstTypeKind != -1 &&
1347+ all_of (EQClasses[C.first ], [&](Instruction *I) {
1348+ return CommonTypeKind (I) == FirstTypeKind;
1349+ })) {
1350+ ClassAllTy[C.first ].set (FirstTypeKind);
1351+ }
13391352 }
13401353
13411354 // Loop over all equivalence classes and try to merge them. Keep track of
@@ -1359,6 +1372,11 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
13591372 if (Ptr1 != Ptr2 || AS1 != AS2 || IsLoad1 != IsLoad2 || TySize1 < TySize2)
13601373 continue ;
13611374
1375+ // An All-FP class should only be merged into another All-FP class.
1376+ if ((ClassAllTy[EC1.first ].test (1 ) && !ClassAllTy[EC2.first ].test (1 )) ||
1377+ (!ClassAllTy[EC1.first ].test (2 ) && ClassAllTy[EC2.first ].test (2 )))
1378+ continue ;
1379+
13621380 // Ensure all instructions in EC2 can be bitcasted into NewTy.
13631381 // / TODO: NewTyBits is needed as stuctured binded variables cannot be
13641382 // / captured by a lambda until C++20.
@@ -1384,8 +1402,8 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
13841402 }
13851403
13861404 for (auto *Inst : EC2.second ) {
1387- auto *Ptr = getLoadStorePointerOperand (Inst);
1388- auto *OrigTy = Inst->getType ();
1405+ Value *Ptr = getLoadStorePointerOperand (Inst);
1406+ Type *OrigTy = Inst->getType ();
13891407 if (OrigTy == NewTy)
13901408 continue ;
13911409 if (auto *LI = dyn_cast<LoadInst>(Inst)) {
@@ -1404,6 +1422,7 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
14041422 SI->getValueOperand ()->getName () + " .cast" );
14051423 auto *NewStore = Builder.CreateStore (
14061424 Cast, getLoadStorePointerOperand (SI), SI->isVolatile ());
1425+ CopyMetaDataFromTo (SI, NewStore);
14071426 SI->eraseFromParent ();
14081427 EQClasses[EC1.first ].emplace_back (NewStore);
14091428 }
@@ -1413,7 +1432,7 @@ void Vectorizer::insertCastsToMergeClasses(EquivalenceClassMap &EQClasses) {
14131432 // basic block. This is important to ensure that the instructions are
14141433 // vectorized in the correct order.
14151434 std::sort (EQClasses[EC1.first ].begin (), EQClasses[EC1.first ].end (),
1416- [](Instruction *A, Instruction *B) {
1435+ [](const Instruction *A, const Instruction *B) {
14171436 return A && B && A->comesBefore (B);
14181437 });
14191438 ClassesToErase.insert (EC2.first );
0 commit comments