[SelectionDAG] Don't create illegally-typed nodes while constant folding

[frasercrmck]

Summary

This patch fixes a (seemingly very rare) crash during vector constant
folding introduced in D113300.

Normally, during legalization, if we create an illegally-typed node during
a failed attempt at constant folding it's cleaned up before being
visited, due to it having no uses.

If, however, an illegally-typed node is created during one round of
legalization and isn't cleaned up, it's possible for a second round of
legalization to create new illegally-typed nodes which add extra uses to
the old illegal nodes. This means that we can end up visiting the old
nodes before they're known to be dead, at which point we crash.

I'm not happy about this fix. Creating illegal types at all seems like a
bad idea, but we all-too-often rely on illegal constants being
successfully folded and being fixed up afterwards. However, we can't
rely on constant folding actually happening, and we don't have a
foolproof way of peering into the future.

Perhaps the correct fix is to revisit the node-iteration order during
legalization, ensuring we visit all uses of nodes before the nodes
themselves. Or alternatively we could try and clean up dead nodes
immediately after failing constant folding.

[craig.topper]

Did this start failing because we no longer pre-check that all the build vector operands are constants like we did in llvm 13?

FoldConstantVectorArithmetic had this code

// All operands must be vector types with the same number of elements as
// the result type and must be either UNDEF or a build vector of constant
// or UNDEF scalars.
if (!llvm::all_of(Ops, IsConstantBuildVectorSplatVectorOrUndef) ||
    !llvm::all_of(Ops, IsScalarOrSameVectorSize))
  return SDValue();

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5617

legalizatoin -> legalization

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5619

Can we use this instead of checking ISD::Constant

!isa<ConstantSDNode>(ScalarOp) && !ScalarOp.isUndef()

That allows undef(fixed the X86 test) and TargetConstant(fixes the mips changes). Does suggest that mips might be using TargetConstant in an odd way.

[frasercrmck]

apply Craig's suggestion

[frasercrmck]

In D122382#3406064, @craig.topper wrote:

Did this start failing because we no longer pre-check that all the build vector operands are constants like we did in llvm 13?

FoldConstantVectorArithmetic had this code

// All operands must be vector types with the same number of elements as
// the result type and must be either UNDEF or a build vector of constant
// or UNDEF scalars.
if (!llvm::all_of(Ops, IsConstantBuildVectorSplatVectorOrUndef) ||
    !llvm::all_of(Ops, IsScalarOrSameVectorSize))
  return SDValue();

Yeah exactly, if I had to guess it'd be D113300 which introduced that change and was up-front about that change in behaviour. Hard to predict this kind of thing falling out of it, though.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5617

Ah, thanks!

5619

Huh, turns out we can, thanks! I regret I didn't dive deeper into those changes; I was already feeling pretty sheepish about this change. I guess we can rely on constant folding succeeding on all nodes when given either undefs and/or constants? Dunno, just feels pretty flaky to me.

[RKSimon]

LGTM - cheers

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

for (SDValue Op : Ops) {
       }
 
       SDValue ScalarOp =
           Op.getOperand(Op.getOpcode() == ISD::SPLAT_VECTOR ? 0 : I);
       EVT ScalarVT = ScalarOp.getValueType();
 
       // Build vector (integer) scalar operands may need implicit
       // truncation - do this before constant folding.
-      if (ScalarVT.isInteger() && ScalarVT.bitsGT(InSVT))
+      if (ScalarVT.isInteger() && ScalarVT.bitsGT(InSVT)) {
+        // Don't create illegally-typed nodes unless they're constants or undef
+        // - if we fail to constant fold we can't guarantee the (dead) nodes
+        // we're creating will be cleaned up before being visited for
+        // legalization.
+        if (NewNodesMustHaveLegalTypes && !ScalarOp.isUndef() &&
+            !isa<ConstantSDNode>(ScalarOp) &&
+            TLI->getTypeAction(*getContext(), InSVT) !=
+                TargetLowering::TypeLegal)
+          return SDValue();
         ScalarOp = getNode(ISD::TRUNCATE, DL, InSVT, ScalarOp);
+      }
 
       ScalarOps.push_back(ScalarOp);
     }
 
     // Constant fold the scalar operands.
     SDValue ScalarResult = getNode(Opcode, DL, SVT, ScalarOps);
 
     // Legalize the (integer) scalar constant if necessary.

