[llvm] [SelectOptimize] Reject Not as BinOp SelectLike aux (fix #226417) (PR #226539)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: Atharva Ajmera (atharvaajmera)

<details>
<summary>Changes</summary>

In `SelectOptimizeImpl::collectSelectGroups`, binary operations (`Add`, `Sub`, `Or`) can be treated as `SelectLike` if one of their operands is an auxiliary instruction associated with a condition.

However, `getTrueOrFalseValue()` only supports `ZExt`, `SExt`, `LShr`, and `AShr` auxiliaries when materializing the true/false values for a `BinOp` `SelectLike`. Other auxiliaries—notably `not` (`xor %c, true`), which is tracked as an auxiliary for grouping inverted selects—cannot be materialized this way. When a binary operation such as `add (xor %c, true), (lshr %v, 0)` picked the `xor` auxiliary, the pass formed an invalid group and later hit an assertion failure in `getTrueOrFalseValue()` during conversion (`Unexpected opcode`).

This patch checks that the auxiliary operand is a `ZExtInst`, `SExtInst`, `LShrOperator`, or `AShrOperator` before forming a `BinOp` `SelectLike`, skipping other auxiliary kinds so no invalid groups are created.

Fixes #<!-- -->226417.

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


2 Files Affected:

- (modified) llvm/lib/CodeGen/SelectOptimize.cpp (+5) 
- (added) llvm/test/CodeGen/AArch64/selectopt-issue226417.ll (+81) 


``````````diff
diff --git a/llvm/lib/CodeGen/SelectOptimize.cpp b/llvm/lib/CodeGen/SelectOptimize.cpp
index ac2ebcc620051..95c4dbad4a4a8 100644
--- a/llvm/lib/CodeGen/SelectOptimize.cpp
+++ b/llvm/lib/CodeGen/SelectOptimize.cpp
@@ -884,6 +884,11 @@ void SelectOptimizeImpl::collectSelectGroups(BasicBlock &BB,
         auto *Op = I->getOperand(Idx);
         auto It = SelectInfo.find(Op);
         if (It != SelectInfo.end() && It->second.IsAuxiliary) {
+          // getTrueOrFalseValue only supports ZExt, SExt, and Shift auxiliaries for BinOps.
+          auto *OpI = dyn_cast<Instruction>(Op);
+          if (!OpI || (!isa<ZExtInst>(OpI) && !isa<SExtInst>(OpI) &&
+                       !isa<LShrOperator>(OpI) && !isa<AShrOperator>(OpI)))
+            continue;
           Cond = It->second.Cond;
           bool Inverted = It->second.IsInverted;
           return SelectInfo.insert({I, {Cond, false, Inverted, Idx}}).first;
diff --git a/llvm/test/CodeGen/AArch64/selectopt-issue226417.ll b/llvm/test/CodeGen/AArch64/selectopt-issue226417.ll
new file mode 100644
index 0000000000000..a4e7dd70b256a
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/selectopt-issue226417.ll
@@ -0,0 +1,81 @@
+; RUN: opt -passes='require<profile-summary>,function(select-optimize)' -mtriple=aarch64-linux-gnu -mcpu=generic -S < %s | FileCheck %s
+
+; Test that xor auxiliaries are not grouped into BinOp SelectLikes (#226417).
+
+define void @not_xor_add_not_selectlike(i32 %n, i1 %c, i1 %v, i32 %a, i32 %b) {
+; CHECK-LABEL: @not_xor_add_not_selectlike(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[OUTER:%.*]]
+; CHECK:       outer:
+; CHECK:       latch:
+; CHECK-NEXT:    [[NC:%.*]] = xor i1 [[C:%.*]], true
+; CHECK-NEXT:    [[SH:%.*]] = lshr i1 [[V:%.*]], 0
+; CHECK-NEXT:    [[S0:%.*]] = select i1 [[C]], i32 [[A:%.*]], i32 [[B:%.*]]
+; CHECK-NEXT:    [[ADD:%.*]] = add i1 [[NC]], [[SH]]
+; CHECK-NEXT:    [[S2:%.*]] = select i1 [[C]], i32 [[A]], i32 [[B]]
+; CHECK-NOT: select.end
+; CHECK-NOT: freeze
+;
+entry:
+  br label %outer
+
+outer:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %latch ]
+  br label %inner
+
+inner:
+  %j = phi i32 [ 0, %outer ], [ %j.next, %inner ]
+  %j.next = add i32 %j, 1
+  %ic = icmp slt i32 %j.next, %n
+  br i1 %ic, label %inner, label %latch
+
+latch:
+  %nc = xor i1 %c, true
+  %sh = lshr i1 %v, 0
+  %s0 = select i1 %c, i32 %a, i32 %b
+  %add = add i1 %nc, %sh
+  %s2 = select i1 %c, i32 %a, i32 %b
+  %i.next = add i32 %i, 1
+  %oc = icmp slt i32 %i.next, %n
+  br i1 %oc, label %outer, label %exit
+
+exit:
+  ret void
+}
+
+define void @not_xor_sub_not_selectlike(i32 %n, i1 %c, i1 %v, i32 %a, i32 %b) {
+; CHECK-LABEL: @not_xor_sub_not_selectlike(
+; CHECK:       latch:
+; CHECK-NEXT:    [[NC:%.*]] = xor i1 [[C:%.*]], true
+; CHECK-NEXT:    [[SH:%.*]] = lshr i1 [[V:%.*]], 0
+; CHECK-NEXT:    [[SUB:%.*]] = sub i1 [[SH]], [[NC]]
+; CHECK-NEXT:    [[S0:%.*]] = select i1 [[C]], i32 [[A:%.*]], i32 [[B:%.*]]
+; CHECK-NEXT:    [[S2:%.*]] = select i1 [[C]], i32 [[A]], i32 [[B]]
+; CHECK-NOT: select.end
+;
+entry:
+  br label %outer
+
+outer:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %latch ]
+  br label %inner
+
+inner:
+  %j = phi i32 [ 0, %outer ], [ %j.next, %inner ]
+  %j.next = add i32 %j, 1
+  %ic = icmp slt i32 %j.next, %n
+  br i1 %ic, label %inner, label %latch
+
+latch:
+  %nc = xor i1 %c, true
+  %sh = lshr i1 %v, 0
+  %sub = sub i1 %sh, %nc
+  %s0 = select i1 %c, i32 %a, i32 %b
+  %s2 = select i1 %c, i32 %a, i32 %b
+  %i.next = add i32 %i, 1
+  %oc = icmp slt i32 %i.next, %n
+  br i1 %oc, label %outer, label %exit
+
+exit:
+  ret void
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/226539