[DAGCombiner] Inverse transform (select c, (and X, 1), 0) -> (and (zext c), X)

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -6299,6 +6299,33 @@ static SDValue foldAndOrOfSETCC(SDNode *LogicOp, SelectionDAG &DAG) {
   return SDValue();
 }
 
+// Combine `(select c, (X & 1), 0)` -> `(and (zext c), X)`.
+// We canonicalize to the `select` form in the middle end, but the `and` form
+// gets better codegen and all tested targets (arm, x86, riscv)
+static SDValue combineSelectAsExtAnd(SDValue Cond, SDValue T, SDValue F,
+                                     const SDLoc &DL, SelectionDAG &DAG) {
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+  if (!isNullConstant(F))
+    return SDValue();
+
+  EVT CondVT = Cond.getValueType();
+  if (TLI.getBooleanContents(CondVT) !=
+      TargetLoweringBase::ZeroOrOneBooleanContent)
+    return SDValue();
+
+  if (T.getOpcode() != ISD::AND)
+    return SDValue();
+
+  if (!isOneConstant(T.getOperand(1)))
+    return SDValue();
+
+  EVT OpVT = T.getValueType();
+
+  SDValue CondMask =
+      OpVT == CondVT ? Cond : DAG.getBoolExtOrTrunc(Cond, DL, OpVT, CondVT);
+  return DAG.getNode(ISD::AND, DL, OpVT, CondMask, T.getOperand(0));
+}
+
 /// This contains all DAGCombine rules which reduce two values combined by
 /// an And operation to a single value. This makes them reusable in the context
 /// of visitSELECT(). Rules involving constants are not included as
@@ -11609,6 +11636,9 @@ SDValue DAGCombiner::visitSELECT(SDNode *N) {
     if (SDValue BinOp = foldSelectOfBinops(N))
       return BinOp;
 
+  if (SDValue R = combineSelectAsExtAnd(N0, N1, N2, DL, DAG))
+    return R;
+
   return SDValue();
 }
 

llvm/test/CodeGen/AArch64/select-to-and-zext.ll

@@ -0,0 +1,88 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-none-elf -verify-machineinstrs %s -o - | FileCheck %s
+
+define i32 @from_cmpeq(i32 %xx, i32 %y) {
+; CHECK-LABEL: from_cmpeq:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and w8, w1, #0x1
+; CHECK-NEXT:    cmp w0, #9
+; CHECK-NEXT:    csel w0, w8, wzr, eq
+; CHECK-NEXT:    ret
+  %x = icmp eq i32 %xx, 9
+  %masked = and i32 %y, 1
+
+  %r = select i1 %x, i32 %masked, i32 0
+  ret i32 %r
+}
+
+define i32 @from_cmpeq_fail_bad_andmask(i32 %xx, i32 %y) {
+; CHECK-LABEL: from_cmpeq_fail_bad_andmask:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and w8, w1, #0x3
+; CHECK-NEXT:    cmp w0, #9
+; CHECK-NEXT:    csel w0, w8, wzr, eq
+; CHECK-NEXT:    ret
+  %x = icmp eq i32 %xx, 9
+  %masked = and i32 %y, 3
+  %r = select i1 %x, i32 %masked, i32 0
+  ret i32 %r
+}
+
+define i32 @from_i1(i1 %x, i32 %y) {
+; CHECK-LABEL: from_i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and w8, w0, w1
+; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ret
+  %masked = and i32 %y, 1
+  %r = select i1 %x, i32 %masked, i32 0
+  ret i32 %r
+}
+
+define i32 @from_trunc_i8(i8 %xx, i32 %y) {
+; CHECK-LABEL: from_trunc_i8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and w8, w0, w1
+; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ret
+  %masked = and i32 %y, 1
+  %x = trunc i8 %xx to i1
+  %r = select i1 %x, i32 %masked, i32 0
+  ret i32 %r
+}
+
+define i32 @from_trunc_i64(i64 %xx, i32 %y) {
+; CHECK-LABEL: from_trunc_i64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and w8, w0, w1
+; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ret
+  %masked = and i32 %y, 1
+  %x = trunc i64 %xx to i1
+  %r = select i1 %x, i32 %masked, i32 0
+  ret i32 %r
+}
+
+define i32 @from_i1_fail_bad_select0(i1 %x, i32 %y) {
+; CHECK-LABEL: from_i1_fail_bad_select0:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and w8, w1, #0x1
+; CHECK-NEXT:    tst w0, #0x1
+; CHECK-NEXT:    csinc w0, w8, wzr, ne
+; CHECK-NEXT:    ret
+  %masked = and i32 %y, 1
+  %r = select i1 %x, i32 %masked, i32 1
+  ret i32 %r
+}
+
+define i32 @from_i1_fail_bad_select1(i1 %x, i32 %y) {
+; CHECK-LABEL: from_i1_fail_bad_select1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and w8, w1, #0x1
+; CHECK-NEXT:    tst w0, #0x1
+; CHECK-NEXT:    csel w0, wzr, w8, ne
+; CHECK-NEXT:    ret
+  %masked = and i32 %y, 1
+  %r = select i1 %x, i32 0, i32 %masked
+  ret i32 %r
+}

