[AArch64] Optimize lowering of 128-bit shifts by a variable

[efriedma]

Summary

This takes advantage of the work that's been done on funnel shifts to make the implementation more readable.

I think we could shave off one more instruction in the funnel shift lowering: if we do the shifts in the opposite order, we can fold the shift by one into the orr. Mostly orthogonal to this patch, though.

The implementation here isn't really target-specific; maybe we should be doing this in target-independent code?

[RKSimon]

I agree there isn't any reason not to try and move this to generic legalization - although we are still finding a few edge legalization cases where funnel/rotates fail on some targets so you might encounter that.

[RKSimon]

reverse ping?

[RKSimon]

The implementation here isn't really target-specific; maybe we should be doing this in target-independent code?

@efriedma I'm hoping D101987's target-independent implementation will make this patch redundant.

[efriedma]

Looks like we end up with essentially the same code with this patch and your proposed patch; I'm happy to just use the target-independent version.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

SDValue AArch64TargetLowering::LowerShiftRightParts(SDValue Op,
   SDLoc dl(Op);
   SDValue ShOpLo = Op.getOperand(0);
   SDValue ShOpHi = Op.getOperand(1);
   SDValue ShAmt = Op.getOperand(2);
   unsigned Opc = (Op.getOpcode() == ISD::SRA_PARTS) ? ISD::SRA : ISD::SRL;
 
   assert(Op.getOpcode() == ISD::SRA_PARTS || Op.getOpcode() == ISD::SRL_PARTS);
 
-  SDValue RevShAmt = DAG.getNode(ISD::SUB, dl, MVT::i64,
-                                 DAG.getConstant(VTBits, dl, MVT::i64), ShAmt);
-  SDValue HiBitsForLo = DAG.getNode(ISD::SHL, dl, VT, ShOpHi, RevShAmt);
-
-  // Unfortunately, if ShAmt == 0, we just calculated "(SHL ShOpHi, 64)" which
-  // is "undef". We wanted 0, so CSEL it directly.
-  SDValue Cmp = emitComparison(ShAmt, DAG.getConstant(0, dl, MVT::i64),
-                               ISD::SETEQ, dl, DAG);
-  SDValue CCVal = DAG.getConstant(AArch64CC::EQ, dl, MVT::i32);
-  HiBitsForLo =
-      DAG.getNode(AArch64ISD::CSEL, dl, VT, DAG.getConstant(0, dl, MVT::i64),
-                  HiBitsForLo, CCVal, Cmp);
-
-  SDValue ExtraShAmt = DAG.getNode(ISD::SUB, dl, MVT::i64, ShAmt,
-                                   DAG.getConstant(VTBits, dl, MVT::i64));
-
-  SDValue LoBitsForLo = DAG.getNode(ISD::SRL, dl, VT, ShOpLo, ShAmt);
-  SDValue LoForNormalShift =
-      DAG.getNode(ISD::OR, dl, VT, LoBitsForLo, HiBitsForLo);
-
-  Cmp = emitComparison(ExtraShAmt, DAG.getConstant(0, dl, MVT::i64), ISD::SETGE,
-                       dl, DAG);
-  CCVal = DAG.getConstant(AArch64CC::GE, dl, MVT::i32);
-  SDValue LoForBigShift = DAG.getNode(Opc, dl, VT, ShOpHi, ExtraShAmt);
-  SDValue Lo = DAG.getNode(AArch64ISD::CSEL, dl, VT, LoForBigShift,
-                           LoForNormalShift, CCVal, Cmp);
-
-  // AArch64 shifts larger than the register width are wrapped rather than
-  // clamped, so we can't just emit "hi >> x".
-  SDValue HiForNormalShift = DAG.getNode(Opc, dl, VT, ShOpHi, ShAmt);
-  SDValue HiForBigShift =
-      Opc == ISD::SRA
-          ? DAG.getNode(Opc, dl, VT, ShOpHi,
-                        DAG.getConstant(VTBits - 1, dl, MVT::i64))
-          : DAG.getConstant(0, dl, VT);
-  SDValue Hi = DAG.getNode(AArch64ISD::CSEL, dl, VT, HiForBigShift,
-                           HiForNormalShift, CCVal, Cmp);
-
-  SDValue Ops[2] = { Lo, Hi };
-  return DAG.getMergeValues(Ops, dl);
+  // Compute the shift with the shift amount masked.
+  SDValue MaskedShAmt = DAG.getNode(ISD::AND, dl, MVT::i64, ShAmt,
+                                    DAG.getConstant(VTBits - 1, dl, MVT::i64));
+  SDValue Hi = DAG.getNode(Opc, dl, MVT::i64, ShOpHi, MaskedShAmt);
+  SDValue Lo = DAG.getNode(ISD::FSHR, dl, MVT::i64, ShOpHi, ShOpLo, ShAmt);
+
+  // Check if the shift is large.
+  SDValue Cmp = emitComparison(ShAmt, DAG.getConstant(VTBits, dl, MVT::i64),
+                               ISD::SETGE, dl, DAG);
+  // Adjust the results.
+  SDValue CCVal = DAG.getConstant(AArch64CC::GE, dl, MVT::i32);
+  Lo = DAG.getNode(AArch64ISD::CSEL, dl, VT, Hi, Lo, CCVal, Cmp);
+  if (Op.getOpcode() == ISD::SRA_PARTS) {
+    Hi = DAG.getNode(AArch64ISD::CSEL, dl, VT,
+                     DAG.getNode(ISD::SRA, dl, MVT::i64, ShOpHi,
+                                 DAG.getConstant(VTBits - 1, dl, MVT::i64)),
+                     Hi, CCVal, Cmp);
+  } else {
+    Hi = DAG.getNode(AArch64ISD::CSEL, dl, VT, DAG.getConstant(0, dl, MVT::i64),
+                     Hi, CCVal, Cmp);
+  }
+  return DAG.getMergeValues({Lo, Hi}, dl);
 }
 
 /// LowerShiftLeftParts - Lower SHL_PARTS, which returns two
 /// i64 values and take a 2 x i64 value to shift plus a shift amount.
 SDValue AArch64TargetLowering::LowerShiftLeftParts(SDValue Op,
                                                    SelectionDAG &DAG) const {
   assert(Op.getNumOperands() == 3 && "Not a double-shift!");
   EVT VT = Op.getValueType();
   unsigned VTBits = VT.getSizeInBits();
   SDLoc dl(Op);
   SDValue ShOpLo = Op.getOperand(0);
   SDValue ShOpHi = Op.getOperand(1);
   SDValue ShAmt = Op.getOperand(2);
 
   assert(Op.getOpcode() == ISD::SHL_PARTS);
-  SDValue RevShAmt = DAG.getNode(ISD::SUB, dl, MVT::i64,
-                                 DAG.getConstant(VTBits, dl, MVT::i64), ShAmt);
-  SDValue LoBitsForHi = DAG.getNode(ISD::SRL, dl, VT, ShOpLo, RevShAmt);
-
-  // Unfortunately, if ShAmt == 0, we just calculated "(SRL ShOpLo, 64)" which
-  // is "undef". We wanted 0, so CSEL it directly.
-  SDValue Cmp = emitComparison(ShAmt, DAG.getConstant(0, dl, MVT::i64),
-                               ISD::SETEQ, dl, DAG);
-  SDValue CCVal = DAG.getConstant(AArch64CC::EQ, dl, MVT::i32);
-  LoBitsForHi =
-      DAG.getNode(AArch64ISD::CSEL, dl, VT, DAG.getConstant(0, dl, MVT::i64),
-                  LoBitsForHi, CCVal, Cmp);
-
-  SDValue ExtraShAmt = DAG.getNode(ISD::SUB, dl, MVT::i64, ShAmt,
-                                   DAG.getConstant(VTBits, dl, MVT::i64));
-  SDValue HiBitsForHi = DAG.getNode(ISD::SHL, dl, VT, ShOpHi, ShAmt);
-  SDValue HiForNormalShift =
-      DAG.getNode(ISD::OR, dl, VT, LoBitsForHi, HiBitsForHi);
-
-  SDValue HiForBigShift = DAG.getNode(ISD::SHL, dl, VT, ShOpLo, ExtraShAmt);
-
-  Cmp = emitComparison(ExtraShAmt, DAG.getConstant(0, dl, MVT::i64), ISD::SETGE,
-                       dl, DAG);
-  CCVal = DAG.getConstant(AArch64CC::GE, dl, MVT::i32);
-  SDValue Hi = DAG.getNode(AArch64ISD::CSEL, dl, VT, HiForBigShift,
-                           HiForNormalShift, CCVal, Cmp);
-
-  // AArch64 shifts of larger than register sizes are wrapped rather than
-  // clamped, so we can't just emit "lo << a" if a is too big.
-  SDValue LoForBigShift = DAG.getConstant(0, dl, VT);
-  SDValue LoForNormalShift = DAG.getNode(ISD::SHL, dl, VT, ShOpLo, ShAmt);
-  SDValue Lo = DAG.getNode(AArch64ISD::CSEL, dl, VT, LoForBigShift,
-                           LoForNormalShift, CCVal, Cmp);
 
-  SDValue Ops[2] = { Lo, Hi };
-  return DAG.getMergeValues(Ops, dl);
+  // Compute the shift with the shift amount masked.
+  SDValue MaskedShAmt = DAG.getNode(ISD::AND, dl, MVT::i64, ShAmt,
+                                    DAG.getConstant(VTBits - 1, dl, MVT::i64));
+  SDValue Lo = DAG.getNode(ISD::SHL, dl, MVT::i64, ShOpLo, MaskedShAmt);
+  SDValue Hi = DAG.getNode(ISD::FSHL, dl, MVT::i64, ShOpHi, ShOpLo, ShAmt);
+
+  // Check if the shift is large.
+  SDValue Cmp = emitComparison(ShAmt, DAG.getConstant(VTBits, dl, MVT::i64),
+                               ISD::SETGE, dl, DAG);
+  // Adjust the results.
+  SDValue CCVal = DAG.getConstant(AArch64CC::GE, dl, MVT::i32);
+  Hi = DAG.getNode(AArch64ISD::CSEL, dl, VT, Lo, Hi, CCVal, Cmp);
+  Lo = DAG.getNode(AArch64ISD::CSEL, dl, VT, DAG.getConstant(0, dl, MVT::i64),
+                   Lo, CCVal, Cmp);
+  return DAG.getMergeValues({Lo, Hi}, dl);
 }
 
 bool AArch64TargetLowering::isOffsetFoldingLegal(
     const GlobalAddressSDNode *GA) const {
   // Offsets are folded in the DAG combine rather than here so that we can
   // intelligently choose an offset based on the uses.
   return false;
 }

llvm/test/CodeGen/AArch64/arm64-long-shift.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=arm64-eabi -mcpu=cyclone | FileCheck %s
 
 define i128 @shl(i128 %r, i128 %s) nounwind readnone {
 ; CHECK-LABEL: shl:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    neg x8, x2
-; CHECK-NEXT:    lsr x8, x0, x8
-; CHECK-NEXT:    cmp x2, #0 // =0
-; CHECK-NEXT:    csel x8, xzr, x8, eq
-; CHECK-NEXT:    lsl x9, x1, x2
+; CHECK-NEXT:    lsl x8, x1, x2
+; CHECK-NEXT:    mvn w9, w2
+; CHECK-NEXT:    lsr x10, x0, #1
+; CHECK-NEXT:    lsr x9, x10, x9
 ; CHECK-NEXT:    orr x8, x8, x9
 ; CHECK-NEXT:    lsl x9, x0, x2
-; CHECK-NEXT:    sub x10, x2, #64 // =64
-; CHECK-NEXT:    cmp x10, #0 // =0
+; CHECK-NEXT:    cmp x2, #64 // =64
 ; CHECK-NEXT:    csel x1, x9, x8, ge
 ; CHECK-NEXT:    csel x0, xzr, x9, ge
 ; CHECK-NEXT:    ret
   %shl = shl i128 %r, %s
   ret i128 %shl
 }
 
 define i128 @shl_mask(i128 %r, i128 %s) nounwind readnone {
; CHECK-NEXT:    ret
   %mask = and i128 %s, 63
   %shl = shl i128 %r, %mask
   ret i128 %shl
 }
 
 define i128 @ashr(i128 %r, i128 %s) nounwind readnone {
 ; CHECK-LABEL: ashr:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    neg x8, x2
-; CHECK-NEXT:    lsl x8, x1, x8
-; CHECK-NEXT:    cmp x2, #0 // =0
-; CHECK-NEXT:    csel x8, xzr, x8, eq
-; CHECK-NEXT:    lsr x9, x0, x2
+; CHECK-NEXT:    lsr x8, x0, x2
+; CHECK-NEXT:    mvn w9, w2
+; CHECK-NEXT:    lsl x10, x1, #1
+; CHECK-NEXT:    lsl x9, x10, x9
 ; CHECK-NEXT:    orr x8, x9, x8
 ; CHECK-NEXT:    asr x9, x1, x2
-; CHECK-NEXT:    sub x10, x2, #64 // =64
-; CHECK-NEXT:    cmp x10, #0 // =0
+; CHECK-NEXT:    cmp x2, #64 // =64
 ; CHECK-NEXT:    csel x0, x9, x8, ge
 ; CHECK-NEXT:    asr x8, x1, #63
 ; CHECK-NEXT:    csel x1, x8, x9, ge
 ; CHECK-NEXT:    ret
   %shr = ashr i128 %r, %s
   ret i128 %shr
 }
 
