[DAGCombiner][AArch64] Prevent SimplifyVCastOp from creating illegal scalar types after type legalization. (#148970)

Fixes #148949

Author: topperc

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -28187,14 +28187,16 @@ SDValue DAGCombiner::SimplifyVCastOp(SDNode *N, const SDLoc &DL) {
       TLI.preferScalarizeSplat(N)) {
     EVT SrcVT = N0.getValueType();
     EVT SrcEltVT = SrcVT.getVectorElementType();
-    SDValue IndexC = DAG.getVectorIdxConstant(Index0, DL);
-    SDValue Elt =
-        DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, SrcEltVT, Src0, IndexC);
-    SDValue ScalarBO = DAG.getNode(Opcode, DL, EltVT, Elt, N->getFlags());
-    if (VT.isScalableVector())
-      return DAG.getSplatVector(VT, DL, ScalarBO);
-    SmallVector<SDValue, 8> Ops(VT.getVectorNumElements(), ScalarBO);
-    return DAG.getBuildVector(VT, DL, Ops);
+    if (!LegalTypes || TLI.isTypeLegal(SrcEltVT)) {
+      SDValue IndexC = DAG.getVectorIdxConstant(Index0, DL);
+      SDValue Elt =
+          DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, SrcEltVT, Src0, IndexC);
+      SDValue ScalarBO = DAG.getNode(Opcode, DL, EltVT, Elt, N->getFlags());
+      if (VT.isScalableVector())
+        return DAG.getSplatVector(VT, DL, ScalarBO);
+      SmallVector<SDValue, 8> Ops(VT.getVectorNumElements(), ScalarBO);
+      return DAG.getBuildVector(VT, DL, Ops);
+    }
   }
 
   return SDValue();

llvm/test/CodeGen/AArch64/pr148949.ll

@@ -0,0 +1,22 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=aarch64 -mattr=+sve | FileCheck %s
+
+define <vscale x 4 x i8> @widget(i1 %arg, <vscale x 4 x i1> %arg1, <vscale x 4 x i8> %arg2, <vscale x 4 x i8> %arg3) {
+; CHECK-LABEL: widget:
+; CHECK:       // %bb.0: // %bb
+; CHECK-NEXT:    mvn w8, w0
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo p1.s, xzr, x8
+; CHECK-NEXT:    mov z1.s, p1/z, #1 // =0x1
+; CHECK-NEXT:    orr z0.d, z1.d, z0.d
+; CHECK-NEXT:    mov z0.s, p0/m, z1.s
+; CHECK-NEXT:    ret
+bb:
+  %insertelement = insertelement <vscale x 4 x i1> zeroinitializer, i1 %arg, i64 0
+  %shufflevector = shufflevector <vscale x 4 x i1> %insertelement, <vscale x 4 x i1> zeroinitializer, <vscale x 4 x i32> zeroinitializer
+  %xor = xor <vscale x 4 x i1> %shufflevector, splat (i1 true)
+  %zext = zext <vscale x 4 x i1> %xor to <vscale x 4 x i8>
+  %select = select <vscale x 4 x i1> %arg1, <vscale x 4 x i8> zeroinitializer, <vscale x 4 x i8> %arg2
+  %or = or <vscale x 4 x i8> %select, %zext
+  ret <vscale x 4 x i8> %or
+}