[AArch64][LV] Adjust costs for low-VF interleaved access

[jacob-crawley]

Addressing regression introduced by #205844 in which a significantly slower SVE tail loop is generated.

The cost model for the case where the interleave factor is larger than the VF has been adjusted to more accurately reflect the cost of the uzp instructions generated by the deinterleave tree, and the cost of legalizing the type of each subvector.

[llvmorg-github-actions]

@llvm/pr-subscribers-backend-aarch64

@llvm/pr-subscribers-llvm-transforms

Author: Jacob Crawley (jacob-crawley)

Changes

Addressing regression introduced by #205844 in which a significantly slower SVE tail loop is generated.

The cost model for the case where the interleave factor is larger than the VF has been adjusted to more accurately reflect the cost of the uzp instructions generated by the deinterleave tree, and the cost of legalizing the type of each subvector.

---

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

3 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp (+13-7)
• (modified) llvm/test/Transforms/LoopVectorize/AArch64/sve-interleave-low-vf-cost.ll (+70-9)
• (modified) llvm/test/Transforms/LoopVectorize/AArch64/sve-interleaved-access-low-vf.ll (+1-1)

diff --git a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
index 9e1a0960617a5..20ffb1d9c342e 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
@@ -5460,13 +5460,19 @@ InstructionCost AArch64TTIImpl::getInterleavedMemoryOpCost(
       }
 
       // llvm.vector.deinterleaveN is lowered as a binary tree of deinterleave2
-      // operations. A binary tree producing Factor leaf vectors has
-      // (Factor -1) inner deinterleave2 nodes. Each deinterleave2 on a pair of
-      // SVE registers emits one uzp1 + one uzp2.
-      // Total shuffle cost: (Factor - 1) deinterleave2 operations, each
-      // processing LT.first legal vector parts,with one uzp shuffle per part.
-      auto LT = getTypeLegalizationCost(VecTy);
-      return MemCost + (Factor - 1) * LT.first;
+      // operations. The tree has Log2(Factor) levels, with Factor UZP/ZIP
+      // operations at each level, giving a total shuffle cost of
+      // Factor * Log2(Factor).
+
+      // For stores, account for an additional legalization cost when
+      // repacking the legalized subvectors into the narrow interleaved
+      // vector.
+      auto LegalizationCost = getTypeLegalizationCost(SubVecTy).first;
+
+      if (Opcode == Instruction::Store)
+        LegalizationCost *= 2;
+
+      return MemCost + (Factor * LegalizationCost) + (Factor * Log2_64(Factor));
     }
   }
 
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/sve-interleave-low-vf-cost.ll b/llvm/test/Transforms/LoopVectorize/AArch64/sve-interleave-low-vf-cost.ll
index 9e99bedb9ed0c..fde4e52737292 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/sve-interleave-low-vf-cost.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/sve-interleave-low-vf-cost.ll
@@ -5,17 +5,18 @@ target triple = "aarch64"
 
 ; Cost-model test for the path where the interleave factor is > the VF.
 ; The cost is modelled as a contiguous load of the wide vector plus a
-; vector.deinterleave4 lowered as a binary tree of (Factor - 1) deinterleave2 shuffles:
+; vector.deinterleave4 lowered as a binary tree of Log2(Factor) levels,
+; each with Factor operations.
 ;
-;   cost = MemCost + (Factor - 1) * LT.first
-;  * VF vscale x 2: wide type <vscale x 8 x i16>, LT.first = 1
-;       => 1 (load) + 3 * 1 (shuffles) = 4
-;  * VF vscale x 4: wide type <vscale x 16 x i16> LT.first = 2
-;       => 2 (loads) + 3 * 2 (shuffles) = 8
+;   cost = MemCost + (Factor * LegalizationCost) + (Factor * Log2(Factor))
+;  * VF vscale x 2: wide type <vscale x 8 x i16>, MemCost = 1, Subvector Legalization Cost = 1
+;       => 1 + (4 * 1) + (4 * 2) = 13
+;  * VF vscale x 4: wide type <vscale x 16 x i16>, MemCost = 2, Subvector Legalization Cost = 1
+;       => 2 + (4 * 1) + (4 * 2) = 14
 
 ; CHECK-LABEL: LV: Checking a loop in 'deinterleave4_nxv2i16_load'
-; CHECK: Cost of 4 for VF vscale x 2: INTERLEAVE-GROUP with factor 4, ir<%ptr.b>
-; CHECK: Cost of 8 for VF vscale x 4: INTERLEAVE-GROUP with factor 4, ir<%ptr.b>
+; CHECK: Cost of 13 for VF vscale x 2: INTERLEAVE-GROUP with factor 4, ir<%ptr.b>
+; CHECK: Cost of 14 for VF vscale x 4: INTERLEAVE-GROUP with factor 4, ir<%ptr.b>
 ; CHECK: LV: Selecting VF: vscale x 2
 define void @deinterleave4_nxv2i16_load(ptr noalias readonly %src, ptr noalias %out, i64 %n) #0 {
 entry:
@@ -69,4 +70,64 @@ exit:
   ret void
 }
 
