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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: 陈子昂 (Michael-Chen-NJU)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp (+9) 
- (added) llvm/test/CodeGen/AArch64/issue-171032.ll (+14) 


``````````diff
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
index b9377fabb8634..5f7fd7cb0e8a6 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
@@ -6206,6 +6206,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
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/171619