[llvm] [GlobalISel] Fix G_SHUFFLE_VECTOR in computeKnownFPClass after d012e0380a8f (PR #192385)

[Jinsong Ji via llvm-commits]

https://github.com/jsji created https://github.com/llvm/llvm-project/pull/192385

Fix crash in computeKnownFPClass when processing G_SHUFFLE_VECTOR that
shuffles from a large vector to a smaller one (e.g., <64 x s32> to <8 x s32>).

The bug was introduced in commit d012e0380a8f which added G_SHUFFLE_VECTOR
handling to computeKnownFPClass. The code incorrectly passed the destination
vector's element count (DstTy.getNumElements()) to getShuffleDemandedElts,
but that function expects the source vector's element count.

This caused an assertion failure:
  Assertion `(-1 <= M) && (M < (SrcWidth * 2))' failed
when mask indices were larger than the destination size.

The fix follows the pattern used in computeKnownBits (line 836) which
correctly gets the source type's element count via:
  MRI.getType(MI.getOperand(1).getReg()).getNumElements()

Added a regression test to verify the fix.


>From 694a26018a075506448761342e81844247cce769 Mon Sep 17 00:00:00 2001
From: Jinsong Ji <jinsong.ji at intel.com>
Date: Wed, 15 Apr 2026 20:17:08 -0700
Subject: [PATCH] [GlobalISel] Fix G_SHUFFLE_VECTOR in computeKnownFPClass
 after d012e0380a8f

Fix crash in computeKnownFPClass when processing G_SHUFFLE_VECTOR that
shuffles from a large vector to a smaller one (e.g., <64 x s32> to <8 x s32>).

The bug was introduced in commit d012e0380a8f which added G_SHUFFLE_VECTOR
handling to computeKnownFPClass. The code incorrectly passed the destination
vector's element count (DstTy.getNumElements()) to getShuffleDemandedElts,
but that function expects the source vector's element count.

This caused an assertion failure:
  Assertion `(-1 <= M) && (M < (SrcWidth * 2))' failed
when mask indices were larger than the destination size.

The fix follows the pattern used in computeKnownBits (line 836) which
correctly gets the source type's element count via:
  MRI.getType(MI.getOperand(1).getReg()).getNumElements()

Added a regression test to verify the fix.
---
 .../CodeGen/GlobalISel/GISelValueTracking.cpp |  3 +-
 .../CodeGen/GlobalISel/KnownFPClassTest.cpp   | 30 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp b/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
index 0a58b611c0eca..edf040dc2dd70 100644
--- a/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
@@ -1790,7 +1790,8 @@ void GISelValueTracking::computeKnownFPClass(Register R,
       assert(DemandedElts == APInt(1, 1));
       DemandedLHS = DemandedRHS = DemandedElts;
     } else {
-      if (!llvm::getShuffleDemandedElts(DstTy.getNumElements(), Shuf.getMask(),
+      unsigned SrcNumElts = MRI.getType(Shuf.getSrc1Reg()).getNumElements();
+      if (!llvm::getShuffleDemandedElts(SrcNumElts, Shuf.getMask(),
                                         DemandedElts, DemandedLHS,
                                         DemandedRHS)) {
         Known.resetAll();
diff --git a/llvm/unittests/CodeGen/GlobalISel/KnownFPClassTest.cpp b/llvm/unittests/CodeGen/GlobalISel/KnownFPClassTest.cpp
index b10fd4621b0fd..33dcf0db32f90 100644
--- a/llvm/unittests/CodeGen/GlobalISel/KnownFPClassTest.cpp
+++ b/llvm/unittests/CodeGen/GlobalISel/KnownFPClassTest.cpp
@@ -1643,3 +1643,33 @@ TEST_F(AArch64GISelMITest, TestFPClassFMASelfSquare) {
   EXPECT_EQ(fcNan | fcPosInf | fcPosNormal, Known.KnownFPClasses);
   EXPECT_EQ(std::nullopt, Known.SignBit);
 }
+
+// Regression test for crash when shuffling from large vector to small vector.
+// This was causing an assertion failure in getShuffleDemandedElts because
+// the destination width was incorrectly used instead of source width.
+TEST_F(AArch64GISelMITest, TestFPClassShuffleVecLargeToSmall) {
+  StringRef MIRString = R"(
+    %ptr:_(p0) = G_IMPLICIT_DEF
+    %vec:_(<64 x s32>) = G_LOAD %ptr(p0) :: (load (<64 x s32>))
+    %undef:_(<64 x s32>) = G_IMPLICIT_DEF
+    %shuf:_(<8 x s32>) = G_SHUFFLE_VECTOR %vec(<64 x s32>), %undef, shufflemask(0, 8, 16, 24, 32, 40, 48, 56)
+    %copy_shuf:_(<8 x s32>) = COPY %shuf
+)";
+
+  setUp(MIRString);
+  if (!TM)
+    GTEST_SKIP();
+
+  Register CopyReg = Copies[Copies.size() - 1];
+  MachineInstr *FinalCopy = MRI->getVRegDef(CopyReg);
+  Register SrcReg = FinalCopy->getOperand(1).getReg();
+
+  GISelValueTracking Info(*MF);
+
+  // This should not crash
+  KnownFPClass Known = Info.computeKnownFPClass(SrcReg);
+
+  // We don't know anything about the FP classes since vec is a load
+  EXPECT_EQ(fcAllFlags, Known.KnownFPClasses);
+  EXPECT_EQ(std::nullopt, Known.SignBit);
+}