[SelectOpt] Fix incorrect IR for SUB when comparison dependent operand is first

[igogo-x86]

No description provided.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Igor Kirillov (igogo-x86)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/119362.diff

2 Files Affected:

• (modified) llvm/lib/CodeGen/SelectOptimize.cpp (+6-1)
• (modified) llvm/test/CodeGen/AArch64/selectopt-cast.ll (+4-10)

diff --git a/llvm/lib/CodeGen/SelectOptimize.cpp b/llvm/lib/CodeGen/SelectOptimize.cpp
index 484705eabbc42e..2ec5af2d864f77 100644
--- a/llvm/lib/CodeGen/SelectOptimize.cpp
+++ b/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) {
diff --git a/llvm/test/CodeGen/AArch64/selectopt-cast.ll b/llvm/test/CodeGen/AArch64/selectopt-cast.ll
index 102b89df32b03b..4afb6af6e757cf 100644
--- a/llvm/test/CodeGen/AArch64/selectopt-cast.ll
+++ b/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

[fhahn]

LGTM, thanks!

[igogo-x86]

Alternatively, we can add support for SUB(zext(), other_op), but in that case, we'll have to generate sub instructions in both branches, and I doubt it has the potential to give performance improvements. Not speaking about how uncommon this operation might be.