[LegalizeTypes] Legalize vector rotate operations

[rsanthir.quic]

Summary

Lower vector rotate operations as long as the legalization occurs outside of LegalizeVectorOps.

This fixes https://bugs.llvm.org/show_bug.cgi?id=47320

[RKSimon]

llvm/include/llvm/CodeGen/TargetLowering.h

4320–4321

Maybe call it AllowVectorOps ?

Add missing doxygen param.

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

6291–6292

clang-format this

[rsanthir.quic]

Changed the name or the param and added appropriate formatting

[rsanthir.quic]

updated usage of param

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

3517

Can you add a /*AllowVectorOps*/ comment to the true argument here and at the other call sites to help with readability?

[rsanthir.quic]

Added comment for readability

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

3517

Yes!

[dmgreen]

llvm/test/CodeGen/AArch64/expand-vector-rot.ll

2

Can you run update_llc_test_checks on this file.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

6291–6292

You need to add the VT.isVector() check as well

[rsanthir.quic]

Ran update_llc_test_checks.py and added missing isVector check

[RKSimon]

LGTM

[RKSimon]

@dmgreen Are you happy for this to be committed?

[dmgreen]

Oh yeah, sure. I'm happy if you were happy. Not crashing is certainly better than crashing.
LGTM.

[RKSimon]

@rsanthir.quic Are you OK to get this committed now?

[rsanthir.quic]

Yes! If you could commit it for me that would be great! Thanks

llvm/include/llvm/CodeGen/TargetLowering.h

   /// Expand funnel shift.
   /// \param N Node to expand
   /// \param Result output after conversion
   /// \returns True, if the expansion was successful, false otherwise
   bool expandFunnelShift(SDNode *N, SDValue &Result, SelectionDAG &DAG) const;
 
   /// Expand rotations.
   /// \param N Node to expand
+  /// \param AllowVectorOps expand vector rotate, this should only be performed
+  ///        if the legalization is happening outside of LegalizeVectorOps
   /// \param Result output after conversion
   /// \returns True, if the expansion was successful, false otherwise
-  bool expandROT(SDNode *N, SDValue &Result, SelectionDAG &DAG) const;
+  bool expandROT(SDNode *N, bool AllowVectorOps, SDValue &Result,
+                 SelectionDAG &DAG) const;
 
   /// Expand float(f32) to SINT(i64) conversion
   /// \param N Node to expand
   /// \param Result output after conversion
   /// \returns True, if the expansion was successful, false otherwise
   bool expandFP_TO_SINT(SDNode *N, SDValue &Result, SelectionDAG &DAG) const;
 
   /// Expand float to UINT conversion

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

bool SelectionDAGLegalize::ExpandNode(SDNode *Node) {
   }
   case ISD::FSHL:
   case ISD::FSHR:
     if (TLI.expandFunnelShift(Node, Tmp1, DAG))
       Results.push_back(Tmp1);
     break;
   case ISD::ROTL:
   case ISD::ROTR:
-    if (TLI.expandROT(Node, Tmp1, DAG))
+    if (TLI.expandROT(Node, true, Tmp1, DAG))
       Results.push_back(Tmp1);
     break;
   case ISD::SADDSAT:
   case ISD::UADDSAT:
   case ISD::SSUBSAT:
   case ISD::USUBSAT:
     Results.push_back(TLI.expandAddSubSat(Node, DAG));
     break;

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

