[SelectOpt] Fix incorrect IR for SUB when comparison dependent operand is first (#119362)

Author: igogo-x86

llvm/lib/CodeGen/SelectOptimize.cpp

@@ -806,7 +806,12 @@ void SelectOptimizeImpl::collectSelectGroups(BasicBlock &BB,
         break;
       }
 
-      for (unsigned Idx = 0; Idx < 2; Idx++) {
+      // Iterate through operands and find dependant on recognised sign
+      // extending auxiliary select-like instructions. The operand index does
+      // not matter for Add and Or. However, for Sub, we can only safely
+      // transform when the operand is second.
+      unsigned Idx = BO->getOpcode() == Instruction::Sub ? 1 : 0;
+      for (; Idx < 2; Idx++) {
         auto *Op = BO->getOperand(Idx);
         auto It = SelectInfo.find(Op);
         if (It != SelectInfo.end() && It->second.IsAuxiliary) {

llvm/test/CodeGen/AArch64/selectopt-cast.ll

@@ -427,22 +427,16 @@ define void @test_sub_zext_first_op(ptr %dst, ptr %src, i64 %j.start, i64 %p, i6
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[IV1:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[SELECT_END:%.*]] ]
-; CHECK-NEXT:    [[J:%.*]] = phi i64 [ [[J_START:%.*]], [[ENTRY]] ], [ [[J_NEXT:%.*]], [[SELECT_END]] ]
-; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ [[I_START:%.*]], [[ENTRY]] ], [ [[J_NEXT]], [[SELECT_END]] ]
+; CHECK-NEXT:    [[IV1:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[J:%.*]] = phi i64 [ [[J_START:%.*]], [[ENTRY]] ], [ [[J_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ [[I_START:%.*]], [[ENTRY]] ], [ [[J_NEXT]], [[LOOP]] ]
 ; CHECK-NEXT:    [[GEP_DST:%.*]] = getelementptr inbounds ptr, ptr [[DST:%.*]], i64 [[IV]]
 ; CHECK-NEXT:    [[L_I:%.*]] = load ptr, ptr [[GEP_DST]], align 8
 ; CHECK-NEXT:    [[GEP_J:%.*]] = getelementptr inbounds ptr, ptr [[DST]], i64 [[J]]
 ; CHECK-NEXT:    [[L_J:%.*]] = load ptr, ptr [[GEP_J]], align 8
 ; CHECK-NEXT:    [[CMP3:%.*]] = icmp ult ptr [[L_I]], [[L_J]]
 ; CHECK-NEXT:    [[DEC:%.*]] = zext i1 [[CMP3]] to i64
-; CHECK-NEXT:    [[CMP3_FROZEN:%.*]] = freeze i1 [[CMP3]]
-; CHECK-NEXT:    br i1 [[CMP3_FROZEN]], label [[SELECT_TRUE_SINK:%.*]], label [[SELECT_END]]
-; CHECK:       select.true.sink:
-; CHECK-NEXT:    [[TMP0:%.*]] = sub nsw i64 1, [[J]]
-; CHECK-NEXT:    br label [[SELECT_END]]
-; CHECK:       select.end:
-; CHECK-NEXT:    [[J_NEXT]] = phi i64 [ [[TMP0]], [[SELECT_TRUE_SINK]] ], [ [[J]], [[LOOP]] ]
+; CHECK-NEXT:    [[J_NEXT]] = sub nsw i64 [[DEC]], [[J]]
 ; CHECK-NEXT:    [[GEP_DST1:%.*]] = getelementptr inbounds ptr, ptr [[DST1:%.*]], i64 [[IV1]]
 ; CHECK-NEXT:    store i64 [[J_NEXT]], ptr [[GEP_DST1]], align 8
 ; CHECK-NEXT:    [[IV_NEXT]] = add i64 [[IV1]], 1