llvm/test/CodeGen/RISCV/select-to-and-zext.ll

@@ -0,0 +1,152 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=riscv32 -verify-machineinstrs < %s \
+; RUN:   | FileCheck %s -check-prefix=RV32I
+; RUN: llc -mtriple=riscv64 -verify-machineinstrs < %s \
+; RUN:   | FileCheck %s -check-prefix=RV64I
+
+
+define i32 @from_cmpeq(i32 %xx, i32 %y) {
+; RV32I-LABEL: from_cmpeq:
+; RV32I:       # %bb.0:
+; RV32I-NEXT:    addi a0, a0, -9
+; RV32I-NEXT:    seqz a0, a0
+; RV32I-NEXT:    and a0, a0, a1
+; RV32I-NEXT:    ret
+;
+; RV64I-LABEL: from_cmpeq:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    sext.w a0, a0
+; RV64I-NEXT:    addi a0, a0, -9
+; RV64I-NEXT:    seqz a0, a0
+; RV64I-NEXT:    and a0, a0, a1
+; RV64I-NEXT:    ret
+  %x = icmp eq i32 %xx, 9
+  %masked = and i32 %y, 1
+
+  %r = select i1 %x, i32 %masked, i32 0
+  ret i32 %r
+}
+
+define i32 @from_cmpeq_fail_bad_andmask(i32 %xx, i32 %y) {
+; RV32I-LABEL: from_cmpeq_fail_bad_andmask:
+; RV32I:       # %bb.0:
+; RV32I-NEXT:    addi a0, a0, -9
+; RV32I-NEXT:    snez a0, a0
+; RV32I-NEXT:    addi a0, a0, -1
+; RV32I-NEXT:    and a0, a1, a0
+; RV32I-NEXT:    andi a0, a0, 3
+; RV32I-NEXT:    ret
+;
+; RV64I-LABEL: from_cmpeq_fail_bad_andmask:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    sext.w a0, a0
+; RV64I-NEXT:    addi a0, a0, -9
+; RV64I-NEXT:    snez a0, a0
+; RV64I-NEXT:    addiw a0, a0, -1
+; RV64I-NEXT:    and a0, a1, a0
+; RV64I-NEXT:    andi a0, a0, 3
+; RV64I-NEXT:    ret
+  %x = icmp eq i32 %xx, 9
+  %masked = and i32 %y, 3
+  %r = select i1 %x, i32 %masked, i32 0
+  ret i32 %r
+}
+
+define i32 @from_i1(i1 %x, i32 %y) {
+; RV32I-LABEL: from_i1:
+; RV32I:       # %bb.0:
+; RV32I-NEXT:    and a0, a0, a1
+; RV32I-NEXT:    andi a0, a0, 1
+; RV32I-NEXT:    ret
+;
+; RV64I-LABEL: from_i1:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    and a0, a0, a1
+; RV64I-NEXT:    andi a0, a0, 1
+; RV64I-NEXT:    ret
+  %masked = and i32 %y, 1
+  %r = select i1 %x, i32 %masked, i32 0
+  ret i32 %r
+}
+
+define i32 @from_trunc_i8(i8 %xx, i32 %y) {
+; RV32I-LABEL: from_trunc_i8:
+; RV32I:       # %bb.0:
+; RV32I-NEXT:    and a0, a0, a1
+; RV32I-NEXT:    andi a0, a0, 1
+; RV32I-NEXT:    ret
+;
+; RV64I-LABEL: from_trunc_i8:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    and a0, a0, a1
+; RV64I-NEXT:    andi a0, a0, 1
+; RV64I-NEXT:    ret
+  %masked = and i32 %y, 1
+  %x = trunc i8 %xx to i1
+  %r = select i1 %x, i32 %masked, i32 0
+  ret i32 %r
+}
+
+define i32 @from_trunc_i64(i64 %xx, i32 %y) {
+; RV32I-LABEL: from_trunc_i64:
+; RV32I:       # %bb.0:
+; RV32I-NEXT:    and a0, a0, a2
+; RV32I-NEXT:    andi a0, a0, 1
+; RV32I-NEXT:    ret
+;
+; RV64I-LABEL: from_trunc_i64:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    and a0, a0, a1
+; RV64I-NEXT:    andi a0, a0, 1
+; RV64I-NEXT:    ret
+  %masked = and i32 %y, 1
+  %x = trunc i64 %xx to i1
+  %r = select i1 %x, i32 %masked, i32 0
+  ret i32 %r
+}
+
+define i32 @from_i1_fail_bad_select0(i1 %x, i32 %y) {
+; RV32I-LABEL: from_i1_fail_bad_select0:
+; RV32I:       # %bb.0:
+; RV32I-NEXT:    andi a0, a0, 1
+; RV32I-NEXT:    bnez a0, .LBB5_2
+; RV32I-NEXT:  # %bb.1:
+; RV32I-NEXT:    li a0, 1
+; RV32I-NEXT:    ret
+; RV32I-NEXT:  .LBB5_2:
+; RV32I-NEXT:    andi a0, a1, 1
+; RV32I-NEXT:    ret
+;
+; RV64I-LABEL: from_i1_fail_bad_select0:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    andi a0, a0, 1
+; RV64I-NEXT:    bnez a0, .LBB5_2
+; RV64I-NEXT:  # %bb.1:
+; RV64I-NEXT:    li a0, 1
+; RV64I-NEXT:    ret
+; RV64I-NEXT:  .LBB5_2:
+; RV64I-NEXT:    andi a0, a1, 1
+; RV64I-NEXT:    ret
+  %masked = and i32 %y, 1
+  %r = select i1 %x, i32 %masked, i32 1
+  ret i32 %r
+}
+
+define i32 @from_i1_fail_bad_select1(i1 %x, i32 %y) {
+; RV32I-LABEL: from_i1_fail_bad_select1:
+; RV32I:       # %bb.0:
+; RV32I-NEXT:    addi a0, a0, -1
+; RV32I-NEXT:    and a0, a1, a0
+; RV32I-NEXT:    andi a0, a0, 1
+; RV32I-NEXT:    ret
+;
+; RV64I-LABEL: from_i1_fail_bad_select1:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    addiw a0, a0, -1
+; RV64I-NEXT:    and a0, a1, a0
+; RV64I-NEXT:    andi a0, a0, 1
+; RV64I-NEXT:    ret
+  %masked = and i32 %y, 1
+  %r = select i1 %x, i32 0, i32 %masked
+  ret i32 %r
+}