SDValue DAGTypeLegalizer::PromoteIntRes_SRL(SDNode *N) {
   if (getTypeAction(RHS.getValueType()) == TargetLowering::TypePromoteInteger)
     RHS = ZExtPromotedInteger(RHS);
   return DAG.getNode(ISD::SRL, SDLoc(N), LHS.getValueType(), LHS, RHS);
 }
 
 SDValue DAGTypeLegalizer::PromoteIntRes_Rotate(SDNode *N) {
   // Lower the rotate to shifts and ORs which can be promoted.
   SDValue Res;
-  TLI.expandROT(N, Res, DAG);
+  TLI.expandROT(N, true, Res, DAG);
   ReplaceValueWith(SDValue(N, 0), Res);
   return SDValue();
 }
 
 SDValue DAGTypeLegalizer::PromoteIntRes_FunnelShift(SDNode *N) {
   SDValue Hi = GetPromotedInteger(N->getOperand(0));
   SDValue Lo = GetPromotedInteger(N->getOperand(1));
   SDValue Amount = GetPromotedInteger(N->getOperand(2));
void DAGTypeLegalizer::ExpandIntRes_VECREDUCE(SDNode *N,
   SDValue Res = TLI.expandVecReduce(N, DAG);
   SplitInteger(Res, Lo, Hi);
 }
 
 void DAGTypeLegalizer::ExpandIntRes_Rotate(SDNode *N,
                                            SDValue &Lo, SDValue &Hi) {
   // Lower the rotate to shifts and ORs which can be expanded.
   SDValue Res;
-  TLI.expandROT(N, Res, DAG);
+  TLI.expandROT(N, true, Res, DAG);
   SplitInteger(Res, Lo, Hi);
 }
 
 void DAGTypeLegalizer::ExpandIntRes_FunnelShift(SDNode *N,
                                                 SDValue &Lo, SDValue &Hi) {
   // Lower the funnel shift to shifts and ORs which can be expanded.
   SDValue Res;
   TLI.expandFunnelShift(N, Res, DAG);

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp

void VectorLegalizer::Expand(SDNode *Node, SmallVectorImpl<SDValue> &Results) {
   case ISD::FSHR:
     if (TLI.expandFunnelShift(Node, Tmp, DAG)) {
       Results.push_back(Tmp);
       return;
     }
     break;
   case ISD::ROTL:
   case ISD::ROTR:
-    if (TLI.expandROT(Node, Tmp, DAG)) {
+    if (TLI.expandROT(Node, true, Tmp, DAG)) {
       Results.push_back(Tmp);
       return;
     }
     break;
   case ISD::FMINNUM:
   case ISD::FMAXNUM:
     if (SDValue Expanded = TLI.expandFMINNUM_FMAXNUM(Node, DAG)) {
       Results.push_back(Expanded);

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

if (IsFSHL) {
       ShY = DAG.getNode(ISD::SRL, DL, VT, Y, ShAmt);
     }
   }
   Result = DAG.getNode(ISD::OR, DL, VT, ShX, ShY);
   return true;
 }
 
 // TODO: Merge with expandFunnelShift.
-bool TargetLowering::expandROT(SDNode *Node, SDValue &Result,
-                               SelectionDAG &DAG) const {
+bool TargetLowering::expandROT(SDNode *Node, bool AllowVectorOps,
+                               SDValue &Result, SelectionDAG &DAG) const {
   EVT VT = Node->getValueType(0);
   unsigned EltSizeInBits = VT.getScalarSizeInBits();
   bool IsLeft = Node->getOpcode() == ISD::ROTL;
   SDValue Op0 = Node->getOperand(0);
   SDValue Op1 = Node->getOperand(1);
   SDLoc DL(SDValue(Node, 0));
 
   EVT ShVT = Op1.getValueType();
   SDValue Zero = DAG.getConstant(0, DL, ShVT);
 
   // If a rotate in the other direction is supported, use it.
   unsigned RevRot = IsLeft ? ISD::ROTR : ISD::ROTL;
   if (isOperationLegalOrCustom(RevRot, VT) && isPowerOf2_32(EltSizeInBits)) {
     SDValue Sub = DAG.getNode(ISD::SUB, DL, ShVT, Zero, Op1);
     Result = DAG.getNode(RevRot, DL, VT, Op0, Sub);
     return true;
   }
 
-  if (VT.isVector() && (!isOperationLegalOrCustom(ISD::SHL, VT) ||
+  if (!AllowVectorOps && (!isOperationLegalOrCustom(ISD::SHL, VT) ||
                           !isOperationLegalOrCustom(ISD::SRL, VT) ||
                           !isOperationLegalOrCustom(ISD::SUB, VT) ||
                           !isOperationLegalOrCustomOrPromote(ISD::OR, VT) ||
                           !isOperationLegalOrCustomOrPromote(ISD::AND, VT)))
     return false;
 
   unsigned ShOpc = IsLeft ? ISD::SHL : ISD::SRL;
   unsigned HsOpc = IsLeft ? ISD::SRL : ISD::SHL;
   SDValue BitWidthMinusOneC = DAG.getConstant(EltSizeInBits - 1, DL, ShVT);
   SDValue ShVal;
   SDValue HsVal;
   if (isPowerOf2_32(EltSizeInBits)) {

llvm/test/CodeGen/AArch64/expand-vector-rot.ll

+; RUN: llc < %s -mtriple=aarch64-linux-android | FileCheck %s
+
+declare <2 x i16> @llvm.fshl.v2i16(<2 x i16>, <2 x i16>, <2 x i16>)
+
+define <2 x i16>  @rotlv2_16(<2 x i16> %vec2_16, <2 x i16> %shift) {
+; CHECK-LABEL: rotlv2_16:
+; CHECK: ushl v
+; CHECK: ushl v
+  %1 = call <2 x i16> @llvm.fshl.v2i16(<2 x i16> %vec2_16, <2 x i16> %vec2_16, <2 x i16> %shift)
+  ret <2 x i16> %1
+}