[SLP] Sort PHIs by ExtractElementInst when relevant

Change-Id: Ife26ac762ee933cfe0b43a4b584caf08c36f0493

Author: jrbyrnes

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

@@ -15742,7 +15742,7 @@ static bool tryToVectorizeSequence(
   // Sort by type, parent, operands.
   stable_sort(Incoming, Comparator);
 
-  // Try to vectorize elements base on their type.
+   // Try to vectorize elements base on their type.
   SmallVector<T *> Candidates;
   for (auto *IncIt = Incoming.begin(), *E = Incoming.end(); IncIt != E;) {
     // Look for the next elements with the same type, parent and operand
@@ -16014,7 +16014,7 @@ bool SLPVectorizerPass::vectorizeChainsInBlock(BasicBlock *BB, BoUpSLP &R) {
           continue;
         return isa<UndefValue>(Opcodes2[I]);
       }
-      if (auto *I1 = dyn_cast<Instruction>(Opcodes1[I]))
+      if (auto *I1 = dyn_cast<Instruction>(Opcodes1[I])) {
         if (auto *I2 = dyn_cast<Instruction>(Opcodes2[I])) {
           DomTreeNodeBase<BasicBlock> *NodeI1 = DT->getNode(I1->getParent());
           DomTreeNodeBase<BasicBlock> *NodeI2 = DT->getNode(I2->getParent());
@@ -16027,11 +16027,47 @@ bool SLPVectorizerPass::vectorizeChainsInBlock(BasicBlock *BB, BoUpSLP &R) {
                  "Different nodes should have different DFS numbers");
           if (NodeI1 != NodeI2)
             return NodeI1->getDFSNumIn() < NodeI2->getDFSNumIn();
+
           InstructionsState S = getSameOpcode({I1, I2}, *TLI);
-          if (S.getOpcode() && !S.isAltShuffle())
+          if (S.getOpcode() && !S.isAltShuffle()) {
+            if (!isa<ExtractElementInst>(I1) || !isa<ExtractElementInst>(I2))
+              continue;
+
+            auto E1 = cast<ExtractElementInst>(I1);
+            auto E2 = cast<ExtractElementInst>(I2);
+            // Sort on ExtractElementInsts primarily by vector operands. Prefer program order of the
+            // vector operands
+            if (E1->getVectorOperand() != E2->getVectorOperand()) {
+              Instruction *V1 = dyn_cast<Instruction>(E1->getVectorOperand());
+              Instruction *V2 = dyn_cast<Instruction>(E2->getVectorOperand());
+              if (!V1 || !V2)
+                continue;
+              if (V1->getParent() != V2->getParent())
+                continue;
+              return V1->comesBefore(V2);
+            }
+            // If we have the same vector operand, try to sort by constant index
+            auto Id1 = E1->getIndexOperand();
+            auto Id2 = E2->getIndexOperand();
+            // Bring constants to the top
+            if (isa<ConstantInt>(Id1) && !isa<ConstantInt>(Id2))
+              return true;
+            if (!isa<ConstantInt>(Id1) && isa<ConstantInt>(Id2))
+              return false;
+            if (isa<ConstantInt>(Id1) && isa<ConstantInt>(Id2)) {
+              auto C1 = cast<ConstantInt>(Id1);
+              auto C2 = cast<ConstantInt>(Id2);
+              // First elements first
+              return C1->getValue().getZExtValue() <
+                      C2->getValue().getZExtValue();
+            }
+
             continue;
+          }
+
           return I1->getOpcode() < I2->getOpcode();
         }
+      }
       if (isa<Constant>(Opcodes1[I]) && isa<Constant>(Opcodes2[I]))
         return Opcodes1[I]->getValueID() < Opcodes2[I]->getValueID();
       if (isa<Instruction>(Opcodes1[I]))
@@ -16120,7 +16156,6 @@ bool SLPVectorizerPass::vectorizeChainsInBlock(BasicBlock *BB, BoUpSLP &R) {
         }
       }
     }
-
     HaveVectorizedPhiNodes = tryToVectorizeSequence<Value>(
         Incoming, PHICompare, AreCompatiblePHIs,
         [this, &R](ArrayRef<Value *> Candidates, bool MaxVFOnly) {

llvm/test/Transforms/SLPVectorizer/AMDGPU/phi-result-use-order.ll

@@ -49,19 +49,19 @@ bb1:
 define <4 x half> @phis_reverse(i1 %cmp1, <4 x half> %in1, <4 x half> %in2)  {
 ; CHECK-LABEL: @phis_reverse(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = shufflevector <4 x half> [[IN1:%.*]], <4 x half> poison, <2 x i32> <i32 2, i32 3>
-; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x half> [[IN1]], <4 x half> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    [[TMP0:%.*]] = shufflevector <4 x half> [[IN1:%.*]], <4 x half> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x half> [[IN1]], <4 x half> poison, <2 x i32> <i32 2, i32 3>
 ; CHECK-NEXT:    br i1 [[CMP1:%.*]], label [[BB1:%.*]], label [[BB0:%.*]]
 ; CHECK:       bb0:
-; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <4 x half> [[IN2:%.*]], <4 x half> poison, <2 x i32> <i32 2, i32 3>
-; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <4 x half> [[IN2]], <4 x half> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <4 x half> [[IN2:%.*]], <4 x half> poison, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <4 x half> [[IN2]], <4 x half> poison, <2 x i32> <i32 2, i32 3>
 ; CHECK-NEXT:    br label [[BB1]]
 ; CHECK:       bb1:
 ; CHECK-NEXT:    [[TMP4:%.*]] = phi <2 x half> [ [[TMP0]], [[ENTRY:%.*]] ], [ [[TMP2]], [[BB0]] ]
 ; CHECK-NEXT:    [[TMP5:%.*]] = phi <2 x half> [ [[TMP1]], [[ENTRY]] ], [ [[TMP3]], [[BB0]] ]
-; CHECK-NEXT:    [[TMP6:%.*]] = shufflevector <2 x half> [[TMP5]], <2 x half> poison, <4 x i32> <i32 0, i32 1, i32 poison, i32 poison>
-; CHECK-NEXT:    [[TMP7:%.*]] = shufflevector <2 x half> [[TMP4]], <2 x half> poison, <4 x i32> <i32 0, i32 1, i32 poison, i32 poison>
-; CHECK-NEXT:    [[TMP8:%.*]] = shufflevector <4 x half> [[TMP6]], <4 x half> [[TMP7]], <4 x i32> <i32 0, i32 1, i32 4, i32 5>
+; CHECK-NEXT:    [[TMP6:%.*]] = shufflevector <2 x half> [[TMP4]], <2 x half> poison, <4 x i32> <i32 0, i32 1, i32 poison, i32 poison>
+; CHECK-NEXT:    [[TMP7:%.*]] = shufflevector <2 x half> [[TMP5]], <2 x half> poison, <4 x i32> <i32 0, i32 1, i32 poison, i32 poison>
+; CHECK-NEXT:    [[TMP8:%.*]] = shufflevector <2 x half> [[TMP4]], <2 x half> [[TMP5]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
 ; CHECK-NEXT:    ret <4 x half> [[TMP8]]
 ;
 entry: