[DAGCombiner] Recognise vector rotations with non-splat constants

[andrew.zhogin]

Summary

Fix for bug 33691: https://bugs.llvm.org/show_bug.cgi?id=33691 .
Such code now will be converted into vector rotate operation:

%1 = lshr <4 x i32> %x, <i32 1, i32 2, i32 3, i32 4>
%2 = shl <4 x i32> %x, <i32 31, i32 30, i32 29, i32 28>
%3 = or <4 x i32> %1, %2

[RKSimon]

Always add llvm-commits as a subscriber for llvm patches so that it has visibility on the mailing list (clang patches should subscribe cfe-commits instead)

[RKSimon]

There are a number of existing tests in files that should have been improved by this:

• llvm\test\CodeGen\X86\combine-rotates.ll
• llvm\test\CodeGen\X86\vector-rotate-128.ll
• llvm\test\CodeGen\X86\vector-rotate-256.ll

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4708

It would be tidier to merge this into the scalar/splat version above. Adding a predicate to pattern match all constants would be useful (and possibly reusable elsewhere - @majnemer has suggested a DAGCombiner equivalent of llvm::PatternMatch in the past).

The bitmasks should be creatable using SelectionDAG::FoldConstantArithmetic from logical shifts of the all ones mask and LHSShiftAmt/RHSShiftAmt.

[andrew.zhogin]

Updated related CodeGen tests.
Implemented unified code for scalar/vector rotr combining.

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4691

This can be done with

SDValue RHSBits = DAG.getNode(ISD::SRL, DL, VT, AllOnes, RHSShiftAmt);

(Taking AllOnes from the value initially assigned to Mask).

4697

This can be done with

SDValue LHSBits = DAG.getNode(ISD::SHL, DL, VT, AllOnes, LHSShiftAmt);

(Taking AllOnes from the value initially assigned to Mask).

[andrew.zhogin]

buildLoHiRotateVecBitmask helper function deleted, code inlined.

[andrew.zhogin]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4697

FoldConstantArithmetic is not necessary here?

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4591

There's no need for this to be part of DAGCombiner - just make it static and pass SelectionDAG& as an argument

4688

Probably better to call this AllOnes not One

4691

Probably safer to use getNode instead of FoldConstantArithmetic, same for below

[andrew.zhogin]

matchRotateSub converted into static function. getNode() instead of FoldConstantArithmetic for shift masks. "AllOnes" instead of "One".

[RKSimon]

LGTM - thanks

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

private:
     SDValue buildSqrtEstimateImpl(SDValue Op, SDNodeFlags Flags, bool Recip);
     SDValue buildSqrtNROneConst(SDValue Op, SDValue Est, unsigned Iterations,
                                 SDNodeFlags Flags, bool Reciprocal);
     SDValue buildSqrtNRTwoConst(SDValue Op, SDValue Est, unsigned Iterations,
                                 SDNodeFlags Flags, bool Reciprocal);
     SDValue MatchBSwapHWordLow(SDNode *N, SDValue N0, SDValue N1,
                                bool DemandHighBits = true);
     SDValue MatchBSwapHWord(SDNode *N, SDValue N0, SDValue N1);
+    SDValue buildLoHiRotateVecBitmask(SDValue Shift, bool Low,
+                                      const SDLoc &DL, EVT VT);
     SDNode *MatchRotatePosNeg(SDValue Shifted, SDValue Pos, SDValue Neg,
                               SDValue InnerPos, SDValue InnerNeg,
                               unsigned PosOpcode, unsigned NegOpcode,
                               const SDLoc &DL);
     SDNode *MatchRotate(SDValue LHS, SDValue RHS, const SDLoc &DL);
     SDValue MatchLoadCombine(SDNode *N);
     SDValue ReduceLoadWidth(SDNode *N);
     SDValue ReduceLoadOpStoreWidth(SDNode *N);
if (matchRotateSub(InnerPos, InnerNeg, VT.getScalarSizeInBits())) {
     bool HasPos = TLI.isOperationLegalOrCustom(PosOpcode, VT);
     return DAG.getNode(HasPos ? PosOpcode : NegOpcode, DL, VT, Shifted,
                        HasPos ? Pos : Neg).getNode();
   }
 
   return nullptr;
 }
 
+// Are constant vectors combinable for rotate.
+// Like <1, 2, 3, 4> + <31, 30, 29, 28> for int32 vectors.
+static bool areNonSplatRotateCombinableVecs(
+    SDValue LShift, SDValue RShift, unsigned EltSizeInBits) {
+  BuildVectorSDNode *LV = dyn_cast<BuildVectorSDNode>(LShift);
+  BuildVectorSDNode *RV = dyn_cast<BuildVectorSDNode>(RShift);
+  if (!LV || !RV) return false;
+  if (LV->getNumOperands() != RV->getNumOperands()) return false;
+
+  for (unsigned i = 0, e = LV->getNumOperands(); i != e; ++i) {
+    SDValue OpL = LV->getOperand(i);
+    SDValue OpR = RV->getOperand(i);
+    if (OpL.isUndef() || OpR.isUndef())
+      return false;
+    ConstantSDNode *LConst = dyn_cast_or_null<ConstantSDNode>(OpL);
+    ConstantSDNode *RConst = dyn_cast_or_null<ConstantSDNode>(OpR);
+    if (!LConst || !RConst ||
+         LConst->getZExtValue() + RConst->getZExtValue() != EltSizeInBits)
+     return false;
+  }
+  return true;
+}
+
+// buildLoHiRotateVecBitmask - Generate vector of low/high bitmasks for vector
+// of constant integer values (shift sizes).
+// <0, 1, 2, 3> --> <0b0..0000, 0b0..0001, 0b0..0011, 0b0..0111> for Low=true
+// <0, 1, 2, 3> --> <0b0000..0, 0b1000..0, 0b1100..0, 0b1110..0> for Low=false
+SDValue DAGCombiner::buildLoHiRotateVecBitmask(SDValue Shift, bool Low,
+                                               const SDLoc &DL, EVT VT) {
+  unsigned EltSizeInBits = VT.getScalarSizeInBits();
+  const BuildVectorSDNode &BV = *cast<BuildVectorSDNode>(Shift);
+
+  SmallVector<SDValue, 8> Ops;
+  for (unsigned i = 0, e = BV.getNumOperands(); i != e; ++i) {
+    ConstantSDNode &Op = *cast<ConstantSDNode>(BV.getOperand(i));
+    APInt Bits = Low ? APInt::getLowBitsSet(EltSizeInBits, Op.getZExtValue())
+                     : APInt::getHighBitsSet(EltSizeInBits, Op.getZExtValue());
+    Ops.push_back(DAG.getConstant(Bits, DL, VT.getScalarType()));
+  }
+  return DAG.getBuildVector(VT, DL, Ops);
+}
+
 // MatchRotate - Handle an 'or' of two operands.  If this is one of the many
 // idioms for rotate, and if the target supports rotation instructions, generate
 // a rot[lr].
 SDNode *DAGCombiner::MatchRotate(SDValue LHS, SDValue RHS, const SDLoc &DL) {
   // Must be a legal type.  Expanded 'n promoted things won't work with rotates.
   EVT VT = LHS.getValueType();
   if (!TLI.isTypeLegal(VT)) return nullptr;
 
if (isConstOrConstSplat(LHSShiftAmt) && isConstOrConstSplat(RHSShiftAmt)) {
     if ((LShVal + RShVal) != EltSizeInBits)
       return nullptr;
 
     SDValue Rot = DAG.getNode(HasROTL ? ISD::ROTL : ISD::ROTR, DL, VT,
                               LHSShiftArg, HasROTL ? LHSShiftAmt : RHSShiftAmt);
 
     // If there is an AND of either shifted operand, apply it to the result.
     if (LHSMask.getNode() || RHSMask.getNode()) {
       SDValue Mask = DAG.getAllOnesConstant(DL, VT);
 
       if (LHSMask.getNode()) {
         APInt RHSBits = APInt::getLowBitsSet(EltSizeInBits, LShVal);
         Mask = DAG.getNode(ISD::AND, DL, VT, Mask,
                            DAG.getNode(ISD::OR, DL, VT, LHSMask,
                                        DAG.getConstant(RHSBits, DL, VT)));
       }
       if (RHSMask.getNode()) {
         APInt LHSBits = APInt::getHighBitsSet(EltSizeInBits, RShVal);
         Mask = DAG.getNode(ISD::AND, DL, VT, Mask,
                            DAG.getNode(ISD::OR, DL, VT, RHSMask,
                                        DAG.getConstant(LHSBits, DL, VT)));
       }
 
       Rot = DAG.getNode(ISD::AND, DL, VT, Rot, Mask);
     }
 
     return Rot.getNode();
   }
+  if (areNonSplatRotateCombinableVecs(LHSShiftAmt, RHSShiftAmt, EltSizeInBits)) {
+    SDValue Rot = DAG.getNode(HasROTL ? ISD::ROTL : ISD::ROTR, DL, VT,
+                              LHSShiftArg, HasROTL ? LHSShiftAmt : RHSShiftAmt);
+    if (LHSMask.getNode() || RHSMask.getNode()) {
+      if (LHSMask.getNode() && RHSMask.getNode()) {
+        // calculating combined vector mask
+        // (LHSMask | 000...111) & (111...000 | RHSMask)
+        SDValue RHSBits = buildLoHiRotateVecBitmask(LHSShiftAmt, true, DL, VT);
+        SDValue LHSBits = buildLoHiRotateVecBitmask(RHSShiftAmt, false, DL, VT);
+
+        SDValue LMask = DAG.getNode(ISD::OR, DL, VT, LHSMask, RHSBits);
+        SDValue RMask = DAG.getNode(ISD::OR, DL, VT, RHSMask, LHSBits);
+
+        SDValue Mask = DAG.getNode(ISD::AND, DL, VT, LMask, RMask);
+        Rot = DAG.getNode(ISD::AND, DL, VT, Rot, Mask);
+      }
+      else if (LHSMask.getNode()) {
+        SDValue RHSBits = buildLoHiRotateVecBitmask(LHSShiftAmt, true, DL, VT);
+        SDValue LMask = DAG.getNode(ISD::OR, DL, VT, LHSMask, RHSBits);
+        Rot = DAG.getNode(ISD::AND, DL, VT, Rot, LMask);
+      }
+      else if (RHSMask.getNode()) {
+        SDValue LHSBits = buildLoHiRotateVecBitmask(RHSShiftAmt, false, DL, VT);
+        SDValue RMask = DAG.getNode(ISD::OR, DL, VT, RHSMask, LHSBits);
+        Rot = DAG.getNode(ISD::AND, DL, VT, Rot, RMask);
+      }
+    }
+    return Rot.getNode();
+  }
 
   // If there is a mask here, and we have a variable shift, we can't be sure
   // that we're masking out the right stuff.
   if (LHSMask.getNode() || RHSMask.getNode())
     return nullptr;
 
   // If the shift amount is sign/zext/any-extended just peel it off.
   SDValue LExtOp0 = LHSShiftAmt;

test/CodeGen/X86/rotate_vec.ll

; CHECK-NEXT:    retq
   %2 = shl <4 x i32> %x, <i32 31, i32 31, i32 31, i32 31>
   %3 = or <4 x i32> %1, %2
   ret <4 x i32> %3
 }
 
 define <4 x i32> @rot_v4i32_non_splat(<4 x i32> %x) {
 ; CHECK-LABEL: rot_v4i32_non_splat:
 ; CHECK:       # BB#0:
-; CHECK-NEXT:    vpsrlvd {{.*}}(%rip), %xmm0, %xmm1
-; CHECK-NEXT:    vpsllvd {{.*}}(%rip), %xmm0, %xmm0
-; CHECK-NEXT:    vpor %xmm0, %xmm1, %xmm0
+; CHECK-NEXT:    vprotd {{.*}}(%rip), %xmm0, %xmm0
 ; CHECK-NEXT:    retq
   %1 = lshr <4 x i32> %x, <i32 1, i32 2, i32 3, i32 4>
   %2 = shl <4 x i32> %x, <i32 31, i32 30, i32 29, i32 28>
   %3 = or <4 x i32> %1, %2
   ret <4 x i32> %3
 }
 
 define <4 x i32> @rot_v4i32_splat_2masks(<4 x i32> %x) {
; CHECK-NEXT:    retq
   %4 = and <4 x i32> %3, <i32 0, i32 4294901760, i32 0, i32 4294901760>
   %5 = or <4 x i32> %2, %4
   ret <4 x i32> %5
 }
 
 define <4 x i32> @rot_v4i32_non_splat_2masks(<4 x i32> %x) {
 ; CHECK-LABEL: rot_v4i32_non_splat_2masks:
 ; CHECK:       # BB#0:
-; CHECK-NEXT:    vpsrlvd {{.*}}(%rip), %xmm0, %xmm1
-; CHECK-NEXT:    vpxor %xmm2, %xmm2, %xmm2
-; CHECK-NEXT:    vpblendw {{.*#+}} xmm1 = xmm2[0],xmm1[1],xmm2[2],xmm1[3],xmm2[4],xmm1[5],xmm2[6],xmm1[7]
-; CHECK-NEXT:    vpsllvd {{.*}}(%rip), %xmm0, %xmm0
-; CHECK-NEXT:    vpblendw {{.*#+}} xmm0 = xmm2[0,1,2],xmm0[3],xmm2[4,5,6],xmm0[7]
-; CHECK-NEXT:    vpor %xmm0, %xmm1, %xmm0
+; CHECK-NEXT:    vprotd {{.*}}(%rip), %xmm0, %xmm0
+; CHECK-NEXT:    vpand {{.*}}(%rip), %xmm0, %xmm0
 ; CHECK-NEXT:    retq
   %1 = lshr <4 x i32> %x, <i32 1, i32 2, i32 3, i32 4>
   %2 = and <4 x i32> %1, <i32 4294901760, i32 4294901760, i32 4294901760, i32 4294901760>
 
   %3 = shl <4 x i32> %x, <i32 31, i32 30, i32 29, i32 28>
   %4 = and <4 x i32> %3, <i32 0, i32 4294901760, i32 0, i32 4294901760>
   %5 = or <4 x i32> %2, %4
   ret <4 x i32> %5
 }