[SelectionDAG] Do not build illegal nodes with users

[ErikHogeman]

When we build a node with illegal type which has a user, it's possible that it can end up being processed by the DAG combiner later before it's removed, which can trigger an assert expecting the types to be legalized already.

[llvmbot]

@llvm/pr-subscribers-llvm-selectiondag

Author: None (ErikHogeman)

Changes

When we build a node with illegal type which has a user, it's possible that it can end up being processed by the DAG combiner later before it's removed, which can trigger an assert expecting the types to be legalized already.

---

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

1 Files Affected:

• (modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp (+7-4)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 29505f444b7650..44ec6f7cab145a 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -6799,14 +6799,17 @@ SDValue SelectionDAG::FoldConstantArithmetic(unsigned Opcode, const SDLoc &DL,
     // Constant fold the scalar operands.
     SDValue ScalarResult = getNode(Opcode, DL, SVT, ScalarOps, Flags);
 
-    // Legalize the (integer) scalar constant if necessary.
-    if (LegalSVT != SVT)
-      ScalarResult = getNode(ExtendCode, DL, LegalSVT, ScalarResult);
-
     // Scalar folding only succeeded if the result is a constant or UNDEF.
     if (!ScalarResult.isUndef() && ScalarResult.getOpcode() != ISD::Constant &&
         ScalarResult.getOpcode() != ISD::ConstantFP)
       return SDValue();
+
+    // Legalize the (integer) scalar constant if necessary. We only do
+    // this once we know the folding succeeded, since otherwise we would
+    // get a node with illegal type which has a user.
+    if (LegalSVT != SVT)
+      ScalarResult = getNode(ExtendCode, DL, LegalSVT, ScalarResult);
+
     ScalarResults.push_back(ScalarResult);
   }
 

[ErikHogeman]

Hi, this code is triggering some asserts in our downstream target. I have tried to reproduce it with a test that can run on an upstream target, but it has been very difficult.

What happens is basically this:

1. The code in FoldConstantArithmetic creates a new SETCC node with an illegal i1 type.
2. It tries to legalize this by inserting an integer extension, I believe the intention being here that this will be constant folded whenever we keep this node.
3. The result is not constant, so the function returns, leaving the node with illegal type still existing in the DAG.
4. Another node is processed by the DAG combiner, and this node gets simplified to another node which is an operand to the new illegal node. This causes the illegal node to be added to the DAG combiner worklist.
5. When the illegal node gets processed, it's not removed because it has a user, which is the integer extension that was added to legalize it.
6. This triggers an assert in SelectionDAGLegalize::LegalizeOp, which expects types to be legalized (because this is running after the type legalizer).
7. This fix avoids the issue by not adding any user of the illegal node in case it wasn't folded to a constant, which makes sure it gets removed in case it somehow ends up again in the DAG combiner worklist.

The tricky part here is to get another node optimized after the i1 node is created, to trigger it to get inserted into the worklist again before it gets removed later. I have some IR where this happens on our downstream target, but I haven't managed to create IR that makes this happen on an upstream target.

So at this point, I only have a code change, no test. I realize that is not ideal, but I still thought I should open this request to see what opinion people have about it. I'm also happy to give it another try to write a test, if someone has a good idea of anything I can try.

[ErikHogeman]

I managed to simplify our case down to this:

define i16 @func(i64 %arg, i64 %arg1, <2 x i16>* %_ptr, i16* %_ptr2) {
entry:
  %i8 = load i16, i16* %_ptr2
  %i9 = sub i16 0, %i8
  %i10 = insertelement <2 x i16> undef, i16 %i9, i64 0
  %i11 = insertelement <2 x i16> %i10, i16 %i9, i64 1
  %i12 = sub <2 x i16> zeroinitializer, %i11
  %i13 = extractelement <2 x i16> %i12, i64 0
  %i14 = add i16 %i13, %i8
  %i15 = insertelement <2 x i16> undef, i16 %i14, i64 0
  %i16 = insertelement <2 x i16> undef, i16 %i8, i64 0
  %i17 = extractelement <2 x i16> %i12, i64 1
  %i18 = insertelement <2 x i16> %i16, i16 %i17, i64 1
  %i19 = insertelement <2 x i16> %i16, i16 %i8, i64 1
  %i20 = add <2 x i16> %i18, %i19
  %i36 = icmp ne <2 x i16> %i20, %i12
  %i37 = extractelement <2 x i1> %i36, i64 0
  %i38 = extractelement <2 x i1> %i36, i64 1
  %i39 = or i1 %i37, %i38
  %i40 = icmp ne i16 %i13, 0
  %i41 = or i1 %i39, %i40
  %i43 = load i16, i16* %_ptr2
  %i44 = sub i16 0, %i43
  %i45 = select i1 %i41, i16 2, i16 %i44
  ret i16 %i45
}

Running selection on this for x86_64 target, it will create the illegal i1 node. However, it happens during DAG combining just after type legalization, but the DAG combiner will only run LegalizeOp if it's AfterLegalizeDAG, and on top of that the illegal node is not added to the worklist again anyway since I didn't get other optimizations to trigger in such a way that this happens, so this prevents me from reproducing it. It's tricky since there's no way (at least I'm aware of?) to write selection DAG IR manually and run specific passes or transforms on it.

If there's any feedback about what I could try with the IR for an upstream target, or whether or not the change itself would be OK without a test, please let me know.

[ErikHogeman]

@bogner , I believe you are listed as the code owner of the SelectionDAG part of LLVM. If/when you have time, could you help me add the relevant people to review this change?

[RKSimon]

@ErikHogeman That IR is building OK for me - what are your command args? Can you repro it in godbolt?

[ErikHogeman]

@ErikHogeman That IR is building OK for me - what are your command args? Can you repro it in godbolt?

Hi, apologies if I wasn't clear in the comments, I have not been able to reproduce it on an upstream target. This IR was the closest I got since it builds the invalid node, but as I wrote it happens too early (needs to happen after LegalizeDAG), plus another optimization needs to insert the illegal node back to the DAG combiner worklist before it gets removed later. I tried to explain what happens in my first comment, so I have been trying to write something that triggers those same steps in an upstream target.

EDIT:

These nodes are built after type legalization, inside 'FoldConstantArithmetic':

Creating new node: t96: i1 = setcc t11, Constant:i16<0>, setne:ch
Creating new node: t97: i16 = sign_extend t96

In our downstream target, t96 is then further inserted into the DAG combiner worklist, and when it runs 'LegalizeOp' this asserts because the type is illegal.

[RKSimon]

LGTM by inspection

[ErikHogeman]

LGTM by inspection

Thanks!

[AZero13]

Should this be backported to 19.x?

[RKSimon]

Given that the only report is for a out of tree target (that can maintain their own minimal local changes in a branch) I don't think this is necessary.

[ErikHogeman]

Given that the only report is for a out of tree target (that can maintain their own minimal local changes in a branch) I don't think this is necessary.

Agreed.