[llvm] [RISCV] Legalize shuffle index after splitting two argument shuffles (PR #79330)

[Philip Reames via llvm-commits]

https://github.com/preames created https://github.com/llvm/llvm-project/pull/79330

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.)

>From fdf385c47efaa54a535da545274aeda90391a739 Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Wed, 24 Jan 2024 08:57:54 -0800
Subject: [PATCH] [RISCV] Legalize shuffle index after splitting two argument
 shuffles

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.)
---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index a38352e8e87f21c..46d75529a2a1515 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