[AArch64] Enhance bit-field-positioning op matcher to see through "any_extend" for pattern "and(any_extend(shl(val, N)), shifted-mask)"

[mingmingl]

Summary

Before this patch (and refactor patch D135843), isBitfieldPositioningOp won't handle and(any_extend(shl(val, N), shifted-mask) (bail out if AND op is not SHL)

After this patch, isBitfieldPositioningOp will see through any_extend to find shl to find possible bit-field-positioning nodes.

https://gcc.godbolt.org/z/3ncGKbGW6 is a four-liner LLVM IR that could be optimized to UBFIZ (see added test case test_and_extended_shift_with_imm in llvm/test/CodeGen/AArch64/bitfield-insert.ll). One existing test case also improves.

[dmgreen]

llvm/test/CodeGen/AArch64/bitfield-insert.ll

598

This is nice. I've seen this come up before. https://godbolt.org/z/hhhTKex4d

[mingmingl]

Changes:

1. Did some changes on local commit log to use D135843 (the refactor) as parent revision (previous A -> B -> C -> D, now A -> D), in a way that after A is upstreamed this one could be rebased.
2. Minor comments update, no functional change.

llvm/test/CodeGen/AArch64/bitfield-insert.ll

598

Yes, and the 2nd function in https://godbolt.org/z/hhhTKex4d are optimized to UBFIZ with this patch.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

2585

Can we create a new variable for the temporary SDValue?

2589

Add an assert that the type of ShlVal is i32?

2594

-> Widens

2615

I think this will always be a MVT::i64 or MVT::i32?

2617

Do we have any tests where this applies?

[mingmingl]

address comments.

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

2585

Renamed this to 'ShlOp0'

2589

Add the type assertion and comment for why ShlVal type is i32 (based on my understanding).

2615

Yes, I agree this should always be i64 or i32 after reading the comment and think about it.

p.s. if I understand correctly, this is because

1. this function gets called during selection (i.e., after legalization) and caller verifies i64 or i32 above; &&
2. AArch64 general purpose registers are for 32 or 64 bit integers, and SimpleVT legalization (from target-lowering-base) class is based on register info.

2617

This is a good point. I think the answer is no (or at least very rare if any) -> changed implementation to bail out if Width >= VT.getSizeInBits() and added comment.

Longer story is that, if AndImm (as a superset of shifted-mask 'NonZeroBits`) from and(op, AndImm) requires more than 32 bits after proper dag combination, the op shouldn't be lowered to any_extend (smaller-type op)

• the higher bits are undefined after any_extend, so previous combination should at least realize the higher 32 bits and results don't really matter.

[dmgreen]

LGTM

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

2615

Yep. Sounds good.

2617

There are quite a lot of words here for "we don't expect this to happen but it would be bad if it did" :)

[mingmingl]

Diff: reworded lengthy comment for why bailing out upon Width >= VT.getSizeInBits. No other change.

Thanks for the reviews and feedback!

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

