[llvm] [SelectionDAG] Use correct result type in visitEXTRACT_VECTOR_ELT (PR #148707)

[via llvm-commits]

https://github.com/dlee992 updated https://github.com/llvm/llvm-project/pull/148707

>From cb71f48f5f380a928bac4d99e60eabc4436f1762 Mon Sep 17 00:00:00 2001
From: dlee992 <lidanuaa at gmail.com>
Date: Mon, 14 Jul 2025 14:04:49 -0500
Subject: [PATCH 1/2] use ScalarVT in visitEXTRACT_VECTOR_ELT

---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 8136f1794775e..f26ee0ace6919 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -23440,10 +23440,10 @@ SDValue DAGCombiner::visitEXTRACT_VECTOR_ELT(SDNode *N) {
 
   // If Idx was -1 above, Elt is going to be -1, so just return undef.
   if (Elt == -1)
-    return DAG.getUNDEF(LVT);
+    return DAG.getUNDEF(ScalarVT);
 
-  if (SDValue Scalarized =
-          TLI.scalarizeExtractedVectorLoad(LVT, DL, VecVT, Index, LN0, DAG)) {
+  if (SDValue Scalarized = TLI.scalarizeExtractedVectorLoad(ScalarVT, DL, VecVT,
+                                                            Index, LN0, DAG)) {
     ++OpsNarrowed;
     return Scalarized;
   }

>From bb5656fc709bbc7eaf01877d0e235a3bd590b178 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Da=20Li=20=28=E6=9D=8E=E8=BE=BE=29?= <lidanuaa at gmail.com>
Date: Sat, 26 Sep 2026 23:52:21 -0500
Subject: [PATCH 2/2] [SelectionDAG] Add a unit test for a wide extract of a
 vector load

In-tree targets only create an EXTRACT_VECTOR_ELT wider than its vector
element when the element type is illegal, while
scalarizeExtractedVectorLoad requires a load of the element type to be
legal, so the fixed combine cannot be reached from IR. Build

  i64 = extract_vector_elt (v4i32 load), 1

on AArch64 directly, run DAGCombiner after legalization and check that
it becomes an i64 extending load of the i32 element at offset 4.
Without the fix the combiner fails the "Type mismatch" assertion.

Co-authored-by: Copilot <223556219+Copilot at users.noreply.github.com>
Copilot-Session: 968e3844-e855-4fba-824c-915937281876
---
 .../AArch64/AArch64SelectionDAGTest.cpp       | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
index c74d15782398a..2e2c472edcdd3 100644
--- a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
+++ b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
@@ -9,6 +9,7 @@
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/AsmParser/Parser.h"
+#include "llvm/CodeGen/FunctionLoweringInfo.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/CodeGen/TargetLowering.h"
@@ -1040,6 +1041,34 @@ TEST_F(AArch64SelectionDAGTest, computeKnownBits_extload_knownnegative) {
   EXPECT_EQ(Known.One, APInt(32, 0xfffffff0));
 }
 
+// The result of an EXTRACT_VECTOR_ELT may be wider than the vector element
+// type. Scalarizing an extract of a loaded vector after legalization must then
+// create an extending load of the result type.
+TEST_F(AArch64SelectionDAGTest, Combine_EXTRACT_VECTOR_ELT_ExtLoad) {
+  // DAGCombiner asks the current block whether to optimize for size.
+  FunctionLoweringInfo FLI;
+  FLI.MBB = MF->CreateMachineBasicBlock(&F->getEntryBlock());
+  MF->push_back(FLI.MBB);
+  DAG->setFunctionLoweringInfo(&FLI);
+
+  SDLoc Loc;
+  SDValue Ptr = DAG->getCopyFromReg(DAG->getEntryNode(), Loc,
+                                    Register::index2VirtReg(0), MVT::i64);
+  SDValue Load = DAG->getLoad(MVT::v4i32, Loc, DAG->getEntryNode(), Ptr,
+                              MachinePointerInfo(), Align(16));
+  HandleSDNode Extract(DAG->getNode(ISD::EXTRACT_VECTOR_ELT, Loc, MVT::i64,
+                                    Load, DAG->getVectorIdxConstant(1, Loc)));
+
+  DAG->Combine(AfterLegalizeDAG, /*BatchAA=*/nullptr, CodeGenOptLevel::Default);
+
+  auto *Scalar = dyn_cast<LoadSDNode>(Extract.getValue());
+  ASSERT_NE(Scalar, nullptr);
+  EXPECT_EQ(Scalar->getValueType(0), MVT::i64);
+  EXPECT_NE(Scalar->getExtensionType(), ISD::NON_EXTLOAD);
+  EXPECT_EQ(Scalar->getMemoryVT(), MVT::i32);
+  EXPECT_EQ(Scalar->getPointerInfo().Offset, 4);
+}
+
 TEST_F(AArch64SelectionDAGTest,
        computeKnownBits_AVGFLOORU_AVGFLOORS_AVGCEILU_AVGCEILS) {
   SDLoc Loc;