; CHECK-NEXT:    ret
   %mask = and i128 %s, 63
   %shr = ashr i128 %r, %mask
   ret i128 %shr
 }
 
 define i128 @lshr(i128 %r, i128 %s) nounwind readnone {
 ; CHECK-LABEL: lshr:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    neg x8, x2
-; CHECK-NEXT:    lsl x8, x1, x8
-; CHECK-NEXT:    cmp x2, #0 // =0
-; CHECK-NEXT:    csel x8, xzr, x8, eq
-; CHECK-NEXT:    lsr x9, x0, x2
+; CHECK-NEXT:    lsr x8, x0, x2
+; CHECK-NEXT:    mvn w9, w2
+; CHECK-NEXT:    lsl x10, x1, #1
+; CHECK-NEXT:    lsl x9, x10, x9
 ; CHECK-NEXT:    orr x8, x9, x8
 ; CHECK-NEXT:    lsr x9, x1, x2
-; CHECK-NEXT:    sub x10, x2, #64 // =64
-; CHECK-NEXT:    cmp x10, #0 // =0
+; CHECK-NEXT:    cmp x2, #64 // =64
 ; CHECK-NEXT:    csel x0, x9, x8, ge
 ; CHECK-NEXT:    csel x1, xzr, x9, ge
 ; CHECK-NEXT:    ret
   %shr = lshr i128 %r, %s
   ret i128 %shr
 }
 
 define i128 @lshr_mask(i128 %r, i128 %s) nounwind readnone {