llvm/test/CodeGen/RISCV/rvv/constant-folding-crash.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=riscv32 -mattr=+v -riscv-v-vector-bits-min=128 -verify-machineinstrs < %s \
+; RUN:   | FileCheck %s --check-prefix RV32
+; RUN: llc -mtriple=riscv64 -mattr=+v -riscv-v-vector-bits-min=128 -verify-machineinstrs < %s \
+; RUN:   | FileCheck %s --check-prefix RV64
+
+; This used to crash during type legalization, where lowering (v4i1 =
+; BUILD_VECTOR) created a (v4i1 = SETCC v4i8) which during constant-folding
+; created illegally-typed i8 nodes. Ultimately, constant-folding failed and so
+; the new illegal nodes had no uses. However, during a second round of
+; legalization, this same pattern was generated from another BUILD_VECTOR. This
+; meant one of the illegally-typed (i8 = Constant<0>) nodes now had two dead
+; uses. Because the Constant and one of the uses were from round 1, they were
+; further up in the node order than the new second use, so the constant was
+; visited while it wasn't "dead". At the point of visiting the constant, we
+; crashed.
+
+define void @constant_folding_crash(i8* %v54, <4 x <4 x i32>*> %lanes.a, <4 x <4 x i32>*> %lanes.b, <4 x i1> %sel) {
+; RV32-LABEL: constant_folding_crash:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    lw a0, 8(a0)
+; RV32-NEXT:    vmv1r.v v10, v0
+; RV32-NEXT:    andi a0, a0, 1
+; RV32-NEXT:    seqz a0, a0
+; RV32-NEXT:    vsetivli zero, 4, e8, mf4, ta, mu
+; RV32-NEXT:    vmv.v.x v11, a0
+; RV32-NEXT:    vmsne.vi v0, v11, 0
+; RV32-NEXT:    vsetvli zero, zero, e32, m1, ta, mu
+; RV32-NEXT:    vmerge.vvm v8, v9, v8, v0
+; RV32-NEXT:    vsetvli zero, zero, e8, mf4, ta, mu
+; RV32-NEXT:    vmv.v.i v9, 0
+; RV32-NEXT:    vsetvli zero, zero, e32, m1, ta, mu
+; RV32-NEXT:    vmv.x.s a0, v8
+; RV32-NEXT:    vsetivli zero, 4, e8, mf4, ta, mu
+; RV32-NEXT:    vmv1r.v v0, v10
+; RV32-NEXT:    vmerge.vim v8, v9, 1, v0
+; RV32-NEXT:    vmv.x.s a1, v8
+; RV32-NEXT:    andi a1, a1, 1
+; RV32-NEXT:    vmv.v.x v8, a1
+; RV32-NEXT:    vmsne.vi v0, v8, 0
+; RV32-NEXT:    vsetvli zero, zero, e32, m1, ta, mu
+; RV32-NEXT:    vmv.v.i v8, 10
+; RV32-NEXT:    vse32.v v8, (a0), v0.t
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: constant_folding_crash:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    ld a0, 8(a0)
+; RV64-NEXT:    vmv1r.v v12, v0
+; RV64-NEXT:    andi a0, a0, 1
+; RV64-NEXT:    seqz a0, a0
+; RV64-NEXT:    vsetivli zero, 4, e8, mf4, ta, mu
+; RV64-NEXT:    vmv.v.x v13, a0
+; RV64-NEXT:    vmsne.vi v0, v13, 0
+; RV64-NEXT:    vsetvli zero, zero, e64, m2, ta, mu
+; RV64-NEXT:    vmerge.vvm v8, v10, v8, v0
+; RV64-NEXT:    vsetvli zero, zero, e8, mf4, ta, mu
+; RV64-NEXT:    vmv.v.i v10, 0
+; RV64-NEXT:    vsetvli zero, zero, e64, m2, ta, mu
+; RV64-NEXT:    vmv.x.s a0, v8
+; RV64-NEXT:    vsetivli zero, 4, e8, mf4, ta, mu
+; RV64-NEXT:    vmv1r.v v0, v12
+; RV64-NEXT:    vmerge.vim v8, v10, 1, v0
+; RV64-NEXT:    vmv.x.s a1, v8
+; RV64-NEXT:    andi a1, a1, 1
+; RV64-NEXT:    vmv.v.x v8, a1
+; RV64-NEXT:    vmsne.vi v0, v8, 0
+; RV64-NEXT:    vsetvli zero, zero, e32, m1, ta, mu
+; RV64-NEXT:    vmv.v.i v8, 10
+; RV64-NEXT:    vse32.v v8, (a0), v0.t
+; RV64-NEXT:    ret
+entry:
+  %sunkaddr = getelementptr i8, i8* %v54, i64 8
+  %v55 = bitcast i8* %sunkaddr to i64*
+  %v56 = load i64, i64* %v55, align 8
+  %trunc = and i64 %v56, 1
+  %cmp = icmp eq i64 %trunc, 0
+  %ptrs = select i1 %cmp, <4 x <4 x i32>*> %lanes.a, <4 x <4 x i32>*> %lanes.b
+  %v67 = extractelement <4 x <4 x i32>*> %ptrs, i64 0
+  %mask = shufflevector <4 x i1> %sel, <4 x i1> undef, <4 x i32> zeroinitializer
+  call void @llvm.masked.store.v4i32.p0v4i32(<4 x i32> <i32 10, i32 10, i32 10, i32 10>, <4 x i32>* %v67, i32 16, <4 x i1> %mask)
+  ret void
+}
+
+declare void @llvm.masked.store.v4i32.p0v4i32(<4 x i32>, <4 x i32>*, i32, <4 x i1>)