[AArch64] Transform shift+and to shift+shift to select more shifted register

[bcl5980]

Summary

and (shl/srl/sra, x, c), mask --> shl (srl/sra, x, c1), c2

[bcl5980]

I will check in the precommit tests if the reviewers look good for the tests.

[dmgreen]

llvm/test/CodeGen/AArch64/shiftregister-from-and.ll

114

This instruction looks wrong, with the shift being more than 32 for w regs.

[mingmingl]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

632

nit: use uint64_t (the same type as the return value of getZExtValue() to avoid warnings.

647–671

nit pick:

To ensure that UBFM/SBFM is used as a right shift (essentially UBFX) here, add assert(BitWidth-1>=NewShiftC && "error message")

654–655

Relatedly, with 1) N as and (srl x, c), mask and 2) LowZBits == 0, N should be selected as UBFX by isBitfieldExtractOp before tablegen pattern (where AArch64DAGToDAGISel::SelectShiftedRegister gets called) applies. Similarly, with and(shl x, c), mask and LowZBits == 0, N is a UBFIZ.

So I think this shouldn't happen be seen in practice; but an early exit (not assertion) looks okay.

657–658

For correctness, similar check is needed when LHSOpcode is ISD::SRL; meanwhile for srl, the condition BitWidth == LowZBits + MaskLen + ShiftAmtC should be sufficient

For example, https://gcc.godbolt.org/z/YvGG3Pov3 shows an IR at trunk; and with the current patch the optimized code changes the result (not correct).

	lsr	x8, x0, #3
	add	x0, x1, x8, lsl #1
	ret

llvm/test/CodeGen/AArch64/shiftregister-from-and.ll

137–139

nit pick:

