[AArch64] Fix regression from “Fold scalar-to-vector shuffles into DUP/FMOV" (#178227)

Revised #166962.

This patch aims to fix the original compile time regression by
restricting the optimisation to run only on non-constant splats. Without
the guard, an infinite loop is caused because the
`CONCAT(SCALAR_TO_VECTOR, zero)` folds back into the same `BUILD_VECTOR`
and immediately re-enters `LowerBUILD_VECTOR`.

This patch was tested with the original TensorFlow reproduction provided
on the PR and shows a (very) slight improvement on compile-time.
This commit is contained in:
Amina Chabane 2026-02-25 19:07:35 +00:00 committed by GitHub
parent 6d5e051be9
commit 5734d978d0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 157 additions and 25 deletions

View File

@ -16167,12 +16167,20 @@ SDValue AArch64TargetLowering::LowerBUILD_VECTOR(SDValue Op,
SmallMapVector<SDValue, unsigned, 16> DifferentValueMap;
unsigned ConsecutiveValCount = 0;
SDValue PrevVal;
auto IsZero = [&](SDValue V) {
return isNullConstant(V) || isNullFPConstant(V);
};
bool MaybeLowHalfZeroHigh =
VT.isFixedLengthVector() && VT.getSizeInBits() == 128;
unsigned HalfElts = MaybeLowHalfZeroHigh ? (NumElts >> 1) : 0;
SDValue LowHalfFirstVal = MaybeLowHalfZeroHigh ? Op.getOperand(0) : SDValue();
for (unsigned i = 0; i < NumElts; ++i) {
SDValue V = Op.getOperand(i);
if (V.getOpcode() != ISD::EXTRACT_VECTOR_ELT)
AllLanesExtractElt = false;
if (V.isUndef()) {
++NumUndefLanes;
MaybeLowHalfZeroHigh = false;
continue;
}
if (i > 0)
@ -16199,6 +16207,14 @@ SDValue AArch64TargetLowering::LowerBUILD_VECTOR(SDValue Op,
ConsecutiveValCount = 0;
PrevVal = V;
}
if (MaybeLowHalfZeroHigh) {
if (i < HalfElts) {
if (V != LowHalfFirstVal)
MaybeLowHalfZeroHigh = false;
} else if (!IsZero(V)) {
MaybeLowHalfZeroHigh = false;
}
}
// Keep different values and its last consecutive count. For example,
//
@ -16229,6 +16245,22 @@ SDValue AArch64TargetLowering::LowerBUILD_VECTOR(SDValue Op,
return DAG.getNode(ISD::SCALAR_TO_VECTOR, DL, VT, Value);
}
if (MaybeLowHalfZeroHigh && LowHalfFirstVal.getNode() &&
!LowHalfFirstVal.isUndef() && !isIntOrFPConstant(LowHalfFirstVal)) {
EVT LaneVT = VT.getVectorElementType();
EVT HalfVT = VT.getHalfNumVectorElementsVT(*DAG.getContext());
SDValue HiZero = LaneVT.isInteger() ? DAG.getConstant(0, DL, HalfVT)
: DAG.getConstantFP(0.0, DL, HalfVT);
SDValue LoHalf =
LaneVT.getSizeInBits() == 64
? DAG.getNode(ISD::SCALAR_TO_VECTOR, DL, HalfVT, LowHalfFirstVal)
: DAG.getNode(AArch64ISD::DUP, DL, HalfVT, LowHalfFirstVal);
return DAG.getNode(ISD::CONCAT_VECTORS, DL, VT, LoHalf, HiZero);
}
if (AllLanesExtractElt) {
SDNode *Vector = nullptr;
bool Even = false;

View File

@ -0,0 +1,18 @@
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
; RUN: llc -mtriple=aarch64-none-linux-gnu %s -o - | FileCheck %s
; Regression test for AArch64 compile-time regression, referring to PR #166962.
define <2 x i64> @_ZN10tensorflow12_GLOBAL__N_125ComputeXWeightsAndIndicesERKNS_17ImageResizerStateEbPNSt3__u6vectorINS0_17WeightsAndIndicesENS4_9allocatorIS6_EEEE(<2 x i64> %0) {
; CHECK-LABEL: _ZN10tensorflow12_GLOBAL__N_125ComputeXWeightsAndIndicesERKNS_17ImageResizerStateEbPNSt3__u6vectorINS0_17WeightsAndIndicesENS4_9allocatorIS6_EEEE:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: adrp x8, .LCPI0_0
; CHECK-NEXT: ldr q1, [x8, :lo12:.LCPI0_0]
; CHECK-NEXT: cmgt v2.2d, v1.2d, v0.2d
; CHECK-NEXT: bif v0.16b, v1.16b, v2.16b
; CHECK-NEXT: ret
entry:
%1 = tail call <2 x i64> @llvm.smin.v2i64(<2 x i64> %0, <2 x i64> <i64 -1, i64 0>)
ret <2 x i64> %1
}
declare <2 x i64> @llvm.smin.v2i64(<2 x i64>, <2 x i64>)

View File

@ -543,9 +543,8 @@ define i8 @addv_zero_lanes_negative_v8i8(ptr %arr) {
define i8 @addv_zero_lanes_v16i8(ptr %arr) {
; CHECK-SD-LABEL: addv_zero_lanes_v16i8:
; CHECK-SD: // %bb.0:
; CHECK-SD-NEXT: movi v0.2d, #0000000000000000
; CHECK-SD-NEXT: ldrb w8, [x0]
; CHECK-SD-NEXT: mov v0.d[0], x8
; CHECK-SD-NEXT: ldr b0, [x0]
; CHECK-SD-NEXT: fmov d0, d0
; CHECK-SD-NEXT: addv b0, v0.16b
; CHECK-SD-NEXT: fmov w0, s0
; CHECK-SD-NEXT: ret
@ -568,9 +567,8 @@ define i8 @addv_zero_lanes_v16i8(ptr %arr) {
define i16 @addv_zero_lanes_v8i16(ptr %arr) {
; CHECK-SD-LABEL: addv_zero_lanes_v8i16:
; CHECK-SD: // %bb.0:
; CHECK-SD-NEXT: movi v0.2d, #0000000000000000
; CHECK-SD-NEXT: ldrh w8, [x0]
; CHECK-SD-NEXT: mov v0.d[0], x8
; CHECK-SD-NEXT: ldr h0, [x0]
; CHECK-SD-NEXT: fmov d0, d0
; CHECK-SD-NEXT: addv h0, v0.8h
; CHECK-SD-NEXT: fmov w0, s0
; CHECK-SD-NEXT: ret
@ -593,9 +591,8 @@ define i16 @addv_zero_lanes_v8i16(ptr %arr) {
define i32 @addv_zero_lanes_v4i32(ptr %arr) {
; CHECK-SD-LABEL: addv_zero_lanes_v4i32:
; CHECK-SD: // %bb.0:
; CHECK-SD-NEXT: movi v0.2d, #0000000000000000
; CHECK-SD-NEXT: ldr w8, [x0]
; CHECK-SD-NEXT: mov v0.d[0], x8
; CHECK-SD-NEXT: ldr s0, [x0]
; CHECK-SD-NEXT: fmov d0, d0
; CHECK-SD-NEXT: addv s0, v0.4s
; CHECK-SD-NEXT: fmov w0, s0
; CHECK-SD-NEXT: ret

View File

@ -822,15 +822,15 @@ define i64 @red_mla_dup_ext_u8_s8_s64(ptr noalias noundef readonly captures(none
; CHECK-SD-NEXT: // %bb.9: // %vec.epilog.iter.check
; CHECK-SD-NEXT: cbz x11, .LBB6_13
; CHECK-SD-NEXT: .LBB6_10: // %vec.epilog.ph
; CHECK-SD-NEXT: movi v0.2d, #0000000000000000
; CHECK-SD-NEXT: mov w11, w1
; CHECK-SD-NEXT: fmov d0, x8
; CHECK-SD-NEXT: movi v1.2d, #0000000000000000
; CHECK-SD-NEXT: sxtb x11, w11
; CHECK-SD-NEXT: sxtb x8, w11
; CHECK-SD-NEXT: movi v3.2d, #0x000000000000ff
; CHECK-SD-NEXT: dup v2.2s, w11
; CHECK-SD-NEXT: mov x11, x10
; CHECK-SD-NEXT: and x10, x9, #0xfffffffc
; CHECK-SD-NEXT: mov v0.d[0], x8
; CHECK-SD-NEXT: fmov d0, d0
; CHECK-SD-NEXT: dup v2.2s, w8
; CHECK-SD-NEXT: sub x8, x11, x10
; CHECK-SD-NEXT: add x11, x0, x11
; CHECK-SD-NEXT: .LBB6_11: // %vec.epilog.vector.body

View File

@ -318,10 +318,7 @@ define <4 x float> @test_insert_2_f32_undef_zero(float %a) {
define <2 x double> @test_insert_v2f64_undef_insert1(double %a) {
; CHECK-LABEL: test_insert_v2f64_undef_insert1:
; CHECK: // %bb.0:
; CHECK-NEXT: movi.2d v1, #0000000000000000
; CHECK-NEXT: // kill: def $d0 killed $d0 def $q0
; CHECK-NEXT: mov.d v1[0], v0[0]
; CHECK-NEXT: mov.16b v0, v1
; CHECK-NEXT: fmov d0, d0
; CHECK-NEXT: ret
%v.0 = insertelement <2 x double > <double undef, double 0.000000e+00>, double %a, i32 0
ret <2 x double> %v.0

View File

@ -339,9 +339,8 @@ define <8 x i8> @load_sext_i32_v8i8(ptr %p) {
define <16 x i8> @load_zext_v16i8(ptr %p) {
; CHECK-SD-LABEL: load_zext_v16i8:
; CHECK-SD: // %bb.0:
; CHECK-SD-NEXT: movi v0.2d, #0000000000000000
; CHECK-SD-NEXT: ldr w8, [x0]
; CHECK-SD-NEXT: mov v0.d[0], x8
; CHECK-SD-NEXT: ldr s0, [x0]
; CHECK-SD-NEXT: fmov d0, d0
; CHECK-SD-NEXT: ret
;
; CHECK-GI-LABEL: load_zext_v16i8:

View File

@ -599,10 +599,10 @@ entry:
define i128 @i128_mask(i128 %x) {
; CHECK-SD-LABEL: i128_mask:
; CHECK-SD: // %bb.0: // %entry
; CHECK-SD-NEXT: movi v0.2d, #0000000000000000
; CHECK-SD-NEXT: and x8, x0, #0xff
; CHECK-SD-NEXT: mov x1, xzr
; CHECK-SD-NEXT: mov v0.d[0], x8
; CHECK-SD-NEXT: fmov d0, x8
; CHECK-SD-NEXT: fmov d0, d0
; CHECK-SD-NEXT: cnt v0.16b, v0.16b
; CHECK-SD-NEXT: addv b0, v0.16b
; CHECK-SD-NEXT: fmov x0, d0

View File

@ -95,10 +95,7 @@ entry:
define <2 x double> @fadd(double noundef %x, double noundef %y) {
; CHECK-LABEL: fadd:
; CHECK: // %bb.0: // %entry
; CHECK-NEXT: movi v2.2d, #0000000000000000
; CHECK-NEXT: fadd d0, d0, d1
; CHECK-NEXT: mov v2.d[0], v0.d[0]
; CHECK-NEXT: mov v0.16b, v2.16b
; CHECK-NEXT: ret
entry:
%add = fadd double %x, %y

View File

@ -0,0 +1,92 @@
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
; RUN: llc -mtriple=aarch64-linux-gnu -o - %s | FileCheck %s
define <2 x i64> @low_vector_splat_v2i64_from_i64(i64 %0){
; CHECK-LABEL: low_vector_splat_v2i64_from_i64:
; CHECK: // %bb.0:
; CHECK-NEXT: fmov d0, x0
; CHECK-NEXT: fmov d0, d0
; CHECK-NEXT: ret
%2 = insertelement <1 x i64> poison, i64 %0, i64 0
%3 = shufflevector <1 x i64> %2, <1 x i64> zeroinitializer, <2 x i32> <i32 0, i32 1>
ret <2 x i64> %3
}
define <4 x i32> @low_vector_splat_v4i32_from_i32(i32 %0) {
; CHECK-LABEL: low_vector_splat_v4i32_from_i32:
; CHECK: // %bb.0:
; CHECK-NEXT: dup v0.2s, w0
; CHECK-NEXT: ret
%2 = insertelement <2 x i32> poison, i32 %0, i64 0
%3 = shufflevector <2 x i32> %2, <2 x i32> poison, <2 x i32> zeroinitializer
%4 = shufflevector <2 x i32> %3, <2 x i32> zeroinitializer, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
ret <4 x i32> %4
}
define <8 x i16> @low_vector_splat_v8i16_from_i16(i16 %0) {
; CHECK-LABEL: low_vector_splat_v8i16_from_i16:
; CHECK: // %bb.0:
; CHECK-NEXT: dup v0.4h, w0
; CHECK-NEXT: ret
%2 = insertelement <4 x i16> poison, i16 %0, i64 0
%3 = shufflevector <4 x i16> %2, <4 x i16> poison, <4 x i32> zeroinitializer
%4 = shufflevector <4 x i16> %3, <4 x i16> zeroinitializer, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
ret <8 x i16> %4
}
define <16 x i8> @low_vector_splat_v16i8_from_i8(i8 %0) {
; CHECK-LABEL: low_vector_splat_v16i8_from_i8:
; CHECK: // %bb.0:
; CHECK-NEXT: dup v0.8b, w0
; CHECK-NEXT: ret
%2 = insertelement <8 x i8> poison, i8 %0, i64 0
%3 = shufflevector <8 x i8> %2, <8 x i8> poison, <8 x i32> zeroinitializer
%4 = shufflevector <8 x i8> %3, <8 x i8> zeroinitializer, <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
ret <16 x i8> %4
}
define <2 x double> @low_vector_splat_v2f64_from_f64(double %0) {
; CHECK-LABEL: low_vector_splat_v2f64_from_f64:
; CHECK: // %bb.0:
; CHECK-NEXT: fmov d0, d0
; CHECK-NEXT: ret
%2 = insertelement <1 x double> poison, double %0, i64 0
%3 = shufflevector <1 x double> %2, <1 x double> zeroinitializer, <2 x i32> <i32 0, i32 1>
ret <2 x double> %3
}
define <4 x float> @low_vector_splat_v4f32_from_f32(float %0) {
; CHECK-LABEL: low_vector_splat_v4f32_from_f32:
; CHECK: // %bb.0:
; CHECK-NEXT: // kill: def $s0 killed $s0 def $q0
; CHECK-NEXT: dup v0.2s, v0.s[0]
; CHECK-NEXT: ret
%2 = insertelement <2 x float> poison, float %0, i64 0
%3 = shufflevector <2 x float> %2, <2 x float> poison, <2 x i32> zeroinitializer
%4 = shufflevector <2 x float> %3, <2 x float> zeroinitializer, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
ret <4 x float> %4
}
define <8 x half> @low_vector_splat_v8f16_from_f16(half %0) {
; CHECK-LABEL: low_vector_splat_v8f16_from_f16:
; CHECK: // %bb.0:
; CHECK-NEXT: // kill: def $h0 killed $h0 def $q0
; CHECK-NEXT: dup v0.4h, v0.h[0]
; CHECK-NEXT: ret
%2 = insertelement <4 x half> poison, half %0, i64 0
%3 = shufflevector <4 x half> %2, <4 x half> poison, <4 x i32> zeroinitializer
%4 = shufflevector <4 x half> %3, <4 x half> zeroinitializer, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
ret <8 x half> %4
}
define <8 x bfloat> @low_vector_splat_v8bf16_from_bf16(bfloat %0) {
; CHECK-LABEL: low_vector_splat_v8bf16_from_bf16:
; CHECK: // %bb.0:
; CHECK-NEXT: // kill: def $h0 killed $h0 def $q0
; CHECK-NEXT: dup v0.4h, v0.h[0]
; CHECK-NEXT: ret
%2 = insertelement <4 x bfloat> poison, bfloat %0, i64 0
%3 = shufflevector <4 x bfloat> %2, <4 x bfloat> poison, <4 x i32> zeroinitializer
%4 = shufflevector <4 x bfloat> %3, <4 x bfloat> zeroinitializer, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
ret <8 x bfloat> %4
}