[llvm] [AArch64][LV] Adjust costs for low-VF interleaved access (PR #209441)

[Sander de Smalen via llvm-commits]

================
@@ -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).
+      llvm::InstructionCost LegalizationCost =
+          getTypeLegalizationCost(SubVecTy).first;
+
+      // For stores, account for an additional legalization cost when
+      // repacking the legalized subvectors into the narrow interleaved
+      // vector.
+      if (Opcode == Instruction::Store)
+        LegalizationCost *= 2;
+
+      return MemCost + (Factor * LegalizationCost) + (Factor * Log2_64(Factor));
----------------
sdesmalen-arm wrote:

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.

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