[AArch64] Enhance 'isBitfieldPositioningOp' to find pattern (shl(and(val,mask), N)

[mingmingl]

Summary

Before this patch (and D135844)

• Given DAG node shl(op, N), isBitfieldPositioningOp uses (optionally shifted)op as the Src (least significant bits of Src are inserted into DstLSB of Dst node)

After this patch

• If op is and(val, mask), isBitfieldPositioningOp tries to see through and and find if val is a simpler source than op.

The motivation is similar (probably symmetric) to isSeveralBitsExtractOpFromShr for isBitfieldExtractOpFromShr

Existing test cases are improved without regressions.

[dmgreen]

This seems similar to a demand bits fold. Can it check that AndMask is a superset of AllOnes >> Shift?

[mingmingl]

Address comments.

[mingmingl]

In D135850#3857910, @dmgreen wrote:

This seems similar to a demand bits fold. Can it check that AndMask is a superset of AllOnes >> Shift?

Checking AndMask is a superset of AllOnes >> Shift makes sense. Done by checking isMask_64((AndImm << Shift) >> Shift).

• This patch previously introduced a bug (fixed in the latest reply) by checking isShiftedMask_64((AndImm << Shift) >> Shift).

Also make refactor patch a parent of this one (bypassing D135844).

[dmgreen]

LGTM

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

 // opcode and that it has a immediate integer right operand.
 // If so Imm will receive the 32 bit value.
 static bool isOpcWithIntImmediate(const SDNode *N, unsigned Opc,
                                   uint64_t &Imm) {
   return N->getOpcode() == Opc &&
          isIntImmediate(N->getOperand(1).getNode(), Imm);
 }
 
+// isIntImmediateEq - This method tests to see if N is a constant operand that
+// is equivalent to 'ImmExpected'.
+static bool isIntImmediateEq(SDValue N, const uint64_t ImmExpected) {
+  uint64_t Imm;
+  if (!isIntImmediate(N.getNode(), Imm))
+    return false;
+  return Imm == ImmExpected;
+}
+
 bool AArch64DAGToDAGISel::SelectInlineAsmMemoryOperand(
     const SDValue &Op, unsigned ConstraintID, std::vector<SDValue> &OutOps) {
   switch(ConstraintID) {
   default:
     llvm_unreachable("Unexpected asm memory constraint");
   case InlineAsm::Constraint_m:
   case InlineAsm::Constraint_o:
   case InlineAsm::Constraint_Q:
static bool isBitfieldPositioningOpFromAnd(SelectionDAG *CurDAG, SDValue Op,
   // which case it is not profitable to insert an extra shift.
   if (ShlImm != DstLSB && !BiggerPattern)
     return false;
 
   Src = getLeftShift(CurDAG, AndOp0.getOperand(0), ShlImm - DstLSB);
   return true;
 }
 
+// For node (shl (and val, mask), N)), returns true if the node is equivalent to
+// UBFIZ.
+static bool isSeveralBitsPositioningOpFromShl(const uint64_t ShlImm, SDValue Op,
+                                              SDValue &Src, int &DstLSB,
+                                              int &Width) {
+  // Caller should have verified that N is a left shift with constant shift
+  // amount; asserts that.
+  assert(Op.getOpcode() == ISD::SHL &&
+         "Op.getNode() should be a SHL node to call this function");
+  assert(isIntImmediateEq(Op.getOperand(1), ShlImm) &&
+         "Op.getNode() should shift ShlImm to call this function");
+
+  uint64_t AndImm = 0;
+  SDValue Op0 = Op.getOperand(0);
+  if (!isOpcWithIntImmediate(Op0.getNode(), ISD::AND, AndImm))
+    return false;
+
+  const uint64_t ShiftedAndImm = ((AndImm << ShlImm) >> ShlImm);
+  if (isMask_64(ShiftedAndImm)) {
+    // AndImm is a superset of (AllOnes >> ShlImm); in other words, AndImm
+    // should end with Mask, and could be prefixed with random bits if those
+    // bits are shifted out.
+    //
+    // For example, xyz11111 (with {x,y,z} being 0 or 1) is fine if ShlImm >= 3;
+    // the AND result corresponding to those bits are shifted out, so it's fine
+    // to not extract them.
+    Width = countTrailingOnes(ShiftedAndImm);
+    DstLSB = ShlImm;
+    Src = Op0.getOperand(0);
+    return true;
+  }
+  return false;
+}
+
 static bool isBitfieldPositioningOpFromShl(SelectionDAG *CurDAG, SDValue Op,
                                            bool BiggerPattern,
                                            const uint64_t NonZeroBits,
                                            SDValue &Src, int &DstLSB,
                                            int &Width) {
   assert(isShiftedMask_64(NonZeroBits) && "Caller guaranteed");
 
   EVT VT = Op.getValueType();
   assert((VT == MVT::i32 || VT == MVT::i64) &&
          "Caller guarantees that type is i32 or i64");
 
   uint64_t ShlImm;
   if (!isOpcWithIntImmediate(Op.getNode(), ISD::SHL, ShlImm))
     return false;
 
   if (!BiggerPattern && !Op.hasOneUse())
     return false;
 
+  if (isSeveralBitsPositioningOpFromShl(ShlImm, Op, Src, DstLSB, Width))
+    return true;
+
   DstLSB = countTrailingZeros(NonZeroBits);
   Width = countTrailingOnes(NonZeroBits >> DstLSB);
 
   if (DstLSB != ShlImm && !BiggerPattern)
     return false;
 
   Src = getLeftShift(CurDAG, Op.getOperand(0), ShlImm - DstLSB);
   return true;

llvm/test/CodeGen/AArch64/arm64-indexed-vector-ldst.ll

; CHECK-NEXT:    ret
   ret i8 %e
 }
 
 define i16 @load_single_extract_variable_index_i16(<8 x i16>* %A, i32 %idx) {
 ; CHECK-LABEL: load_single_extract_variable_index_i16:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    sub sp, sp, #16
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-NEXT:    ; kill: def $w1 killed $w1 def $x1
-; CHECK-NEXT:    and x8, x1, #0x7
+; CHECK-NEXT:    mov x8, sp
 ; CHECK-NEXT:    ldr q0, [x0]
-; CHECK-NEXT:    mov x9, sp
-; CHECK-NEXT:    bfi x9, x8, #1, #3
+; CHECK-NEXT:    ; kill: def $w1 killed $w1 def $x1
+; CHECK-NEXT:    bfi x8, x1, #1, #3
 ; CHECK-NEXT:    str q0, [sp]
-; CHECK-NEXT:    ldrh w0, [x9]
+; CHECK-NEXT:    ldrh w0, [x8]
 ; CHECK-NEXT:    add sp, sp, #16
 ; CHECK-NEXT:    ret
   %lv = load <8 x i16>, <8 x i16>* %A
   %e = extractelement <8 x i16> %lv, i32 %idx
   ret i16 %e
 }
 
 define i32 @load_single_extract_variable_index_i32(<4 x i32>* %A, i32 %idx) {

llvm/test/CodeGen/AArch64/arm64-neon-copy.ll

; CHECK-NEXT:    ret
   ret <8 x i8> %vecinit14
 }
 
 define <4 x i16> @test_extracts_inserts_varidx_extract(<8 x i16> %x, i32 %idx) {
 ; CHECK-LABEL: test_extracts_inserts_varidx_extract:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    sub sp, sp, #16
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    mov x8, sp
 ; CHECK-NEXT:    // kill: def $w0 killed $w0 def $x0
-; CHECK-NEXT:    and x8, x0, #0x7
-; CHECK-NEXT:    mov x9, sp
 ; CHECK-NEXT:    str q0, [sp]
-; CHECK-NEXT:    bfi x9, x8, #1, #3
-; CHECK-NEXT:    ldr h1, [x9]
+; CHECK-NEXT:    bfi x8, x0, #1, #3
+; CHECK-NEXT:    ldr h1, [x8]
 ; CHECK-NEXT:    mov v1.h[1], v0.h[1]
 ; CHECK-NEXT:    mov v1.h[2], v0.h[2]
 ; CHECK-NEXT:    mov v1.h[3], v0.h[3]
 ; CHECK-NEXT:    fmov d0, d1
 ; CHECK-NEXT:    add sp, sp, #16
 ; CHECK-NEXT:    ret
   %tmp = extractelement <8 x i16> %x, i32 %idx
   %tmp2 = insertelement <4 x i16> undef, i16 %tmp, i32 0
   %tmp3 = extractelement <8 x i16> %x, i32 1
   %tmp4 = insertelement <4 x i16> %tmp2, i16 %tmp3, i32 1
   %tmp5 = extractelement <8 x i16> %x, i32 2
   %tmp6 = insertelement <4 x i16> %tmp4, i16 %tmp5, i32 2
   %tmp7 = extractelement <8 x i16> %x, i32 3
   %tmp8 = insertelement <4 x i16> %tmp6, i16 %tmp7, i32 3
   ret <4 x i16> %tmp8
 }
 
 define <4 x i16> @test_extracts_inserts_varidx_insert(<8 x i16> %x, i32 %idx) {
 ; CHECK-LABEL: test_extracts_inserts_varidx_insert:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    sub sp, sp, #16
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    add x8, sp, #8
 ; CHECK-NEXT:    // kill: def $w0 killed $w0 def $x0
-; CHECK-NEXT:    and x8, x0, #0x3
-; CHECK-NEXT:    add x9, sp, #8
-; CHECK-NEXT:    bfi x9, x8, #1, #2
-; CHECK-NEXT:    str h0, [x9]
+; CHECK-NEXT:    bfi x8, x0, #1, #2
+; CHECK-NEXT:    str h0, [x8]
 ; CHECK-NEXT:    ldr d1, [sp, #8]
 ; CHECK-NEXT:    mov v1.h[1], v0.h[1]
 ; CHECK-NEXT:    mov v1.h[2], v0.h[2]
 ; CHECK-NEXT:    mov v1.h[3], v0.h[3]
 ; CHECK-NEXT:    fmov d0, d1
 ; CHECK-NEXT:    add sp, sp, #16
 ; CHECK-NEXT:    ret
   %tmp = extractelement <8 x i16> %x, i32 0

llvm/test/CodeGen/AArch64/arm64-nvcast.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=arm64-apple-ios -aarch64-enable-collect-loh=false | FileCheck %s
 
 define void @test(float * %p1, i32 %v1) {
 ; CHECK-LABEL: test:
 ; CHECK:       ; %bb.0: ; %entry
 ; CHECK-NEXT:    sub sp, sp, #16
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-NEXT:    ; kill: def $w1 killed $w1 def $x1
-; CHECK-NEXT:    and x8, x1, #0x3
-; CHECK-NEXT:    mov x9, sp
 ; CHECK-NEXT:    fmov.2d v0, #2.00000000
-; CHECK-NEXT:    bfi x9, x8, #2, #2
+; CHECK-NEXT:    mov x8, sp
+; CHECK-NEXT:    ; kill: def $w1 killed $w1 def $x1
+; CHECK-NEXT:    bfi x8, x1, #2, #2
 ; CHECK-NEXT:    str q0, [sp]
-; CHECK-NEXT:    ldr s0, [x9]
+; CHECK-NEXT:    ldr s0, [x8]
 ; CHECK-NEXT:    str s0, [x0]
 ; CHECK-NEXT:    add sp, sp, #16
 ; CHECK-NEXT:    ret
 entry:
   %v2 = extractelement <3 x float> <float 0.000000e+00, float 2.000000e+00, float 0.000000e+00>, i32 %v1
   store float %v2, float* %p1, align 4
   ret void
 }
 
 define void @test2(float * %p1, i32 %v1) {
 ; CHECK-LABEL: test2:
 ; CHECK:       ; %bb.0: ; %entry
 ; CHECK-NEXT:    sub sp, sp, #16
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-NEXT:    ; kill: def $w1 killed $w1 def $x1
-; CHECK-NEXT:    and x8, x1, #0x3
-; CHECK-NEXT:    mov x9, sp
 ; CHECK-NEXT:    movi.16b v0, #63
-; CHECK-NEXT:    bfi x9, x8, #2, #2
+; CHECK-NEXT:    mov x8, sp
+; CHECK-NEXT:    ; kill: def $w1 killed $w1 def $x1
+; CHECK-NEXT:    bfi x8, x1, #2, #2
 ; CHECK-NEXT:    str q0, [sp]
-; CHECK-NEXT:    ldr s0, [x9]
+; CHECK-NEXT:    ldr s0, [x8]
 ; CHECK-NEXT:    str s0, [x0]
 ; CHECK-NEXT:    add sp, sp, #16
 ; CHECK-NEXT:    ret
 entry:
   %v2 = extractelement <3 x float> <float 0.7470588088035583, float 0.7470588088035583, float 0.7470588088035583>, i32 %v1
   store float %v2, float* %p1, align 4
   ret void
 }

llvm/test/CodeGen/AArch64/pr58350.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=aarch64-none-linux-gnu | FileCheck %s
 
 ; This used to hit an assertion caused by dagcombine merge store.
 ; When the store memVT is v1f32 and the other store to be merged
 ; is f32, we need to build vector for the f32 store.
 
 define void @f(<1 x float> %a, i64 %b) {
 ; CHECK-LABEL: f:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    sub sp, sp, #16
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    adrp x8, .LCPI0_0
-; CHECK-NEXT:    and x9, x0, #0x1
-; CHECK-NEXT:    mov x10, sp
+; CHECK-NEXT:    mov x9, sp
+; CHECK-NEXT:    bfi x9, x0, #2, #1
 ; CHECK-NEXT:    // kill: def $d0 killed $d0 def $q0
 ; CHECK-NEXT:    ldr d1, [x8, :lo12:.LCPI0_0]
-; CHECK-NEXT:    bfi x10, x9, #2, #1
 ; CHECK-NEXT:    str d1, [sp]
-; CHECK-NEXT:    ldr s1, [x10]
+; CHECK-NEXT:    ldr s1, [x9]
 ; CHECK-NEXT:    mov v1.s[1], v0.s[0]
 ; CHECK-NEXT:    str d1, [sp, #8]
 ; CHECK-NEXT:    add sp, sp, #16
 ; CHECK-NEXT:    ret
   %P = alloca i64
   %E = extractelement <2 x float> <float 0.5, float 1.0>, i64 %b
   %G = getelementptr <1 x float>, ptr %P, i64 1
   store float %E, ptr %P
   store <1 x float> %a, ptr %G
   ret void
 }

llvm/test/CodeGen/AArch64/vec_uaddo.ll

 define <4 x i32> @uaddo_v4i1(<4 x i1> %a0, <4 x i1> %a1, <4 x i1>* %p2) nounwind {
 ; CHECK-LABEL: uaddo_v4i1:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v2.4h, #1
 ; CHECK-NEXT:    and v1.8b, v1.8b, v2.8b
 ; CHECK-NEXT:    and v0.8b, v0.8b, v2.8b
 ; CHECK-NEXT:    add v0.4h, v0.4h, v1.4h
 ; CHECK-NEXT:    umov w8, v0.h[1]
-; CHECK-NEXT:    umov w9, v0.h[2]
-; CHECK-NEXT:    umov w10, v0.h[0]
+; CHECK-NEXT:    umov w9, v0.h[0]
+; CHECK-NEXT:    umov w10, v0.h[2]
 ; CHECK-NEXT:    umov w11, v0.h[3]
 ; CHECK-NEXT:    and v1.8b, v0.8b, v2.8b
 ; CHECK-NEXT:    cmeq v0.4h, v1.4h, v0.4h
-; CHECK-NEXT:    and w8, w8, #0x1
-; CHECK-NEXT:    and w9, w9, #0x1
+; CHECK-NEXT:    bfi w9, w8, #1, #1
+; CHECK-NEXT:    bfi w9, w10, #2, #1
 ; CHECK-NEXT:    mvn v0.8b, v0.8b
+; CHECK-NEXT:    bfi w9, w11, #3, #29
+; CHECK-NEXT:    and w8, w9, #0xf
 ; CHECK-NEXT:    sshll v0.4s, v0.4h, #0
-; CHECK-NEXT:    bfi w10, w8, #1, #1
-; CHECK-NEXT:    bfi w10, w9, #2, #1
-; CHECK-NEXT:    bfi w10, w11, #3, #29
-; CHECK-NEXT:    and w8, w10, #0xf
 ; CHECK-NEXT:    strb w8, [x0]
 ; CHECK-NEXT:    ret
   %t = call {<4 x i1>, <4 x i1>} @llvm.uadd.with.overflow.v4i1(<4 x i1> %a0, <4 x i1> %a1)
   %val = extractvalue {<4 x i1>, <4 x i1>} %t, 0
   %obit = extractvalue {<4 x i1>, <4 x i1>} %t, 1
   %res = sext <4 x i1> %obit to <4 x i32>
   store <4 x i1> %val, <4 x i1>* %p2
   ret <4 x i32> %res

llvm/test/CodeGen/AArch64/vec_umulo.ll

 
 define <4 x i32> @umulo_v4i1(<4 x i1> %a0, <4 x i1> %a1, <4 x i1>* %p2) nounwind {
 ; CHECK-LABEL: umulo_v4i1:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    fmov d2, d0
 ; CHECK-NEXT:    movi v0.2d, #0000000000000000
 ; CHECK-NEXT:    and v1.8b, v2.8b, v1.8b
 ; CHECK-NEXT:    umov w8, v1.h[1]
-; CHECK-NEXT:    umov w9, v1.h[2]
-; CHECK-NEXT:    umov w10, v1.h[0]
+; CHECK-NEXT:    umov w9, v1.h[0]
+; CHECK-NEXT:    umov w10, v1.h[2]
 ; CHECK-NEXT:    umov w11, v1.h[3]
-; CHECK-NEXT:    and w8, w8, #0x1
-; CHECK-NEXT:    and w9, w9, #0x1
-; CHECK-NEXT:    bfi w10, w8, #1, #1
-; CHECK-NEXT:    bfi w10, w9, #2, #1
-; CHECK-NEXT:    bfi w10, w11, #3, #29
-; CHECK-NEXT:    and w8, w10, #0xf
+; CHECK-NEXT:    bfi w9, w8, #1, #1
+; CHECK-NEXT:    bfi w9, w10, #2, #1
+; CHECK-NEXT:    bfi w9, w11, #3, #29
+; CHECK-NEXT:    and w8, w9, #0xf
 ; CHECK-NEXT:    strb w8, [x0]
 ; CHECK-NEXT:    ret
   %t = call {<4 x i1>, <4 x i1>} @llvm.umul.with.overflow.v4i1(<4 x i1> %a0, <4 x i1> %a1)
   %val = extractvalue {<4 x i1>, <4 x i1>} %t, 0
   %obit = extractvalue {<4 x i1>, <4 x i1>} %t, 1
   %res = sext <4 x i1> %obit to <4 x i32>
   store <4 x i1> %val, <4 x i1>* %p2
   ret <4 x i32> %res