[AArch64] Peek through freeze in setcc if it has one_use, and is comparing against a constant

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -13492,6 +13492,27 @@ SDValue DAGCombiner::visitSELECT_CC(SDNode *N) {
   ISD::CondCode CC = cast<CondCodeSDNode>(N4)->get();
   SDLoc DL(N);
 
+  // Peek through freeze if it has one use, following the same pattern as SETCC
+  // optimization. This is safe for comparisons that are not "always true" or
+  // "always false" when the operand is poison.
+  if (N0->getOpcode() == ISD::FREEZE && N0.hasOneUse()) {
+    // Check if this comparison could be "always true" or "always false" when
+    // the operand is poison, which would make peeking through freeze unsafe.
+    ConstantSDNode *N1C = dyn_cast<ConstantSDNode>(N1);
+    if (!N1C || !IsAlwaysTrueOrFalseForFreeze(CC, N1C)) {
+      N0 = N0->getOperand(0);
+    }
+  }
+  if (N1->getOpcode() == ISD::FREEZE && N1.hasOneUse()) {
+    // Check if this comparison could be "always true" or "always false" when
+    // the operand is poison, which would make peeking through freeze unsafe.
+    ConstantSDNode *N0C = dyn_cast<ConstantSDNode>(N0);
+    if (!N0C ||
+        !IsAlwaysTrueOrFalseForFreeze(ISD::getSetCCSwappedOperands(CC), N0C)) {
+      N1 = N1->getOperand(0);
+    }
+  }
+
   // fold select_cc lhs, rhs, x, x, cc -> x
   if (N2 == N3)
     return N2;
@@ -13544,6 +13565,46 @@ SDValue DAGCombiner::visitSETCC(SDNode *N) {
   SDValue N0 = N->getOperand(0), N1 = N->getOperand(1);
   SDLoc DL(N);
 
+  // Is 'X Cond C' always true or false?
+  auto IsAlwaysTrueOrFalse = [](ISD::CondCode Cond, ConstantSDNode *C) {
+    bool False = (Cond == ISD::SETULT && C->isZero()) ||
+                 (Cond == ISD::SETLT && C->isMinSignedValue()) ||
+                 (Cond == ISD::SETUGT && C->isAllOnes()) ||
+                 (Cond == ISD::SETGT && C->isMaxSignedValue());
+    bool True = (Cond == ISD::SETULE && C->isAllOnes()) ||
+                (Cond == ISD::SETLE && C->isMaxSignedValue()) ||
+                (Cond == ISD::SETUGE && C->isZero()) ||
+                (Cond == ISD::SETGE && C->isMinSignedValue());
+    return True || False;
+  };
+
+  // Peek through freeze if it has one use. This is safe for comparisons that
+  // are not "always true" or "always false" when the operand is poison. Unsafe
+  // cases include:
+  // - X < 0 (SETULT), X >= 0 (SETUGE) - always false/true when X is poison
+  // - X < MIN_SIGNED (SETLT), X >= MIN_SIGNED (SETGE) - always false/true when
+  // X is poison
+  // - X > MAX_UNSIGNED (SETUGT), X <= MAX_UNSIGNED (SETULE) - always false/true
+  // when X is poison
+  // - X > MAX_SIGNED (SETGT), X <= MAX_SIGNED (SETLE) - always false/true when
+  // X is poison
+  if (N0->getOpcode() == ISD::FREEZE && N0.hasOneUse()) {
+    // Check if this comparison could be "always true" or "always false" when
+    // the operand is poison, which would make peeking through freeze unsafe.
+    ConstantSDNode *N1C = dyn_cast<ConstantSDNode>(N1);
+    if (!N1C || !IsAlwaysTrueOrFalse(Cond, N1C)) {
+      N0 = N0->getOperand(0);
+    }
+  }
+  if (N1->getOpcode() == ISD::FREEZE && N1.hasOneUse()) {
+    // Check if this comparison could be "always true" or "always false" when
+    // the operand is poison, which would make peeking through freeze unsafe.
+    ConstantSDNode *N0C = dyn_cast<ConstantSDNode>(N0);
+    if (!N0C || !IsAlwaysTrueOrFalse(ISD::getSetCCSwappedOperands(Cond), N0C)) {
+      N1 = N1->getOperand(0);
+    }
+  }
+
   if (SDValue Combined = SimplifySetCC(VT, N0, N1, Cond, DL, !PreferSetCC)) {
     // If we prefer to have a setcc, and we don't, we'll try our best to
     // recreate one using rebuildSetCC.
@@ -19287,56 +19348,6 @@ SDValue DAGCombiner::visitBRCOND(SDNode *N) {
                        N1->getOperand(0), N2, N->getFlags());
   }
 
-  // Variant of the previous fold where there is a SETCC in between:
-  //   BRCOND(SETCC(FREEZE(X), CONST, Cond))
-  // =>
-  //   BRCOND(FREEZE(SETCC(X, CONST, Cond)))
-  // =>
-  //   BRCOND(SETCC(X, CONST, Cond))
-  // This is correct if FREEZE(X) has one use and SETCC(FREEZE(X), CONST, Cond)
-  // isn't equivalent to true or false.
-  // For example, SETCC(FREEZE(X), -128, SETULT) cannot be folded to
-  // FREEZE(SETCC(X, -128, SETULT)) because X can be poison.
-  if (N1->getOpcode() == ISD::SETCC && N1.hasOneUse()) {
-    SDValue S0 = N1->getOperand(0), S1 = N1->getOperand(1);
-    ISD::CondCode Cond = cast<CondCodeSDNode>(N1->getOperand(2))->get();
-    ConstantSDNode *S0C = dyn_cast<ConstantSDNode>(S0);
-    ConstantSDNode *S1C = dyn_cast<ConstantSDNode>(S1);
-    bool Updated = false;
-
-    // Is 'X Cond C' always true or false?
-    auto IsAlwaysTrueOrFalse = [](ISD::CondCode Cond, ConstantSDNode *C) {
-      bool False = (Cond == ISD::SETULT && C->isZero()) ||
-                   (Cond == ISD::SETLT && C->isMinSignedValue()) ||
-                   (Cond == ISD::SETUGT && C->isAllOnes()) ||
-                   (Cond == ISD::SETGT && C->isMaxSignedValue());
-      bool True = (Cond == ISD::SETULE && C->isAllOnes()) ||
-                  (Cond == ISD::SETLE && C->isMaxSignedValue()) ||
-                  (Cond == ISD::SETUGE && C->isZero()) ||
-                  (Cond == ISD::SETGE && C->isMinSignedValue());
-      return True || False;
-    };
-
-    if (S0->getOpcode() == ISD::FREEZE && S0.hasOneUse() && S1C) {
-      if (!IsAlwaysTrueOrFalse(Cond, S1C)) {
-        S0 = S0->getOperand(0);
-        Updated = true;
-      }
-    }
-    if (S1->getOpcode() == ISD::FREEZE && S1.hasOneUse() && S0C) {
-      if (!IsAlwaysTrueOrFalse(ISD::getSetCCSwappedOperands(Cond), S0C)) {
-        S1 = S1->getOperand(0);
-        Updated = true;
-      }
-    }
-
-    if (Updated)
-      return DAG.getNode(
-          ISD::BRCOND, SDLoc(N), MVT::Other, Chain,
-          DAG.getSetCC(SDLoc(N1), N1->getValueType(0), S0, S1, Cond), N2,
-          N->getFlags());
-  }
-
   // If N is a constant we could fold this into a fallthrough or unconditional
   // branch. However that doesn't happen very often in normal code, because
   // Instcombine/SimplifyCFG should have handled the available opportunities.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -3604,19 +3604,25 @@ static SDValue emitComparison(SDValue LHS, SDValue RHS, ISD::CondCode CC,
     Opcode = AArch64ISD::ADDS;
     LHS = LHS.getOperand(1);
   } else if (isNullConstant(RHS) && !isUnsignedIntSetCC(CC)) {
-    if (LHS.getOpcode() == ISD::AND) {
+    SDValue AndOp = LHS;
+    // Peek through freeze to find AND operation
+    if (LHS.getOpcode() == ISD::FREEZE && LHS.hasOneUse()) {
+      AndOp = LHS.getOperand(0);
+    }
+
+    if (AndOp.getOpcode() == ISD::AND) {
       // Similarly, (CMP (and X, Y), 0) can be implemented with a TST
       // (a.k.a. ANDS) except that the flags are only guaranteed to work for one
       // of the signed comparisons.
       const SDValue ANDSNode =
           DAG.getNode(AArch64ISD::ANDS, DL, DAG.getVTList(VT, FlagsVT),
-                      LHS.getOperand(0), LHS.getOperand(1));
-      // Replace all users of (and X, Y) with newly generated (ands X, Y)
-      DAG.ReplaceAllUsesWith(LHS, ANDSNode);
+                      AndOp.getOperand(0), AndOp.getOperand(1));
+      // Replace all users of the AND operation with newly generated (ands X, Y)
+      DAG.ReplaceAllUsesWith(AndOp, ANDSNode);
       return ANDSNode.getValue(1);
-    } else if (LHS.getOpcode() == AArch64ISD::ANDS) {
+    } else if (AndOp.getOpcode() == AArch64ISD::ANDS) {
       // Use result of ANDS
-      return LHS.getValue(1);
+      return AndOp.getValue(1);
     }
   }
 

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

@@ -5203,19 +5203,25 @@ MachineInstr *AArch64InstructionSelector::tryFoldIntegerCompare(
   // Produce this if the compare is signed:
   //
   // tst x, y
-  if (!CmpInst::isUnsigned(P) && LHSDef &&
-      LHSDef->getOpcode() == TargetOpcode::G_AND) {
-    // Make sure that the RHS is 0.
-    auto ValAndVReg = getIConstantVRegValWithLookThrough(RHS.getReg(), MRI);
-    if (!ValAndVReg || ValAndVReg->Value != 0)
-      return nullptr;
+  if (!CmpInst::isUnsigned(P) && LHSDef) {
+    MachineInstr *AndDef = LHSDef;
+    // Peek through freeze to find AND operation
+    if (LHSDef->getOpcode() == TargetOpcode::G_FREEZE &&
+        MRI.hasOneUse(LHSDef->getOperand(0).getReg())) {
+      AndDef = getDefIgnoringCopies(LHSDef->getOperand(0).getReg(), MRI);
+    }
 
-    return emitTST(LHSDef->getOperand(1),
-                   LHSDef->getOperand(2), MIRBuilder);
-  }
+    if (AndDef && AndDef->getOpcode() == TargetOpcode::G_AND) {
+      // Make sure that the RHS is 0.
+      auto ValAndVReg = getIConstantVRegValWithLookThrough(RHS.getReg(), MRI);
+      if (!ValAndVReg || ValAndVReg->Value != 0)
+        return nullptr;
 
-  return nullptr;
-}
+      return emitTST(AndDef->getOperand(1), AndDef->getOperand(2), MIRBuilder);
+    }
+
+    return nullptr;
+  }
 
 bool AArch64InstructionSelector::selectShuffleVector(
     MachineInstr &I, MachineRegisterInfo &MRI) {

llvm/test/CodeGen/AArch64/icmp-ult-eq-fold.ll

@@ -161,6 +161,25 @@ define i1 @lt64_u16_and_23(i64 %0) {
   ret i1 %3
 }
 
+define i1 @test_disjoint(i1 %0, i32 %1, i32 %2) {
+; CHECK-LABEL: test_disjoint:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    orr w8, w2, #0x800000
+; CHECK-NEXT:    lsr w8, w8, w1
+; CHECK-NEXT:    tst w8, #0x1
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    orr w8, w0, w8
+; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ret
+entry:
+  %3 = or disjoint i32 %2, 8388608
+  %4 = shl nuw i32 1, %1
+  %5 = and i32 %3, %4
+  %6 = icmp eq i32 %5, 0
+  %7 = select i1 %0, i1 true, i1 %6
+  ret i1 %7
+}
+
 ; negative test
 define i1 @lt3_u8(i8 %0) {
 ; CHECK-LABEL: lt3_u8: