From 02f8519502447de6ef69a85fa8de5732dd59d853 Mon Sep 17 00:00:00 2001
From: David Green <david.green@arm.com>
Date: Mon, 9 May 2022 09:36:22 +0100
Subject: [PATCH] [DAG] Prevent infinite loop combining bitcast shuffle

This prevents an infinite loop from D123801, where code trying to reduce
the total number of bitcasts, but also handling constants, could create
the opposite transform. Prevent the transform in these case to let the
bitcast of a constant transform naturally.

Fixes #55345
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp |  6 +++
 llvm/test/CodeGen/ARM/vector-DAGCombine.ll    |  1 +
 .../CodeGen/X86/shuffle-combine-crash-4.ll    | 42 +++++++++++++++++++
 3 files changed, 49 insertions(+)
 create mode 100644 llvm/test/CodeGen/X86/shuffle-combine-crash-4.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 1b84ad0efcc7..d76e77a7fe13 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -22043,6 +22043,12 @@ static SDValue combineShuffleOfBitcast(ShuffleVectorSDNode *SVN,
       (!Op1.isUndef() && (Op1.getOpcode() != ISD::BITCAST ||
                           Op1.getOperand(0).getValueType() != InVT)))
     return SDValue();
+  if ((ISD::isBuildVectorOfConstantSDNodes(Op0.getOperand(0).getNode()) ||
+       ISD::isBuildVectorOfConstantFPSDNodes(Op0.getOperand(0).getNode())) &&
+      (Op1.isUndef() ||
+       ISD::isBuildVectorOfConstantSDNodes(Op1.getOperand(0).getNode()) ||
+       ISD::isBuildVectorOfConstantFPSDNodes(Op1.getOperand(0).getNode())))
+    return SDValue();
 
   int VTLanes = VT.getVectorNumElements();
   int InLanes = InVT.getVectorNumElements();
diff --git a/llvm/test/CodeGen/ARM/vector-DAGCombine.ll b/llvm/test/CodeGen/ARM/vector-DAGCombine.ll
index 1a9f44bd52e9..f274f331a507 100644
--- a/llvm/test/CodeGen/ARM/vector-DAGCombine.ll
+++ b/llvm/test/CodeGen/ARM/vector-DAGCombine.ll
@@ -56,6 +56,7 @@ define <4 x i32> @test_vmovrrd_combine() nounwind {
 ; CHECK-NEXT:    bne .LBB3_2
 ; CHECK-NEXT:  @ %bb.1: @ %bb1.preheader
 ; CHECK-NEXT:    vmov.i32 q8, #0x0
+; CHECK-NEXT:    vext.8 q8, q8, q8, #4
 ; CHECK-NEXT:  .LBB3_2: @ %bb2
 ; CHECK-NEXT:    vmov r0, r1, d16
 ; CHECK-NEXT:    vmov r2, r3, d17
diff --git a/llvm/test/CodeGen/X86/shuffle-combine-crash-4.ll b/llvm/test/CodeGen/X86/shuffle-combine-crash-4.ll
new file mode 100644
index 000000000000..8ea99b313838
--- /dev/null
+++ b/llvm/test/CodeGen/X86/shuffle-combine-crash-4.ll
@@ -0,0 +1,42 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu -mattr=+avx | FileCheck %s
+
+; Make sure that we do not infinitely-loop combining shuffle vectors.
+
+@test24_id5239 = dso_local local_unnamed_addr global i64 0, align 8
+define void @infiloop() {
+; CHECK-LABEL: infiloop:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushq %rbp
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset %rbp, -16
+; CHECK-NEXT:    movq %rsp, %rbp
+; CHECK-NEXT:    .cfi_def_cfa_register %rbp
+; CHECK-NEXT:    andq $-32, %rsp
+; CHECK-NEXT:    subq $64, %rsp
+; CHECK-NEXT:    movabsq $506097522914230528, %rax # imm = 0x706050403020100
+; CHECK-NEXT:    movq %rax, test24_id5239(%rip)
+; CHECK-NEXT:    vmovaps {{.*#+}} ymm0 = [4,1,6,7,6,7,2,3,6,7,4,5,6,7,2,3,6,7,2,3,2,3,2,3,4,5,4,5,2,3,0,1]
+; CHECK-NEXT:    vmovaps %ymm0, (%rsp)
+; CHECK-NEXT:    movq %rbp, %rsp
+; CHECK-NEXT:    popq %rbp
+; CHECK-NEXT:    .cfi_def_cfa %rsp, 8
+; CHECK-NEXT:    vzeroupper
+; CHECK-NEXT:    retq
+entry:
+  %id5230 = alloca <32 x i8>, align 32
+  store i8 0, ptr @test24_id5239, align 8
+  store i8 1, ptr getelementptr inbounds (i8, ptr @test24_id5239, i64 1), align 1
+  store i8 2, ptr getelementptr inbounds (i8, ptr @test24_id5239, i64 2), align 2
+  store i8 3, ptr getelementptr inbounds (i8, ptr @test24_id5239, i64 3), align 1
+  store i8 4, ptr getelementptr inbounds (i8, ptr @test24_id5239, i64 4), align 4
+  store i8 5, ptr getelementptr inbounds (i8, ptr @test24_id5239, i64 5), align 1
+  store i8 6, ptr getelementptr inbounds (i8, ptr @test24_id5239, i64 6), align 2
+  store i8 7, ptr getelementptr inbounds (i8, ptr @test24_id5239, i64 7), align 1
+  %0 = load <4 x i16>, ptr @test24_id5239, align 8
+  %shuffle = shufflevector <4 x i16> %0, <4 x i16> poison, <16 x i32> <i32 0, i32 3, i32 3, i32 1, i32 3, i32 2, i32 3, i32 1, i32 3, i32 1, i32 1, i32 1, i32 2, i32 2, i32 1, i32 0>
+  %1 = bitcast <16 x i16> %shuffle to <32 x i8>
+  %2 = or <32 x i8> %1, <i8 4, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0>
+  store volatile <32 x i8> %2, ptr %id5230, align 32
+  ret void
+}