[AArch64][SelectionDAG] fix infinite loop caused by legalizing & combining CONCAT_VECTORS

[FLZ101]

Summary

Legalizing in AArch64TargetLowering::LowerCONCAT_VECTORS() and combining in DAGCombiner::visitCONCAT_VECTORS() could cause an infinite loop, as described in https://github.com/llvm/llvm-project/issues/63322.

This commit fixes that issue by conditionally skipping the combining.

[RKSimon]

Please regenerate the diff with context

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23176

We don't usually refer to issue# in code - we add that to the tests instead. The description here should be more generic

23180

The comment doesn't appear to match the code - the code just prevents folding scalable vectors after LegalDAG.

llvm/test/CodeGen/AArch64/aarch64-sve-concat_vectors.ll

2 ↗

(On Diff #532788)

Better to use the update_llc_test_checks script

18 ↗

(On Diff #532788)

remove the unnecessary metadata etc.

[dmgreen]

Can you upload with context? https://llvm.org/docs/Phabricator.html#requesting-a-review-via-the-web-interface. It makes the reviews easier to read.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23171

I'm not sure this comment adds a lot to the existing code. It might make a better summary for the patch.

23182

Does this need to check for isScalableVector, or can it just check for !LegalDAG? Or does that alter other tests?

llvm/test/CodeGen/AArch64/aarch64-sve-concat_vectors.ll

1 ↗

(On Diff #532788)

This likely needs a -mtriple=aarch64-none-eabi, in order to work when the default target is not aarch64.
And is it possible to use -mattr=+sve2 instead of a specific cpu?

4 ↗

(On Diff #532788)

I usually remove ; Function Attrs:..

18 ↗

(On Diff #532788)

Can this be removed?

[FLZ101]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23180

The legalizing happens only when the following conditions are met:

LegalDAG is true

DAGCombiner::Run():

// If this combine is running after legalizing the DAG, re-legalize any
// nodes pulled off the worklist.
if (LegalDAG) {
  SmallSetVector<SDNode *, 16> UpdatedNodes;
  bool NIsValid = DAG.LegalizeOp(N, UpdatedNodes);

  for (SDNode *LN : UpdatedNodes)
    AddToWorklistWithUsers(LN);

  if (!NIsValid)
    continue;
}

input of LowerCONCAT_VECTORS is scalable

AArch64TargetLowering::LowerCONCAT_VECTORS():

assert(Op.getValueType().isScalableVector() &&
       isTypeLegal(Op.getValueType()) &&
       "Expected legal scalable vector type!");

So if we skip the combining here when any of the following conditions is met, the infinite loop will not happen:

• LegalDAG is true
• the combining won't produce a scalable vector

I also want targets that does not support SVE are not affected by this change. However I am not sure whether the code does exactly what I want.

23182

The legalizing happens only when the following conditions are met:

LegalDAG is true

DAGCombiner::Run():

// If this combine is running after legalizing the DAG, re-legalize any
// nodes pulled off the worklist.
if (LegalDAG) {
  SmallSetVector<SDNode *, 16> UpdatedNodes;
  bool NIsValid = DAG.LegalizeOp(N, UpdatedNodes);

  for (SDNode *LN : UpdatedNodes)
    AddToWorklistWithUsers(LN);

  if (!NIsValid)
    continue;
}

input of LowerCONCAT_VECTORS is scalable

AArch64TargetLowering::LowerCONCAT_VECTORS():

assert(Op.getValueType().isScalableVector() &&
       isTypeLegal(Op.getValueType()) &&
       "Expected legal scalable vector type!");

So if we skip the combining here when any of the following conditions is met, the infinite loop will not happen:

• LegalDAG is true
• the combining won't produce a scalable vector

I also want targets that does not support SVE are not affected by this change. However I am not sure whether the code does exactly what I want.

[RKSimon]

llvm/test/CodeGen/AArch64/dag-combine-concat-vectors.ll

2

Add a FileCheck and run update_llc_test_checks to generate codegen to test

[FLZ101]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23180

the combining won't produce a scalable vector

should be

the combining will produce a scalable vector

23182

the combining won't produce a scalable vector

should be

the combining will produce a scalable vector

llvm/test/CodeGen/AArch64/dag-combine-concat-vectors.ll

2

Is there a way to check whether the case could be finished in for example 10s ? I think that would be more appropriate

[RKSimon]

llvm/test/CodeGen/AArch64/dag-combine-concat-vectors.ll

2

We don't tend to bother with timelimit tests as (IIRC) some buildbots struggle to handle them correctly.

[FLZ101]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23180

So if we skip the combining here when any of the following conditions is met

should be

So if we skip the combining here when the following conditions are met

23182

So if we skip the combining here when any of the following conditions is met

should be

So if we skip the combining here when the following conditions are met

[FLZ101]

Add comments

[FLZ101]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23180

Comments updated

23182

Comments updated

[FLZ101]

Fix a format issue

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23167

Drop the "To avoid the infinite loop..." sentence - its just repeating what is explained earlier.

[FLZ101]

Update comments

[RKSimon]

LGTM - cheers

[MaskRay]

llvm/test/CodeGen/AArch64/dag-combine-concat-vectors.ll

2

Run PATH=$build_dir/bin:$PATH llvm-project/llvm/utils/update_llc_test_checks.py dag-combine-concat-vectors.ll to regenerate this test.

https://maskray.me/blog/2021-08-08-toolchain-testing "A regression test which simply runs the compiler and expects it not to crash has low utilization. It is often better to additionally test some properties of the produced object file."

The comment can be simplified:

Test that we do not end in an infinite loop https://github.com/llvm/llvm-project/issues/63322

[FLZ101]

Update a test case

[FLZ101]

In D153316#4448220, @RKSimon wrote:

LGTM - cheers

Thanks

llvm/test/CodeGen/AArch64/dag-combine-concat-vectors.ll

2

Very helpful. Thanks

2

Run PATH=$build_dir/bin:$PATH llvm-project/llvm/utils/update_llc_test_checks.py dag-combine-concat-vectors.ll to regenerate this test.

https://maskray.me/blog/2021-08-08-toolchain-testing "A regression test which simply runs the compiler and expects it not to crash has low utilization. It is often better to additionally test some properties of the produced object file."

The comment can be simplified:

Test that we do not end in an infinite loop https://github.com/llvm/llvm-project/issues/63322

[FLZ101]

I do not have commit access, @MaskRay Could you please commit this patch on my behalf?

Below is the name and email address I would like to use in the Author property of the commit:

FLZ101
[e-mail]

[MaskRay]

In D153316#4451191, @FLZ101 wrote:

I do not have commit access, @MaskRay Could you please commit this patch on my behalf?

Below is the name and email address I would like to use in the Author property of the commit:

FLZ101
[e-mail]

I'll wait a bit in case there is more feedback. Thanks for the patience.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitCONCAT_VECTORS(SDNode *N) {
   // Optimize concat_vectors where all but the first of the vectors are undef.
   if (all_of(drop_begin(N->ops()),
              [](const SDValue &Op) { return Op.isUndef(); })) {
     SDValue In = N->getOperand(0);
     assert(In.getValueType().isVector() && "Must concat vectors");
 
     // If the input is a concat_vectors, just make a larger concat by padding
     // with smaller undefs.
-    if (In.getOpcode() == ISD::CONCAT_VECTORS && In.hasOneUse()) {
+    //
+    // Legalizing in AArch64TargetLowering::LowerCONCAT_VECTORS() and combining
+    // here could cause an infinite loop. That legalizing happens when LegalDAG
+    // is true and input of AArch64TargetLowering::LowerCONCAT_VECTORS() is
+    // scalable.
+    if (In.getOpcode() == ISD::CONCAT_VECTORS && In.hasOneUse() &&
+        !(LegalDAG && In.getValueType().isScalableVector())) {
       unsigned NumOps = N->getNumOperands() * In.getNumOperands();
       SmallVector<SDValue, 4> Ops(In->op_begin(), In->op_end());
       Ops.resize(NumOps, DAG.getUNDEF(Ops[0].getValueType()));
       return DAG.getNode(ISD::CONCAT_VECTORS, SDLoc(N), VT, Ops);
     }
 
     SDValue Scalar = peekThroughOneUseBitcasts(In);
 
     // concat_vectors(scalar_to_vector(scalar), undef) ->
     //     scalar_to_vector(scalar)
     if (!LegalOperations && Scalar.getOpcode() == ISD::SCALAR_TO_VECTOR &&
          Scalar.hasOneUse()) {
       EVT SVT = Scalar.getValueType().getVectorElementType();
       if (SVT == Scalar.getOperand(0).getValueType())
         Scalar = Scalar.getOperand(0);
     }
 
     // concat_vectors(scalar, undef) -> scalar_to_vector(scalar)
     if (!Scalar.getValueType().isVector()) {
       // If the bitcast type isn't legal, it might be a trunc of a legal type;
       // look through the trunc so we can still do the transform:
       //   concat_vectors(trunc(scalar), undef) -> scalar_to_vector(scalar)

llvm/test/CodeGen/AArch64/dag-combine-concat-vectors.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc -mtriple=aarch64 -mattr=+sve < %s | FileCheck %s
+
+; Test that we do not end in an infinite loop (https://github.com/llvm/llvm-project/issues/63322)
+
+declare void @llvm.masked.scatter.nxv16i8.nxv16p0(<vscale x 16 x i8>, <vscale x 16 x ptr>, i32 immarg, <vscale x 16 x i1>)
+
+define fastcc i8 @allocno_reload_assign() {
+; CHECK-LABEL: allocno_reload_assign:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z3.b, #0 // =0x0
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    uunpklo z4.h, z3.b
+; CHECK-NEXT:    uunpkhi z7.h, z3.b
+; CHECK-NEXT:    uunpklo z2.s, z4.h
+; CHECK-NEXT:    uunpkhi z4.s, z4.h
+; CHECK-NEXT:    uunpklo z6.s, z7.h
+; CHECK-NEXT:    uunpkhi z16.s, z7.h
+; CHECK-NEXT:    ptrue p1.b
+; CHECK-NEXT:    mov z0.d, #0 // =0x0
+; CHECK-NEXT:    uunpklo z1.d, z2.s
+; CHECK-NEXT:    uunpkhi z2.d, z2.s
+; CHECK-NEXT:    uunpklo z3.d, z4.s
+; CHECK-NEXT:    uunpkhi z4.d, z4.s
+; CHECK-NEXT:    uunpklo z5.d, z6.s
+; CHECK-NEXT:    uunpkhi z6.d, z6.s
+; CHECK-NEXT:    uunpklo z7.d, z16.s
+; CHECK-NEXT:    uunpkhi z16.d, z16.s
+; CHECK-NEXT:  .LBB0_1: // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    fmov d17, xzr
+; CHECK-NEXT:    cmpeq p2.d, p0/z, z17.d, #0
+; CHECK-NEXT:    uzp1 p2.s, p2.s, p0.s
+; CHECK-NEXT:    uzp1 p2.h, p2.h, p0.h
+; CHECK-NEXT:    uzp1 p2.b, p2.b, p0.b
+; CHECK-NEXT:    mov z17.b, p2/z, #1 // =0x1
+; CHECK-NEXT:    fmov w8, s17
+; CHECK-NEXT:    sbfx x8, x8, #0, #1
+; CHECK-NEXT:    whilelo p2.b, xzr, x8
+; CHECK-NEXT:    not p2.b, p1/z, p2.b
+; CHECK-NEXT:    punpklo p3.h, p2.b
+; CHECK-NEXT:    punpkhi p2.h, p2.b
+; CHECK-NEXT:    punpklo p4.h, p3.b
+; CHECK-NEXT:    punpkhi p3.h, p3.b
+; CHECK-NEXT:    punpklo p5.h, p4.b
+; CHECK-NEXT:    punpkhi p4.h, p4.b
+; CHECK-NEXT:    st1b { z1.d }, p5, [z0.d]
+; CHECK-NEXT:    punpklo p5.h, p2.b
+; CHECK-NEXT:    st1b { z2.d }, p4, [z0.d]
+; CHECK-NEXT:    punpklo p4.h, p3.b
+; CHECK-NEXT:    punpkhi p2.h, p2.b
+; CHECK-NEXT:    punpkhi p3.h, p3.b
+; CHECK-NEXT:    st1b { z3.d }, p4, [z0.d]
+; CHECK-NEXT:    punpklo p4.h, p5.b
+; CHECK-NEXT:    st1b { z4.d }, p3, [z0.d]
+; CHECK-NEXT:    punpkhi p3.h, p5.b
+; CHECK-NEXT:    st1b { z5.d }, p4, [z0.d]
+; CHECK-NEXT:    punpklo p4.h, p2.b
+; CHECK-NEXT:    punpkhi p2.h, p2.b
+; CHECK-NEXT:    st1b { z6.d }, p3, [z0.d]
+; CHECK-NEXT:    st1b { z7.d }, p4, [z0.d]
+; CHECK-NEXT:    st1b { z16.d }, p2, [z0.d]
+; CHECK-NEXT:    b .LBB0_1
+  br label %1
+
+1:                                                ; preds = %1, %0
+  call void @llvm.masked.scatter.nxv16i8.nxv16p0(<vscale x 16 x i8> zeroinitializer, <vscale x 16 x ptr> zeroinitializer, i32 0, <vscale x 16 x i1> xor (<vscale x 16 x i1> shufflevector (<vscale x 16 x i1> icmp eq (<vscale x 16 x ptr> insertelement (<vscale x 16 x ptr> poison, ptr null, i64 0), <vscale x 16 x ptr> zeroinitializer), <vscale x 16 x i1> poison, <vscale x 16 x i32> zeroinitializer), <vscale x 16 x i1> shufflevector (<vscale x 16 x i1> insertelement (<vscale x 16 x i1> poison, i1 true, i32 0), <vscale x 16 x i1> poison, <vscale x 16 x i32> zeroinitializer)))
+  br label %1
+}
+
+uselistorder <vscale x 16 x i1> poison, { 1, 2, 0 }