[llvm] Improve selection of conditional branch on amdgcn.ballot!=0 condition in SelectionDAG. (PR #68714)

[via llvm-commits]

================
@@ -13584,6 +13585,56 @@ SDValue SITargetLowering::performClampCombine(SDNode *N,
   return SDValue(CSrc, 0);
 }
 
+SDValue SITargetLowering::performBRCondCombine(SDNode *N,
+                                               DAGCombinerInfo &DCI) const {
+  if (!DCI.isAfterLegalizeDAG())
+    return SDValue(N, 0);
+
+  SDValue Cond = N->getOperand(1);
+  if (Cond.getOpcode() == ISD::SETCC &&
+      Cond->getOperand(0)->getOpcode() == AMDGPUISD::SETCC) {
+
+    // %VCMP = i32/i64 AMDGPUISD::SETCC ...
+    // %C = ISD::SETCC %VCMP, 0, setne/seteq
+    // BRCOND %BB, %C
+    // =>
+    // %VCMP = i32/i64 AMDGPUISD::SETCC ...
+    // BRCONDZ %BB, %VCMP, setne/seteq
+
+    auto CC = cast<CondCodeSDNode>(Cond->getOperand(2))->get();
+    auto *CRHS = dyn_cast<ConstantSDNode>(Cond->getOperand(1));
+    if ((CC == ISD::SETEQ || CC == ISD::SETNE) && CRHS && CRHS->isZero()) {
+
+      auto VCMP = Cond->getOperand(0);
+      auto VCMP_CC = cast<CondCodeSDNode>(VCMP.getOperand(2))->get();
+      auto *VCMP_CRHS = dyn_cast<ConstantSDNode>(VCMP.getOperand(1));
+      auto Src = VCMP;
+      if (VCMP_CC == ISD::SETNE && VCMP_CRHS && VCMP_CRHS->isZero()) {
+
+        // Special case for amdgcn.ballot:
+        // %VCMPSrc = ISD::SETCC or a logical combination of ISD::SETCCs
+        // %VCMP = i32/i64 AMDGPUISD::SETCC (ext %VCMPSrc), 0, setne
+        // %C = ISD::SETCC %VCMP, 0, setne/seteq
+        // BRCOND %BB, %C
+        // =>
+        // BRCONDZ %BB, %VCMPSrc, setne/seteq
----------------
vpykhtin wrote:

On the other thought these two patterns have a node in common:

**%C = ISD::SETCC %VCMP, 0, setne/seteq**
BRCOND %BB, %C
=>
BRCONDZ %BB, %VCMP, setne/seteq

%VCMPSrc = ISD::SETCC or a logical combination of ISD::SETCCs
%VCMP = i32/i64 AMDGPUISD::SETCC (ext %VCMPSrc), 0, setne
**%C = ISD::SETCC %VCMP, 0, setne/seteq**
=>
%C = ISD::SETCC or a logical combination of ISD::SETCCs

SelectionDAG starts with BRCOND of the first pattern and erases the ISD::SETCC which is the root for the second pattern. I tried to perform DAG combine on the %C value of the first rule so that the second rule could perform its transformation and return result to the first rule but this requires to create artificial result _%C = ISD::SETCC %VCMPSrc, 0, setne_ so that the first rule pattern could match. It looks like this should work as a single pattern.

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