[RISCV] Support extraction of misaligned subvectors

This patch extends the support for RVV EXTRACT_SUBVECTOR to cover those
which don't align to a vector register boundary. It accomplishes this by
extracting the nearest register-sized subvector (a subregister
operation), then sliding the vector down with VSLIDEDOWN and extracting
the subvector from the first position (a COPY operation).

Since this procedure involves the use of VSCALE and multiplication, the
handling of such operations is done during lowering to simplify the
implementation and make use of DAG combining. This necessitated moving
some helper functions from RISCVISelDAGToDAG to RISCVTargetLowering.

Reviewed By: craig.topper

Differential Revision: https://reviews.llvm.org/D96959

Author: frasercrmck

llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp

@@ -13,6 +13,7 @@
 #include "RISCVISelDAGToDAG.h"
 #include "MCTargetDesc/RISCVMCTargetDesc.h"
 #include "MCTargetDesc/RISCVMatInt.h"
+#include "RISCVISelLowering.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/IR/IntrinsicsRISCV.h"
 #include "llvm/Support/Alignment.h"
@@ -62,64 +63,6 @@ static SDNode *selectImm(SelectionDAG *CurDAG, const SDLoc &DL, int64_t Imm,
   return Result;
 }
 
-static RISCVVLMUL getLMUL(MVT VT) {
-  switch (VT.getSizeInBits().getKnownMinValue() / 8) {
-  default:
-    llvm_unreachable("Invalid LMUL.");
-  case 1:
-    return RISCVVLMUL::LMUL_F8;
-  case 2:
-    return RISCVVLMUL::LMUL_F4;
-  case 4:
-    return RISCVVLMUL::LMUL_F2;
-  case 8:
-    return RISCVVLMUL::LMUL_1;
-  case 16:
-    return RISCVVLMUL::LMUL_2;
-  case 32:
-    return RISCVVLMUL::LMUL_4;
-  case 64:
-    return RISCVVLMUL::LMUL_8;
-  }
-}
-
-static unsigned getRegClassIDForLMUL(RISCVVLMUL LMul) {
-  switch (LMul) {
-  default:
-    llvm_unreachable("Invalid LMUL.");
-  case RISCVVLMUL::LMUL_F8:
-  case RISCVVLMUL::LMUL_F4:
-  case RISCVVLMUL::LMUL_F2:
-  case RISCVVLMUL::LMUL_1:
-    return RISCV::VRRegClassID;
-  case RISCVVLMUL::LMUL_2:
-    return RISCV::VRM2RegClassID;
-  case RISCVVLMUL::LMUL_4:
-    return RISCV::VRM4RegClassID;
-  case RISCVVLMUL::LMUL_8:
-    return RISCV::VRM8RegClassID;
-  }
-}
-
-static unsigned getSubregIndexByMVT(MVT VT, unsigned Index) {
-  RISCVVLMUL LMUL = getLMUL(VT);
-  if (LMUL == RISCVVLMUL::LMUL_F8 || LMUL == RISCVVLMUL::LMUL_F4 ||
-      LMUL == RISCVVLMUL::LMUL_F2 || LMUL == RISCVVLMUL::LMUL_1) {
-    static_assert(RISCV::sub_vrm1_7 == RISCV::sub_vrm1_0 + 7,
-                  "Unexpected subreg numbering");
-    return RISCV::sub_vrm1_0 + Index;
-  } else if (LMUL == RISCVVLMUL::LMUL_2) {
-    static_assert(RISCV::sub_vrm2_3 == RISCV::sub_vrm2_0 + 3,
-                  "Unexpected subreg numbering");
-    return RISCV::sub_vrm2_0 + Index;
-  } else if (LMUL == RISCVVLMUL::LMUL_4) {
-    static_assert(RISCV::sub_vrm4_1 == RISCV::sub_vrm4_0 + 1,
-                  "Unexpected subreg numbering");
-    return RISCV::sub_vrm4_0 + Index;
-  }
-  llvm_unreachable("Invalid vector type.");
-}
-
 static SDValue createTupleImpl(SelectionDAG &CurDAG, ArrayRef<SDValue> Regs,
                                unsigned RegClassID, unsigned SubReg0) {
   assert(Regs.size() >= 2 && Regs.size() <= 8);
@@ -187,7 +130,7 @@ void RISCVDAGToDAGISel::selectVLSEG(SDNode *Node, bool IsMasked,
   MVT VT = Node->getSimpleValueType(0);
   unsigned ScalarSize = VT.getScalarSizeInBits();
   MVT XLenVT = Subtarget->getXLenVT();
-  RISCVVLMUL LMUL = getLMUL(VT);
+  RISCVVLMUL LMUL = RISCVTargetLowering::getLMUL(VT);
   SDValue SEW = CurDAG->getTargetConstant(ScalarSize, DL, XLenVT);
   unsigned CurOp = 2;
   SmallVector<SDValue, 7> Operands;
@@ -218,10 +161,11 @@ void RISCVDAGToDAGISel::selectVLSEG(SDNode *Node, bool IsMasked,
     CurDAG->setNodeMemRefs(Load, {MemOp->getMemOperand()});
 
   SDValue SuperReg = SDValue(Load, 0);
-  for (unsigned I = 0; I < NF; ++I)
+  for (unsigned I = 0; I < NF; ++I) {
+    unsigned SubRegIdx = RISCVTargetLowering::getSubregIndexByMVT(VT, I);
     ReplaceUses(SDValue(Node, I),
-                CurDAG->getTargetExtractSubreg(getSubregIndexByMVT(VT, I), DL,
-                                               VT, SuperReg));
+                CurDAG->getTargetExtractSubreg(SubRegIdx, DL, VT, SuperReg));
+  }
 
   ReplaceUses(SDValue(Node, NF), SDValue(Load, 1));
   CurDAG->RemoveDeadNode(Node);
@@ -233,7 +177,7 @@ void RISCVDAGToDAGISel::selectVLSEGFF(SDNode *Node, bool IsMasked) {
   MVT VT = Node->getSimpleValueType(0);
   MVT XLenVT = Subtarget->getXLenVT();
   unsigned ScalarSize = VT.getScalarSizeInBits();
-  RISCVVLMUL LMUL = getLMUL(VT);
+  RISCVVLMUL LMUL = RISCVTargetLowering::getLMUL(VT);
   SDValue SEW = CurDAG->getTargetConstant(ScalarSize, DL, XLenVT);
 
   unsigned CurOp = 2;
@@ -265,10 +209,11 @@ void RISCVDAGToDAGISel::selectVLSEGFF(SDNode *Node, bool IsMasked) {
     CurDAG->setNodeMemRefs(Load, {MemOp->getMemOperand()});
 
   SDValue SuperReg = SDValue(Load, 0);
-  for (unsigned I = 0; I < NF; ++I)
+  for (unsigned I = 0; I < NF; ++I) {
+    unsigned SubRegIdx = RISCVTargetLowering::getSubregIndexByMVT(VT, I);
     ReplaceUses(SDValue(Node, I),
-                CurDAG->getTargetExtractSubreg(getSubregIndexByMVT(VT, I), DL,
-                                               VT, SuperReg));
+                CurDAG->getTargetExtractSubreg(SubRegIdx, DL, VT, SuperReg));
+  }
 
   ReplaceUses(SDValue(Node, NF), SDValue(ReadVL, 0));   // VL
   ReplaceUses(SDValue(Node, NF + 1), SDValue(Load, 1)); // Chain
@@ -282,7 +227,7 @@ void RISCVDAGToDAGISel::selectVLXSEG(SDNode *Node, bool IsMasked,
   MVT VT = Node->getSimpleValueType(0);
   unsigned ScalarSize = VT.getScalarSizeInBits();
   MVT XLenVT = Subtarget->getXLenVT();
-  RISCVVLMUL LMUL = getLMUL(VT);
+  RISCVVLMUL LMUL = RISCVTargetLowering::getLMUL(VT);
   SDValue SEW = CurDAG->getTargetConstant(ScalarSize, DL, XLenVT);
   unsigned CurOp = 2;
   SmallVector<SDValue, 7> Operands;
@@ -307,7 +252,7 @@ void RISCVDAGToDAGISel::selectVLXSEG(SDNode *Node, bool IsMasked,
   assert(VT.getVectorElementCount() == IndexVT.getVectorElementCount() &&
          "Element count mismatch");
 
-  RISCVVLMUL IndexLMUL = getLMUL(IndexVT);
+  RISCVVLMUL IndexLMUL = RISCVTargetLowering::getLMUL(IndexVT);
   unsigned IndexScalarSize = IndexVT.getScalarSizeInBits();
   const RISCV::VLXSEGPseudo *P = RISCV::getVLXSEGPseudo(
       NF, IsMasked, IsOrdered, IndexScalarSize, static_cast<unsigned>(LMUL),
@@ -319,10 +264,11 @@ void RISCVDAGToDAGISel::selectVLXSEG(SDNode *Node, bool IsMasked,
     CurDAG->setNodeMemRefs(Load, {MemOp->getMemOperand()});
 
   SDValue SuperReg = SDValue(Load, 0);
-  for (unsigned I = 0; I < NF; ++I)
+  for (unsigned I = 0; I < NF; ++I) {
+    unsigned SubRegIdx = RISCVTargetLowering::getSubregIndexByMVT(VT, I);
     ReplaceUses(SDValue(Node, I),
-                CurDAG->getTargetExtractSubreg(getSubregIndexByMVT(VT, I), DL,
-                                               VT, SuperReg));
+                CurDAG->getTargetExtractSubreg(SubRegIdx, DL, VT, SuperReg));
+  }
 
   ReplaceUses(SDValue(Node, NF), SDValue(Load, 1));
   CurDAG->RemoveDeadNode(Node);
@@ -339,7 +285,7 @@ void RISCVDAGToDAGISel::selectVSSEG(SDNode *Node, bool IsMasked,
   MVT VT = Node->getOperand(2)->getSimpleValueType(0);
   unsigned ScalarSize = VT.getScalarSizeInBits();
   MVT XLenVT = Subtarget->getXLenVT();
-  RISCVVLMUL LMUL = getLMUL(VT);
+  RISCVVLMUL LMUL = RISCVTargetLowering::getLMUL(VT);
   SDValue SEW = CurDAG->getTargetConstant(ScalarSize, DL, XLenVT);
   SmallVector<SDValue, 8> Regs(Node->op_begin() + 2, Node->op_begin() + 2 + NF);
   SDValue StoreVal = createTuple(*CurDAG, Regs, NF, LMUL);
@@ -376,7 +322,7 @@ void RISCVDAGToDAGISel::selectVSXSEG(SDNode *Node, bool IsMasked,
   MVT VT = Node->getOperand(2)->getSimpleValueType(0);
   unsigned ScalarSize = VT.getScalarSizeInBits();
   MVT XLenVT = Subtarget->getXLenVT();
-  RISCVVLMUL LMUL = getLMUL(VT);
+  RISCVVLMUL LMUL = RISCVTargetLowering::getLMUL(VT);
   SDValue SEW = CurDAG->getTargetConstant(ScalarSize, DL, XLenVT);
   SmallVector<SDValue, 7> Operands;
   SmallVector<SDValue, 8> Regs(Node->op_begin() + 2, Node->op_begin() + 2 + NF);
@@ -397,7 +343,7 @@ void RISCVDAGToDAGISel::selectVSXSEG(SDNode *Node, bool IsMasked,
   assert(VT.getVectorElementCount() == IndexVT.getVectorElementCount() &&
          "Element count mismatch");
 
-  RISCVVLMUL IndexLMUL = getLMUL(IndexVT);
+  RISCVVLMUL IndexLMUL = RISCVTargetLowering::getLMUL(IndexVT);
   unsigned IndexScalarSize = IndexVT.getScalarSizeInBits();
   const RISCV::VSXSEGPseudo *P = RISCV::getVSXSEGPseudo(
       NF, IsMasked, IsOrdered, IndexScalarSize, static_cast<unsigned>(LMUL),
@@ -411,47 +357,6 @@ void RISCVDAGToDAGISel::selectVSXSEG(SDNode *Node, bool IsMasked,
   ReplaceNode(Node, Store);
 }
 
-static unsigned getRegClassIDForVecVT(MVT VT) {
-  if (VT.getVectorElementType() == MVT::i1)
-    return RISCV::VRRegClassID;
-  return getRegClassIDForLMUL(getLMUL(VT));
-}
-
-// Attempt to decompose a subvector insert/extract between VecVT and
-// SubVecVT via subregister indices. Returns the subregister index that
-// can perform the subvector insert/extract with the given element index, as
-// well as the index corresponding to any leftover subvectors that must be
-// further inserted/extracted within the register class for SubVecVT.
-static std::pair<unsigned, unsigned>
-decomposeSubvectorInsertExtractToSubRegs(MVT VecVT, MVT SubVecVT,
-                                         unsigned InsertExtractIdx,
-                                         const RISCVRegisterInfo *TRI) {
-  static_assert((RISCV::VRM8RegClassID > RISCV::VRM4RegClassID &&
-                 RISCV::VRM4RegClassID > RISCV::VRM2RegClassID &&
-                 RISCV::VRM2RegClassID > RISCV::VRRegClassID),
-                "Register classes not ordered");
-  unsigned VecRegClassID = getRegClassIDForVecVT(VecVT);
-  unsigned SubRegClassID = getRegClassIDForVecVT(SubVecVT);
-  // Try to compose a subregister index that takes us from the incoming
-  // LMUL>1 register class down to the outgoing one. At each step we half
-  // the LMUL:
-  //   nxv16i32@12 -> nxv2i32: sub_vrm4_1_then_sub_vrm2_1_then_sub_vrm1_0
-  // Note that this is not guaranteed to find a subregister index, such as
-  // when we are extracting from one VR type to another.
-  unsigned SubRegIdx = RISCV::NoSubRegister;
-  for (const unsigned RCID :
-       {RISCV::VRM4RegClassID, RISCV::VRM2RegClassID, RISCV::VRRegClassID})
-    if (VecRegClassID > RCID && SubRegClassID <= RCID) {
-      VecVT = VecVT.getHalfNumVectorElementsVT();
-      bool IsHi =
-          InsertExtractIdx >= VecVT.getVectorElementCount().getKnownMinValue();
-      SubRegIdx = TRI->composeSubRegIndices(SubRegIdx,
-                                            getSubregIndexByMVT(VecVT, IsHi));
-      if (IsHi)
-        InsertExtractIdx -= VecVT.getVectorElementCount().getKnownMinValue();
-    }
-  return {SubRegIdx, InsertExtractIdx};
-}
 
 void RISCVDAGToDAGISel::Select(SDNode *Node) {
   // If we have a custom node, we have already selected.
@@ -726,8 +631,8 @@ void RISCVDAGToDAGISel::Select(SDNode *Node) {
       assert(VT.getVectorElementCount() == IndexVT.getVectorElementCount() &&
              "Element count mismatch");
 
-      RISCVVLMUL LMUL = getLMUL(VT);
-      RISCVVLMUL IndexLMUL = getLMUL(IndexVT);
+      RISCVVLMUL LMUL = RISCVTargetLowering::getLMUL(VT);
+      RISCVVLMUL IndexLMUL = RISCVTargetLowering::getLMUL(IndexVT);
       unsigned IndexScalarSize = IndexVT.getScalarSizeInBits();
       const RISCV::VLX_VSXPseudo *P = RISCV::getVLXPseudo(
           IsMasked, IsOrdered, IndexScalarSize, static_cast<unsigned>(LMUL),
@@ -855,8 +760,8 @@ void RISCVDAGToDAGISel::Select(SDNode *Node) {
       assert(VT.getVectorElementCount() == IndexVT.getVectorElementCount() &&
              "Element count mismatch");
 
-      RISCVVLMUL LMUL = getLMUL(VT);
-      RISCVVLMUL IndexLMUL = getLMUL(IndexVT);
+      RISCVVLMUL LMUL = RISCVTargetLowering::getLMUL(VT);
+      RISCVVLMUL IndexLMUL = RISCVTargetLowering::getLMUL(IndexVT);
       unsigned IndexScalarSize = IndexVT.getScalarSizeInBits();
       const RISCV::VLX_VSXPseudo *P = RISCV::getVSXPseudo(
           IsMasked, IsOrdered, IndexScalarSize, static_cast<unsigned>(LMUL),
@@ -895,7 +800,7 @@ void RISCVDAGToDAGISel::Select(SDNode *Node) {
     // For now, keep the two paths separate.
     if (VT.isScalableVector() && SubVecVT.isScalableVector()) {
       bool IsFullVecReg = false;
-      switch (getLMUL(SubVecVT)) {
+      switch (RISCVTargetLowering::getLMUL(SubVecVT)) {
       default:
         break;
       case RISCVVLMUL::LMUL_1:
@@ -915,10 +820,11 @@ void RISCVDAGToDAGISel::Select(SDNode *Node) {
       const auto *TRI = Subtarget->getRegisterInfo();
       unsigned SubRegIdx;
       std::tie(SubRegIdx, Idx) =
-          decomposeSubvectorInsertExtractToSubRegs(VT, SubVecVT, Idx, TRI);
+          RISCVTargetLowering::decomposeSubvectorInsertExtractToSubRegs(
+              VT, SubVecVT, Idx, TRI);
 
       // If the Idx hasn't been completely eliminated then this is a subvector
-      // extract which doesn't naturally align to a vector register. These must
+      // insert which doesn't naturally align to a vector register. These must
       // be handled using instructions to manipulate the vector registers.
       if (Idx != 0)
         break;
@@ -936,7 +842,7 @@ void RISCVDAGToDAGISel::Select(SDNode *Node) {
       if (!Node->getOperand(0).isUndef())
         break;
 
-      unsigned RegClassID = getRegClassIDForVecVT(VT);
+      unsigned RegClassID = RISCVTargetLowering::getRegClassIDForVecVT(VT);
 
       SDValue RC =
           CurDAG->getTargetConstant(RegClassID, DL, Subtarget->getXLenVT());
@@ -961,7 +867,8 @@ void RISCVDAGToDAGISel::Select(SDNode *Node) {
       const auto *TRI = Subtarget->getRegisterInfo();
       unsigned SubRegIdx;
       std::tie(SubRegIdx, Idx) =
-          decomposeSubvectorInsertExtractToSubRegs(InVT, VT, Idx, TRI);
+          RISCVTargetLowering::decomposeSubvectorInsertExtractToSubRegs(
+              InVT, VT, Idx, TRI);
 
       // If the Idx hasn't been completely eliminated then this is a subvector
       // extract which doesn't naturally align to a vector register. These must
@@ -972,8 +879,10 @@ void RISCVDAGToDAGISel::Select(SDNode *Node) {
       // If we haven't set a SubRegIdx, then we must be going between LMUL<=1
       // types (VR -> VR). This can be done as a copy.
       if (SubRegIdx == RISCV::NoSubRegister) {
-        unsigned InRegClassID = getRegClassIDForVecVT(InVT);
-        assert(getRegClassIDForVecVT(VT) == RISCV::VRRegClassID &&
+        unsigned InRegClassID =
+            RISCVTargetLowering::getRegClassIDForVecVT(InVT);
+        assert(RISCVTargetLowering::getRegClassIDForVecVT(VT) ==
+                   RISCV::VRRegClassID &&
                InRegClassID == RISCV::VRRegClassID &&
                "Unexpected subvector extraction");
         SDValue RC =
@@ -993,7 +902,7 @@ void RISCVDAGToDAGISel::Select(SDNode *Node) {
       if (Idx != 0)
         break;
 
-      unsigned InRegClassID = getRegClassIDForVecVT(InVT);
+      unsigned InRegClassID = RISCVTargetLowering::getRegClassIDForVecVT(InVT);
 
       SDValue RC =
           CurDAG->getTargetConstant(InRegClassID, DL, Subtarget->getXLenVT());