-attributes #0 = { "target-features"="+sve" }
\ No newline at end of file
+; Check that the increased low-VF interleaved-store cost prevents selection of
+; an SVE epilogue.
+
+; For VF vscale x 4:
+;   load cost  = 2 + (4 * 1) + (4 * 2) = 14
+;   store cost = 1 + (4 * 2) + (4 * 2) = 17
+;
+; This makes the fixed VF 8 epilogue preferable to VF vscale x 4.
+;
+; CHECK-LABEL: LV: Checking a loop in 'deinterleave4_nxv4i16_load_interleave4_nxv4i8_store'
+; CHECK: Cost of 14 for VF vscale x 4: INTERLEAVE-GROUP with factor 4
+; CHECK: Cost of 17 for VF vscale x 4: INTERLEAVE-GROUP with factor 4
+; CHECK: LV: Selecting VF: vscale x 16
+; CHECK: LEV: Vectorizing epilogue loop with VF = 8
+define void @deinterleave4_nxv4i16_load_interleave4_nxv4i8_store(
+    ptr readonly %src, ptr writeonly %out, i32 %n) #0 {
+entry:
+  %empty = icmp eq i32 %n, 0
+  br i1 %empty, label %exit, label %loop
+
+loop:
+  %src.iv = phi ptr [ %src.next, %loop ], [ %src, %entry ]
+  %out.iv = phi ptr [ %out.next, %loop ], [ %out, %entry ]
+  %iv = phi i32 [ %iv.next, %loop ], [ %n, %entry ]
+
+  %ptr.g = getelementptr inbounds i16, ptr %src.iv, i64 1
+  %ptr.r = getelementptr inbounds i16, ptr %src.iv, i64 2
+  %ptr.a = getelementptr inbounds i16, ptr %src.iv, i64 3
+  %load.b = load i16, ptr %src.iv, align 2
+  %load.g = load i16, ptr %ptr.g, align 2
+  %load.r = load i16, ptr %ptr.r, align 2
+  %load.a = load i16, ptr %ptr.a, align 2
+
+  %shift.b = lshr i16 %load.b, 8
+  %shift.g = lshr i16 %load.g, 8
+  %shift.r = lshr i16 %load.r, 8
+  %shift.a = lshr i16 %load.a, 8
+  %trunc.b = trunc nuw i16 %shift.b to i8
+  %trunc.g = trunc nuw i16 %shift.g to i8
+  %trunc.r = trunc nuw i16 %shift.r to i8
+  %trunc.a = trunc nuw i16 %shift.a to i8
+
+  %out.g = getelementptr inbounds i8, ptr %out.iv, i64 1
+  %out.r = getelementptr inbounds i8, ptr %out.iv, i64 2
+  %out.a = getelementptr inbounds i8, ptr %out.iv, i64 3
+  store i8 %trunc.b, ptr %out.iv, align 1
+  store i8 %trunc.g, ptr %out.g, align 1
+  store i8 %trunc.r, ptr %out.r, align 1
+  store i8 %trunc.a, ptr %out.a, align 1
+
+  %src.next = getelementptr inbounds i16, ptr %src.iv, i64 4
+  %out.next = getelementptr inbounds i8, ptr %out.iv, i64 4
+  %iv.next = add nsw i32 %iv, -1
+  %done = icmp eq i32 %iv.next, 0
+  br i1 %done, label %exit, label %loop
+
+exit:
+  ret void
+}
+
+attributes #0 = { "target-features"="+sve" }
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/sve-interleaved-access-low-vf.ll b/llvm/test/Transforms/LoopVectorize/AArch64/sve-interleaved-access-low-vf.ll
index 035a9f80fae4a..7b6fe0a9c6b1f 100644
--- a/llvm/test/Transforms/LoopVectorize/AArch64/sve-interleaved-access-low-vf.ll
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/sve-interleaved-access-low-vf.ll
@@ -1,4 +1,4 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals none --filter-out-after "^scalar.ph" --version 6
 ; RUN: opt -passes=loop-vectorize -S < %s | FileCheck %s
 
 target triple = "aarch64"

[Stylie777]

LGTM, thanks!

[sdesmalen-arm]

Looking at https://godbolt.org/z/qeezP1c3n, I think the only case where codegen is currently bad (and could be improved with simply 3 x uzp1's for the concatenation of the interleaved results) is the case where after legalising the input/output types, the input element type is 4 x the size of the output element type.

i.e.

define <vscale x 8 x i16> @interleave4_nxv2i16(<vscale x 2 x i16> %vec0, <vscale x 2 x i16> %vec1, <vscale x 2 x i16> %vec2, <vscale x 2 x i16> %vec3) {
  %retval = call <vscale x 8 x i16> @llvm.vector.interleave4.nxv8i16(<vscale x 2 x i16> %vec0, <vscale x 2 x i16> %vec1, <vscale x 2 x i16> %vec2, <vscale x 2 x i16> %vec3)
  ret <vscale x 8 x i16> %retval
}

define <vscale x 16 x i8> @interleave4_nxv4i8(<vscale x 4 x i8> %vec0, <vscale x 4 x i8> %vec1, <vscale x 4 x i8> %vec2, <vscale x 4 x i8> %vec3) {
  %retval = call <vscale x 16 x i8> @llvm.vector.interleave4.nxv8i8(<vscale x 4 x i8> %vec0, <vscale x 4 x i8> %vec1, <vscale x 4 x i8> %vec2, <vscale x 4 x i8> %vec3)
  ret <vscale x 16 x i8> %retval
}

where the input types are legalised to <vscale x 2 x i64> and <vscale x 4 x i32> respectively.

For other cases, the formula used to calculate the cost seems fine.

[MacDue]

You can compare the codegen of i16 and half to see with/without the legalization to <vscale x 2 x i64>: https://godbolt.org/z/PTY6nYxo3

[sdesmalen-arm]

nit: maybe add FIXME to it, to make it clear this must be removed.

[MacDue]

LGTM

[MacDue]

/cherry-pick 79e05f4

[llvmbot]

/pull-request #210696