[llvm] [LoongArch] Fix crash on SIGN_EXTEND of non-simple vector types (PR #222250)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-loongarch

Author: wanglei (wangleiat)

<details>
<summary>Changes</summary>

Custom lowering of ISD::SIGN_EXTEND assumed a simple vector source,
which does not hold for e.g. the <4 x i24> mask promoted from an i24
vector compare. Bail out so the type legalizer falls back to the
default path; also guard lowerSIGN_EXTEND_VECTOR_INREG.

Fixes #<!-- -->222088


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


2 Files Affected:

- (modified) llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp (+14-4) 
- (added) llvm/test/CodeGen/LoongArch/lsx/pr222088.ll (+53) 


``````````diff
diff --git a/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp b/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
index 07f6a438bbb1b..9f3f16426e0e4 100644
--- a/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
+++ b/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
@@ -1021,6 +1021,11 @@ SDValue LoongArchTargetLowering::lowerSIGN_EXTEND_VECTOR_INREG(
     SDValue Op, SelectionDAG &DAG) const {
   SDLoc DL(Op);
   SDValue Src = Op.getOperand(0);
+
+  // Only whole 128-bit vectors of simple element types can be lowered this
+  // way; fall back to the default lowering for other sources.
+  if (!Src.getValueType().isSimple())
+    return SDValue();
   MVT SrcVT = Src.getSimpleValueType();
   MVT DstVT = Op.getSimpleValueType();
 
@@ -6003,15 +6008,20 @@ void LoongArchTargetLowering::ReplaceNodeResults(
 
     EVT DstVT = N->getValueType(0);
     SDValue Src = N->getOperand(0);
-    MVT SrcVT = Src.getSimpleValueType();
+
+    // The VSLTI + VILVL/VILVH expansion below only supports simple vector
+    // sources no wider than 128 bits. Fall back to the default legalization
+    // for anything else, e.g. when the mask of an i24 vector compare is
+    // sign-extended, the promoted <4 x i24> source is not a simple type.
+    EVT SrcEVT = Src.getValueType();
+    if (!SrcEVT.isSimple() || SrcEVT.getSizeInBits() > 128)
+      return;
+    MVT SrcVT = SrcEVT.getSimpleVT();
 
     unsigned SrcEltBits = SrcVT.getScalarSizeInBits();
     unsigned DstEltBits = DstVT.getScalarSizeInBits();
     unsigned NumElts = DstVT.getVectorNumElements();
 
-    if (SrcVT.getSizeInBits() > 128)
-      return;
-
     if (!DstVT.isVector() || DstVT.getSizeInBits() <= 128)
       return;
 
diff --git a/llvm/test/CodeGen/LoongArch/lsx/pr222088.ll b/llvm/test/CodeGen/LoongArch/lsx/pr222088.ll
new file mode 100644
index 0000000000000..020e5e7793c75
--- /dev/null
+++ b/llvm/test/CodeGen/LoongArch/lsx/pr222088.ll
@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc --mtriple=loongarch64 --mattr=+lsx %s -o - | FileCheck %s
+
+define void @sext_v4i24_mask(ptr %src, ptr %dst) {
+; CHECK-LABEL: sext_v4i24_mask:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    ld.b $a2, $a0, 5
+; CHECK-NEXT:    ld.hu $a3, $a0, 3
+; CHECK-NEXT:    ld.b $a4, $a0, 2
+; CHECK-NEXT:    ld.hu $a5, $a0, 0
+; CHECK-NEXT:    slli.d $a2, $a2, 16
+; CHECK-NEXT:    or $a2, $a3, $a2
+; CHECK-NEXT:    slli.d $a3, $a4, 16
+; CHECK-NEXT:    or $a3, $a5, $a3
+; CHECK-NEXT:    ld.b $a4, $a0, 8
+; CHECK-NEXT:    vinsgr2vr.w $vr0, $a3, 0
+; CHECK-NEXT:    vinsgr2vr.w $vr0, $a2, 1
+; CHECK-NEXT:    ld.hu $a2, $a0, 6
+; CHECK-NEXT:    slli.d $a3, $a4, 16
+; CHECK-NEXT:    ld.b $a4, $a0, 11
+; CHECK-NEXT:    ld.hu $a0, $a0, 9
+; CHECK-NEXT:    or $a2, $a2, $a3
+; CHECK-NEXT:    vinsgr2vr.w $vr0, $a2, 2
+; CHECK-NEXT:    slli.d $a2, $a4, 16
+; CHECK-NEXT:    or $a0, $a0, $a2
+; CHECK-NEXT:    vinsgr2vr.w $vr0, $a0, 3
+; CHECK-NEXT:    vrepli.b $vr1, 0
+; CHECK-NEXT:    vslt.w $vr0, $vr1, $vr0
+; CHECK-NEXT:    vpickve2gr.w $a0, $vr0, 0
+; CHECK-NEXT:    vinsgr2vr.d $vr2, $a0, 0
+; CHECK-NEXT:    vpickve2gr.w $a0, $vr0, 1
+; CHECK-NEXT:    vinsgr2vr.d $vr2, $a0, 1
+; CHECK-NEXT:    vslli.d $vr2, $vr2, 40
+; CHECK-NEXT:    vsrai.d $vr2, $vr2, 40
+; CHECK-NEXT:    vpickve2gr.w $a0, $vr0, 2
+; CHECK-NEXT:    vinsgr2vr.d $vr3, $a0, 0
+; CHECK-NEXT:    vpickve2gr.w $a0, $vr0, 3
+; CHECK-NEXT:    vinsgr2vr.d $vr3, $a0, 1
+; CHECK-NEXT:    vslli.d $vr0, $vr3, 40
+; CHECK-NEXT:    vsrai.d $vr0, $vr0, 40
+; CHECK-NEXT:    vrepli.d $vr3, 5
+; CHECK-NEXT:    vbitsel.v $vr0, $vr1, $vr3, $vr0
+; CHECK-NEXT:    vbitsel.v $vr1, $vr1, $vr3, $vr2
+; CHECK-NEXT:    vst $vr1, $a1, 0
+; CHECK-NEXT:    vst $vr0, $a1, 16
+; CHECK-NEXT:    ret
+entry:
+  %1 = load <4 x i24>, ptr %src
+  %2 = icmp sgt <4 x i24> %1, zeroinitializer
+  %3 = select <4 x i1> %2, <4 x i64> splat (i64 5), <4 x i64> zeroinitializer
+  store <4 x i64> %3, ptr %dst
+  ret void
+}

``````````

</details>


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