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

[Igor Kirillov via llvm-commits]

https://github.com/igogo-x86 created https://github.com/llvm/llvm-project/pull/119362

None

>From 29a8c02263030b254277e7f1e7fd23eaefee4235 Mon Sep 17 00:00:00 2001
From: Igor Kirillov <igor.kirillov at arm.com>
Date: Tue, 10 Dec 2024 11:39:59 +0000
Subject: [PATCH] [SelectOpt] Fix incorrect IR for SUB when comparison
 dependent operand is first

---
 llvm/lib/CodeGen/SelectOptimize.cpp         |  7 ++++++-
 llvm/test/CodeGen/AArch64/selectopt-cast.ll | 14 ++++----------
 2 files changed, 10 insertions(+), 11 deletions(-)

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