static bool isBitfieldPositioningOpFromAnd(SelectionDAG *CurDAG, SDValue Op,
   // 1) and 2) don't agree so something must be wrong (e.g., in
   // 'SelectionDAG::computeKnownBits')
   assert((~AndImm & NonZeroBits) == 0 &&
          "Something must be wrong (e.g., in SelectionDAG::computeKnownBits)");
 
   SDValue AndOp0 = Op.getOperand(0);
 
   uint64_t ShlImm;
-  if (!isOpcWithIntImmediate(AndOp0.getNode(), ISD::SHL, ShlImm))
+  SDValue ShlOp0;
+  if (isOpcWithIntImmediate(AndOp0.getNode(), ISD::SHL, ShlImm)) {
+    // For pattern "and(shl(val, N), shifted-mask)", 'ShlOp0' is set to 'val'.
+    ShlOp0 = AndOp0.getOperand(0);
+  } else if (VT == MVT::i64 && AndOp0.getOpcode() == ISD::ANY_EXTEND &&
+             isOpcWithIntImmediate(AndOp0.getOperand(0).getNode(), ISD::SHL,
+                                   ShlImm)) {
+    // For pattern "and(any_extend(shl(val, N)), shifted-mask)"
+
+    // ShlVal == shl(val, N), which is a left shift on a smaller type.
+    SDValue ShlVal = AndOp0.getOperand(0);
+
+    // Since this is after type legalization and ShlVal is extended to MVT::i64,
+    // expect VT to be MVT::i32.
+    assert((ShlVal.getValueType() == MVT::i32) && "Expect VT to be MVT::i32.");
+
+    // Widens 'val' to MVT::i64 as the source of bit field positioning.
+    ShlOp0 = Widen(CurDAG, ShlVal.getOperand(0));
+  } else
     return false;
 
-  // Bail out if the SHL has more than one use, since then we'll end up
-  // generating SHL+UBFIZ instead of just keeping SHL+AND.
+  // For !BiggerPattern, bail out if the AndOp0 has more than one use, since
+  // then we'll end up generating AndOp0+UBFIZ instead of just keeping
+  // AndOp0+AND.
   if (!BiggerPattern && !AndOp0.hasOneUse())
     return false;
 
   DstLSB = countTrailingZeros(NonZeroBits);
   Width = countTrailingOnes(NonZeroBits >> DstLSB);
 
+  // Bail out on large Width. This happens when no proper combining / constant
+  // folding was performed.
+  if (Width >= (int)VT.getSizeInBits()) {
+    // If VT is i64, Width > 64 is insensible since NonZeroBits is uint64_t, and
+    // Width == 64 indicates a missed dag-combine from "(and val, AllOnes)" to
+    // "val".
+    // If VT is i32, what Width >= 32 means:
+    // - For "(and (any_extend(shl val, N)), shifted-mask)", the`and` Op
+    //   demands at least 'Width' bits (after dag-combiner). This together with
+    //   `any_extend` Op (undefined higher bits) indicates missed combination
+    //   when lowering the 'and' IR instruction to an machine IR instruction.
+    LLVM_DEBUG(
+        dbgs()
+        << "Found large Width in bit-field-positioning -- this indicates no "
+           "proper combining / constant folding was performed\n");
+    return false;
+  }
+
   // BFI encompasses sufficiently many nodes that it's worth inserting an extra
   // LSL/LSR if the mask in NonZeroBits doesn't quite match up with the ISD::SHL
   // amount.  BiggerPattern is true when this pattern is being matched for BFI,
   // BiggerPattern is false when this pattern is being matched for UBFIZ, in
   // which case it is not profitable to insert an extra shift.
   if (ShlImm != uint64_t(DstLSB) && !BiggerPattern)
     return false;
 
-  Src = getLeftShift(CurDAG, AndOp0.getOperand(0), ShlImm - DstLSB);
+  Src = getLeftShift(CurDAG, ShlOp0, 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) {

llvm/test/CodeGen/AArch64/bitfield-insert.ll

; CHECK-NEXT:    ret
   store i64 %both, i64* %bf
 
   ret <2 x i32> %vec
 }
 
 define i64 @test_truncated_shift(i64 %x, i64 %y) {
 ; CHECK-LABEL: test_truncated_shift:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    lsl w8, w1, #25
-; CHECK-NEXT:    lsr x8, x8, #25
-; CHECK-NEXT:    bfi x0, x8, #25, #5
+; CHECK-NEXT:    // kill: def $w1 killed $w1 killed $x1 def $x1
+; CHECK-NEXT:    bfi x0, x1, #25, #5
 ; CHECK-NEXT:    ret
 entry:
   %and = and i64 %x, -1040187393
   %shl4 = shl i64 %y, 25
   %and5 = and i64 %shl4, 1040187392
   %or = or i64 %and5, %and
   ret i64 %or
 }
+
+define i64 @test_and_extended_shift_with_imm(i64 %0) {
+; CHECK-LABEL: test_and_extended_shift_with_imm:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $w0 killed $w0 killed $x0 def $x0
+; CHECK-NEXT:    ubfiz x0, x0, #7, #8
+; CHECK-NEXT:    ret
+  %2 = shl i64 %0, 7
+  %3 = and i64 %2, 32640  ; #0x7f80
+  ret i64 %3
+}