[llvm] [AArch64] Fix assertion failure during promotion of EXTEND_VECTOR_INREG. (PR #171619)

[via llvm-commits]

https://github.com/Michael-Chen-NJU created https://github.com/llvm/llvm-project/pull/171619

This fixes an assertion failure in `SelectionDAG::getNode` on AArch64 during Type Legalization.

The crash is triggered by an `ANY_EXTEND_VECTOR_INREG` operation involving small vector types (e.g., v16i1).

The crash occurs when the Type Legalizer processes a vector extend operation where the input vector uses small elements, specifically in the case of a ShuffleVector that generates a mask vector.

1.  **Original Node**: `any_extend_vector_inreg (v16i1) -> v2i16`. (This is physically valid: 16 bits < 32 bits).
2.  **Promotion Issue**: When both the input and result types are promoted for legality:
    * The **Result** (`v2i16`) is promoted to a larger legal type, e.g., `v2i32` (**64 bits**).
    * The **Input** (`v16i1`) is promoted to `v16i8` (**128 bits**) due to the necessary scalar promotion of `i1` to `i8`.
3.  The legalizer then attempts to create the new node: `any_extend_vector_inreg (v16i8) -> v2i32`.
4.  Since $128 \text{ bits} > 64 \text{ bits}$, the physical constraint of the `EXTEND_VECTOR_INREG` operation is violated, causing the assertion to fail.

### Solution

In `DAGTypeLegalizer::PromoteIntRes_EXTEND_VECTOR_INREG`, when the size of the promoted input vector (`Promoted`) is found to be greater than the size of the promoted result vector (`NVT`):

We explicitly truncates the promoted input to the size of the result type (`NVT`), ensuring the final `*_EXTEND_VECTOR_INREG` node satisfies the size constraint before it is created.

This behavior aligns with the fact that `*_EXTEND_VECTOR_INREG` typically only requires the low-order lanes of the input vector.

**Test Added**: `llvm/test/CodeGen/AArch64/issue-171032.ll`

Fixes: #171032

>From 0fa0ed67b4410db36e5ba9225219536861b9b1a7 Mon Sep 17 00:00:00 2001
From: Michael-Chen-NJU <2802328816 at qq.com>
Date: Wed, 10 Dec 2025 21:30:50 +0800
Subject: [PATCH] [AArch64] Fix assertion failure during promotion of
 EXTEND_VECTOR_INREG.

---
 .../CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp  |  9 +++++++++
 llvm/test/CodeGen/AArch64/issue-171032.ll          | 14 ++++++++++++++
 2 files changed, 23 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/issue-171032.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
index 08c6b90cc8a74..d97d1181cd4c0 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
@@ -6216,6 +6216,15 @@ SDValue DAGTypeLegalizer::PromoteIntRes_EXTEND_VECTOR_INREG(SDNode *N) {
       default:
         llvm_unreachable("Node has unexpected Opcode");
     }
+    if (Promoted.getValueType().getSizeInBits() > NVT.getSizeInBits()) {
+      unsigned NewSize = NVT.getSizeInBits();
+      EVT ExtractVT = EVT::getVectorVT(
+          *DAG.getContext(), Promoted.getValueType().getScalarType(),
+          NewSize / Promoted.getScalarValueSizeInBits());
+
+      Promoted = DAG.getNode(ISD::EXTRACT_SUBVECTOR, dl, ExtractVT, Promoted,
+                             DAG.getVectorIdxConstant(0, dl));
+    }
     return DAG.getNode(N->getOpcode(), dl, NVT, Promoted);
   }
 
diff --git a/llvm/test/CodeGen/AArch64/issue-171032.ll b/llvm/test/CodeGen/AArch64/issue-171032.ll
new file mode 100644
index 0000000000000..fe3c65f24f8d3
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/issue-171032.ll
@@ -0,0 +1,14 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-unknown-linux-musl < %s | FileCheck %s
+
+define <64 x i1> @func_20(<2 x i1> %0) {
+; CHECK-LABEL: func_20:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    // kill: def $d0 killed $d0 def $q0
+; CHECK-NEXT:    mov v0.b[1], v0.b[4]
+; CHECK-NEXT:    str d0, [x8]
+; CHECK-NEXT:    ret
+entry:
+  %.splat = shufflevector <2 x i1> %0, <2 x i1> zeroinitializer, <64 x i32> <i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 1, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison>
+  ret <64 x i1> %.splat
+}