[RISCV] Improve legalization of e8 m8 VL>256 shuffles

[preames]

If we can't produce a large enough index vector in i8, we may need to legalize
the shuffle (via scalarization - which in turn gets lowered into stack usage).
This change makes two related changes:

• Deferring legalization until we actually need to generate the vrgather
  instruction. With the new recursive structure, this only happens when
  doing the fallback for one of the arms.
• Check the actual mask values for something outside of the representable
  range.

Both are covered by recently added tests.

[llvmbot]

@llvm/pr-subscribers-backend-risc-v

Author: Philip Reames (preames)

Changes

If we can't produce a large enough index vector in i8, we may need to legalize the shuffle (via splitting or scalarization). We were doing this before sub-dividing, but the actual vselect doesn't have this legality property. If our two arms can be generated without resorting to the vrgather path, we were failing for no reason.

This is a functional change, but I didn't include a test (mostly because the existing logic didn't seem to have a test either.)

---

Full diff: https://github.com/llvm/llvm-project/pull/79330.diff

1 Files Affected:

• (modified) llvm/lib/Target/RISCV/RISCVISelLowering.cpp (+8-8)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index a38352e8e87f21..46d75529a2a151 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -4965,14 +4965,6 @@ static SDValue lowerVECTOR_SHUFFLE(SDValue Op, SelectionDAG &DAG,
   if (SDValue V = lowerVECTOR_SHUFFLEAsRotate(SVN, DAG, Subtarget))
     return V;
 
-  if (VT.getScalarSizeInBits() == 8 && VT.getVectorNumElements() > 256) {
-    // On such a large vector we're unable to use i8 as the index type.
-    // FIXME: We could promote the index to i16 and use vrgatherei16, but that
-    // may involve vector splitting if we're already at LMUL=8, or our
-    // user-supplied maximum fixed-length LMUL.
-    return SDValue();
-  }
-
   // As a backup, shuffles can be lowered via a vrgather instruction, possibly
   // merged with a second vrgather.
   SmallVector<int> ShuffleMaskLHS, ShuffleMaskRHS;
@@ -5002,6 +4994,14 @@ static SDValue lowerVECTOR_SHUFFLE(SDValue Op, SelectionDAG &DAG,
   // single source permutation.  Note that all the splat variants
   // are handled above.
   if (V2.isUndef()) {
+    if (VT.getScalarSizeInBits() == 8 && VT.getVectorNumElements() > 256) {
+      // On such a large vector we're unable to use i8 as the index type.
+      // FIXME: We could promote the index to i16 and use vrgatherei16, but that
+      // may involve vector splitting if we're already at LMUL=8, or our
+      // user-supplied maximum fixed-length LMUL.
+      return SDValue();
+    }
+
     unsigned GatherVVOpc = RISCVISD::VRGATHER_VV_VL;
     MVT IndexVT = VT.changeTypeToInteger();
     // Since we can't introduce illegal index types at this stage, use i16 and

[topperc]

The description here is misleading. I believe returning SDValue() here causes the shuffle to get expanded through memory. It will go through the Expand path in LegalizeDAG.

[topperc]

The description here is misleading. I believe returning SDValue() here causes the shuffle to get expanded through memory. It will go through the Expand path in LegalizeDAG.

It doesn't go through memory. I guess I was thinking of build_vector. I'm not sure what shuffle expansion does. I guess it uses extracts and_vector?

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[preames]

I went ahead and added test coverage for both the original code and the new case. In the process, I discovered that one my earlier refactorings actually broke this logic - I think my mental state was assuming the prior version of this patch despite it not having landed - and I landed a separate change to get us back to a working baseline.

Pushed a reworked change, please take a look at the description from scratch.

[lukel97]

LGTM

[topperc]

Why do we need enumerate?