i16 will be legalized to i32 [1], and the patch optimizes shiftedreg_from_and_negative_type (trunk https://gcc.godbolt.org/z/cnac5989e). This doesn't seem to be a negative test . Probably use a vector like https://gcc.godbolt.org/z/T4szTsEav

[1] llc -O3 -mtriple=aarch64 -debug -print-after-all test.ll gives debug logs, and https://gist.github.com/minglotus/763173f099bec76c720ec9ce7c181471 shows the legalization step.

[bcl5980]

address comments.

[bcl5980]

test comments update.

[bcl5980]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

657–658

BitWidth <= LowZBits + MaskLen + ShiftAmtC should be better.

llvm/test/CodeGen/AArch64/shiftregister-from-and.ll

114

fixed the code and rename the test to @shiftedreg_from_and_negative_andc5

[bcl5980]

remove redundant condition

[mingmingl]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

689–694

Selecting and (shl/srl/sra, x, c), mask into shl (srl/sra, x, c1), c2 (and folding shl) is simplifying one instruction away but could change the pipeline (latency and throughput as well) of the instruciton.

For example, arithmetic operations with shifted operand could use M pipeline [1] on neoverse n1, and M pipeline is used for all ALU operations with imm-shifted operand for cortex a57.

I wonder if improvements are seen in some benchmarks from this patch?

[1] for neoverse n1, M pipeline is used for {ADDS, SUBS} with "Arithmetic, LSR/ASR/ROR shift or LSL shift > 4".

[bcl5980]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

689–694

I agree that it may not get any timing improvement except lsl < 4. I also don't believe this patch can get improvement in any real benchmarks. Actually the current shifted register select function also does not consider for the thing you mentioned. If you really worry about that I can limit the lsl constant < 4.But personally I don't want to do that.

[bcl5980]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

689–694

And for high-end CPU, it's really hard to model by just execute ports. 2 instructions will use 2 Rob entry, 2 physical registers, more front end decode, more code size. At least we should not consider these detail things when these pattern with the same latency and only some targets may have different throughput. If we really want to consider these detail micro architecture trade off, the better way is always do here and maybe split it in machine combiner that we already have detail schmodel.

[dmgreen]

Performance can depend on many factors, and often less instructions is a win in itself. Can you separate out and rebase on the tests to just show the differences?

[bcl5980]

rebase code

[mingmingl]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

689–694

Thanks for the discussion! It makes sense to me to consider rewriting ALU with shifted operand into two simpler instructions in machine-combiner pass, and that less instructions is a win in itself (regarding David's comment)

To provide some context why the original ask, there are motivating test cases (micro benchmarks) that want something opposite (i.e., split ALU with shifted operand into two simpler instructions); a manual hack to use dummy inline assembly (asm volatile("" : "+r" (variable));) shows speedups for neoverse n1 (sole M pipeline, and could become bottleneck) -> and I was thinking about machine-combiner as one place (to implement split with some critical path analysis?) as well.

[bcl5980]

So how about this patch？

[mingmingl]

In D138904#3970087, @bcl5980 wrote:

So how about this patch？

I'm not too concerned about having fewer instructions and using M pipeline, since I also concur having fewer instructions makes general sense (as you and David pointed out). So this transformation looks good to me . Also note on some aarch64 processors (e.g., neoverse n1), logical instructions with imm-shifted operand is a net win compared with two instructions; and some aarch64 processors have more than one M pipeline (neoverse n2), so fewer chances of M pipeline being a bottleneck.

Tests on microbenchmarks (llvm-test-suite, etc) might give us more confidence (the numbers depend on the specific processor though), but I don't strongly prefer a performance test.

[bcl5980]

In D138904#3971306, @mingmingl wrote:

In D138904#3970087, @bcl5980 wrote:

So how about this patch？

I'm not too concerned about having fewer instructions and using M pipeline, since I also concur having fewer instructions makes general sense (as you and David pointed out). So this transformation looks good to me . Also note on some aarch64 processors (e.g., neoverse n1), logical instructions with imm-shifted operand is a net win compared with two instructions; and some aarch64 processors have more than one M pipeline (neoverse n2), so fewer chances of M pipeline being a bottleneck.

Tests on microbenchmarks (llvm-test-suite, etc) might give us more confidence (the numbers depend on the specific processor though), but I don't strongly prefer a performance test.

I try llvm-test-suite but unfortunately it looks no test trigger the code.

[dmgreen]

From what I can tell I think the code is OK. LGTM

This doesn't come up often, but I do see it triggering. I'm not sure what the motivating case is but it seemed OK for performance in the tests I ran.

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

633–634

If RHS isnt used anywhere else:
ConstantSDNode *RHSC = dyn_cast<ConstantSDNode>(N.getOperand(1));

658

Some more comments explaining these conditions would be good.

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

public:
   bool tryWriteRegister(SDNode *N);
 
 // Include the pieces autogenerated from the target description.
 #include "AArch64GenDAGISel.inc"
 
 private:
   bool SelectShiftedRegister(SDValue N, bool AllowROR, SDValue &Reg,
                              SDValue &Shift);
+  bool SelectShiftedRegisterFromAnd(SDValue N, SDValue &Reg, SDValue &Shift);
   bool SelectAddrModeIndexed7S(SDValue N, unsigned Size, SDValue &Base,
                                SDValue &OffImm) {
     return SelectAddrModeIndexedBitWidth(N, true, 7, Size, Base, OffImm);
   }
   bool SelectAddrModeIndexedBitWidth(SDValue N, bool IsSignedImm, unsigned BW,
                                      unsigned Size, SDValue &Base,
                                      SDValue &OffImm);
   bool SelectAddrModeIndexed(SDValue N, unsigned Size, SDValue &Base,
if (Subtarget->hasLSLFast() && V.getOpcode() == ISD::ADD) {
     if (RHS.getOpcode() == ISD::SHL && isWorthFoldingSHL(RHS))
       return true;
   }
 
   // It hurts otherwise, since the value will be reused.
   return false;
 }
 
+/// and (shl/srl/sra, x, c), mask --> shl (srl/sra, x, c1), c2
+/// to select more shifted register
+bool AArch64DAGToDAGISel::SelectShiftedRegisterFromAnd(SDValue N, SDValue &Reg,
+                                                       SDValue &Shift) {
+  EVT VT = N.getValueType();
+  if (VT != MVT::i32 && VT != MVT::i64)
+    return false;
+
+  if (N->getOpcode() != ISD::AND || !N->hasOneUse())
+    return false;
+  SDValue LHS = N.getOperand(0);
+  if (!LHS->hasOneUse())
+    return false;
+
+  unsigned LHSOpcode = LHS->getOpcode();
+  if (LHSOpcode != ISD::SHL && LHSOpcode != ISD::SRL && LHSOpcode != ISD::SRA)
+    return false;
+
+  ConstantSDNode *ShiftAmtNode = dyn_cast<ConstantSDNode>(LHS.getOperand(1));
+  if (!ShiftAmtNode)
+    return false;
+
+  uint64_t ShiftAmtC = ShiftAmtNode->getZExtValue();
+  SDValue RHS = N.getOperand(1);
+  ConstantSDNode *RHSC = dyn_cast<ConstantSDNode>(RHS);
+  if (!RHSC)
+    return false;
+
+  APInt AndMask = RHSC->getAPIntValue();
+  unsigned LowZBits, MaskLen;
+  if (!AndMask.isShiftedMask(LowZBits, MaskLen))
+    return false;
+
+  unsigned BitWidth = N.getValueSizeInBits();
+  SDLoc DL(LHS);
+  uint64_t NewShiftC;
+  unsigned NewShiftOp;
+  if (LHSOpcode == ISD::SHL) {
+    if (LowZBits <= ShiftAmtC || (BitWidth != LowZBits + MaskLen))
+      return false;
+
+    NewShiftC = LowZBits - ShiftAmtC;
+    NewShiftOp = VT == MVT::i64 ? AArch64::UBFMXri : AArch64::UBFMWri;
+  } else {
+    if (LowZBits == 0)
+      return false;
+
+    NewShiftC = LowZBits + ShiftAmtC;
+    if (NewShiftC >= BitWidth)
+      return false;
+
+    if (LHSOpcode == ISD::SRA && (BitWidth != (LowZBits + MaskLen)))
+      return false;
+
+    if (LHSOpcode == ISD::SRL && (BitWidth > (NewShiftC + MaskLen)))
+      return false;
+
+    if (LHSOpcode == ISD::SRL)
+      NewShiftOp = VT == MVT::i64 ? AArch64::UBFMXri : AArch64::UBFMWri;
+    else
+      NewShiftOp = VT == MVT::i64 ? AArch64::SBFMXri : AArch64::SBFMWri;
+  }
+
+  assert(NewShiftC < BitWidth && "Invalid shift amount");
+  SDValue NewShiftAmt = CurDAG->getTargetConstant(NewShiftC, DL, VT);
+  SDValue BitWidthMinus1 = CurDAG->getTargetConstant(BitWidth - 1, DL, VT);
+  Reg = SDValue(CurDAG->getMachineNode(NewShiftOp, DL, VT, LHS->getOperand(0),
+                                       NewShiftAmt, BitWidthMinus1),
+                0);
+  unsigned ShVal = AArch64_AM::getShifterImm(AArch64_AM::LSL, LowZBits);
+  Shift = CurDAG->getTargetConstant(ShVal, DL, MVT::i32);
+  return true;
+}
+
 /// SelectShiftedRegister - Select a "shifted register" operand.  If the value
 /// is not shifted, set the Shift operand to default of "LSL 0".  The logical
 /// instructions allow the shifted register to be rotated, but the arithmetic
 /// instructions do not.  The AllowROR parameter specifies whether ROR is
 /// supported.
 bool AArch64DAGToDAGISel::SelectShiftedRegister(SDValue N, bool AllowROR,
                                                 SDValue &Reg, SDValue &Shift) {
+  if (SelectShiftedRegisterFromAnd(N, Reg, Shift))
+    return true;
+
   AArch64_AM::ShiftExtendType ShType = getShiftTypeForNode(N);
   if (ShType == AArch64_AM::InvalidShiftExtend)
     return false;
   if (!AllowROR && ShType == AArch64_AM::ROR)
     return false;
 
   if (ConstantSDNode *RHS = dyn_cast<ConstantSDNode>(N.getOperand(1))) {
     unsigned BitSize = N.getValueSizeInBits();
     unsigned Val = RHS->getZExtValue() & (BitSize - 1);

llvm/test/CodeGen/AArch64/shiftregister-from-and.ll

; CHECK-NEXT:    ret
   ret i64 %r
 }
 
 ; logic shift reg pattern: eon
 
 define i64 @eon_shiftedreg_from_and(i64 %a, i64 %b) {
 ; CHECK-LABEL: eon_shiftedreg_from_and:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    lsl x8, x0, #36
-; CHECK-NEXT:    and x8, x8, #0xffe0000000000000
-; CHECK-NEXT:    eon x0, x8, x1
+; CHECK-NEXT:    lsr x8, x0, #17
+; CHECK-NEXT:    eon x0, x1, x8, lsl #53
 ; CHECK-NEXT:    ret
   %shl = shl i64 %a, 36
   %and = and i64 %shl, -9007199254740992
   %xor = xor i64 %and, -1
   %r = xor i64 %b, %xor
   ret i64 %r
 }
 
 ; logic shift reg pattern: eor
 
 define i64 @eor_shiftedreg_from_and(i64 %a, i64 %b) {
 ; CHECK-LABEL: eor_shiftedreg_from_and:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    lsr x8, x0, #23
-; CHECK-NEXT:    and x8, x8, #0x1ffff000000
-; CHECK-NEXT:    eor x0, x8, x1
+; CHECK-NEXT:    lsr x8, x0, #47
+; CHECK-NEXT:    eor x0, x1, x8, lsl #24
 ; CHECK-NEXT:    ret
   %lshr = lshr i64 %a, 23
   %and = and i64 %lshr, 2199006478336
   %or = xor i64 %and, %b
   ret i64 %or
 }
 
 ; logic shift reg pattern: mvn

 }
 
 ; logic shift reg pattern: orr
 ; srl constant bitwidth == (lowbits + masklen + shiftamt)
 
 define i64 @orr_shiftedreg_from_and(i64 %a, i64 %b) {
 ; CHECK-LABEL: orr_shiftedreg_from_and:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    lsr x8, x0, #23
-; CHECK-NEXT:    and x8, x8, #0x1ffff000000
-; CHECK-NEXT:    orr x0, x8, x1
+; CHECK-NEXT:    lsr x8, x0, #47
+; CHECK-NEXT:    orr x0, x1, x8, lsl #24
 ; CHECK-NEXT:    ret
   %lshr = lshr i64 %a, 23
   %and = and i64 %lshr, 2199006478336 ; 0x1ffff000000
   %or = or i64 %and, %b
   ret i64 %or
 }
 
 ; logic shift reg pattern: orr
 ; srl constant bitwidth < (lowbits + masklen + shiftamt)
 
 define i64 @orr_shiftedreg_from_and_mask2(i64 %a, i64 %b) {
 ; CHECK-LABEL: orr_shiftedreg_from_and_mask2:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    lsr x8, x0, #23
-; CHECK-NEXT:    and x8, x8, #0x1ffff000000
-; CHECK-NEXT:    orr x0, x8, x1
+; CHECK-NEXT:    lsr x8, x0, #47
+; CHECK-NEXT:    orr x0, x1, x8, lsl #24
 ; CHECK-NEXT:    ret
   %lshr = lshr i64 %a, 23
   %and = and i64 %lshr, 4398029733888 ; 0x3ffff000000
   %or = or i64 %and, %b
   ret i64 %or
 }
 
 
 ; arithmetic shift reg pattern: add
 
 define i32 @add_shiftedreg_from_and(i32 %a, i32 %b) {
 ; CHECK-LABEL: add_shiftedreg_from_and:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    asr w8, w0, #3
-; CHECK-NEXT:    and w8, w8, #0xff000000
-; CHECK-NEXT:    add w0, w8, w1
+; CHECK-NEXT:    asr w8, w0, #27
+; CHECK-NEXT:    add w0, w1, w8, lsl #24
 ; CHECK-NEXT:    ret
   %ashr = ashr i32 %a, 3
   %and = and i32 %ashr, -16777216
   %add = add i32 %and, %b
   ret i32 %add
 }
 
 ; arithmetic shift reg pattern: sub
 
 define i64 @sub_shiftedreg_from_and_shl(i64 %a, i64 %b) {
 ; CHECK-LABEL: sub_shiftedreg_from_and_shl:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    lsl x8, x0, #36
-; CHECK-NEXT:    and x8, x8, #0xffe0000000000000
-; CHECK-NEXT:    sub x0, x1, x8
+; CHECK-NEXT:    lsr x8, x0, #17
+; CHECK-NEXT:    sub x0, x1, x8, lsl #53
 ; CHECK-NEXT:    ret
   %shl = shl i64 %a, 36
   %and = and i64 %shl, -9007199254740992
   %sub = sub i64 %b, %and
   ret i64 %sub
 }
 
 ; negative test: type is not i32 or i64