[BasicTTIImpl] Fix getCastInstrCost for scalable vectors by querying for ElementCount.

[sdesmalen]

Summary

This fixes an overly restrictive assumption that the vector is a FixedVectorType,
in code that tries to calculate the cost of a cast operation when splitting
a too-wide vector. The algorithm works the same for scalable vectors, so this
patch removes the cast<FixedVectorType>.

[frasercrmck]

LGTM (though I can't speak to whether the costs make sense for AArch64)

[david-arm]

llvm/test/Analysis/CostModel/AArch64/cast.ll

711

Hi @sdesmalen, given you've added support in the cost function for scalable vectors of the type <vscale x 1 x ElTy> perhaps it's worth adding a test for that too? i.e.

%loadnxv1i32 = load <vscale x 1 x i32>, <vscale x 1 x i32>* undef

[sdesmalen]

llvm/test/Analysis/CostModel/AArch64/cast.ll

711

I could, but I'm not sure how helpful that is, because we can't currently vectorize <vscale x 1 x <eltty>>. For this the cost should return Invalid, but it just returns some value now. This is probably something to fix in a follow-up patch at some point. It also wouldn't help test the code-change in this patch.

[david-arm]

LGTM!

llvm/test/Analysis/CostModel/AArch64/cast.ll

711

OK that's fine. I think it'd be good to fix getCastInstrCost to deal with the <vscale x 1 x> case properly then at some point, since at the moment it just lets them through. I guess this is a general problem for a lot of our cost model right now.

llvm/include/llvm/CodeGen/BasicTTIImpl.h

if (DstVTy && SrcVTy) {
       // cost of the split itself. Count that as 1, to be consistent with
       // TLI->getTypeLegalizationCost().
       bool SplitSrc =
           TLI->getTypeAction(Src->getContext(), TLI->getValueType(DL, Src)) ==
           TargetLowering::TypeSplitVector;
       bool SplitDst =
           TLI->getTypeAction(Dst->getContext(), TLI->getValueType(DL, Dst)) ==
           TargetLowering::TypeSplitVector;
-      if ((SplitSrc || SplitDst) &&
-          cast<FixedVectorType>(SrcVTy)->getNumElements() > 1 &&
-          cast<FixedVectorType>(DstVTy)->getNumElements() > 1) {
+      if ((SplitSrc || SplitDst) && SrcVTy->getElementCount().isVector() &&
+          DstVTy->getElementCount().isVector()) {
         Type *SplitDstTy = VectorType::getHalfElementsVectorType(DstVTy);
         Type *SplitSrcTy = VectorType::getHalfElementsVectorType(SrcVTy);
         T *TTI = static_cast<T *>(this);
         // If both types need to be split then the split is free.
         unsigned SplitCost =
             (!SplitSrc || !SplitDst) ? TTI->getVectorSplitCost() : 0;
         return SplitCost +
                (2 * TTI->getCastInstrCost(Opcode, SplitDstTy, SplitSrcTy, CCH,

llvm/test/Analysis/CostModel/AArch64/cast.ll

;
   %d = bitcast float undef to i32
   %e = bitcast i64 undef to double
   %f = bitcast double undef to i64
   %g = bitcast half undef to i16
   %h = bitcast i16 undef to half
   ret i32 undef
 }
 
-define i32 @load_extends() {
+define i32 @load_extends() #0 {
 ; CHECK-LABEL: 'load_extends'
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %loadi8 = load i8, i8* undef
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %loadi16 = load i16, i16* undef
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %loadi32 = load i32, i32* undef
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 16 for instruction: %loadv2i8 = load <2 x i8>, <2 x i8>* undef
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 64 for instruction: %loadv4i8 = load <4 x i8>, <4 x i8>* undef
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %loadv8i8 = load <8 x i8>, <8 x i8>* undef
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %loadv2i16 = load <2 x i16>, <2 x i16>* undef
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %loadv4i16 = load <4 x i16>, <4 x i16>* undef
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %loadv2i32 = load <2 x i32>, <2 x i32>* undef
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %loadv4i32 = load <4 x i32>, <4 x i32>* undef
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %loadnxv2i32 = load <vscale x 2 x i32>, <vscale x 2 x i32>* undef
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %loadnxv4i32 = load <vscale x 4 x i32>, <vscale x 4 x i32>* undef
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %r0 = sext i8 %loadi8 to i16
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %r1 = zext i8 %loadi8 to i16
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %r2 = sext i8 %loadi8 to i32
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %r3 = zext i8 %loadi8 to i32
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %r4 = sext i8 %loadi8 to i64
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %r5 = zext i8 %loadi8 to i64
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %r6 = sext i16 %loadi16 to i32
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %r7 = zext i16 %loadi16 to i32

 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %v6 = sext <4 x i16> %loadv4i16 to <4 x i32>
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %v7 = zext <4 x i16> %loadv4i16 to <4 x i32>
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %v8 = sext <2 x i16> %loadv2i16 to <2 x i64>
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %v9 = zext <2 x i16> %loadv2i16 to <2 x i64>
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %v10 = sext <2 x i32> %loadv2i32 to <2 x i64>
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %v11 = zext <2 x i32> %loadv2i32 to <2 x i64>
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %v12 = sext <4 x i32> %loadv4i32 to <4 x i64>
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 2 for instruction: %v13 = zext <4 x i32> %loadv4i32 to <4 x i64>
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %v14 = sext <vscale x 2 x i32> %loadnxv2i32 to <vscale x 2 x i64>
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %v15 = zext <vscale x 2 x i32> %loadnxv2i32 to <vscale x 2 x i64>
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %v16 = sext <vscale x 4 x i32> %loadnxv4i32 to <vscale x 4 x i64>
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %v17 = zext <vscale x 4 x i32> %loadnxv4i32 to <vscale x 4 x i64>
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: ret i32 undef
 ;
   %loadi8 = load i8, i8* undef
   %loadi16 = load i16, i16* undef
   %loadi32 = load i32, i32* undef
   %loadv2i8 = load <2 x i8>, <2 x i8>* undef
   %loadv4i8 = load <4 x i8>, <4 x i8>* undef
   %loadv8i8 = load <8 x i8>, <8 x i8>* undef
   %loadv2i16 = load <2 x i16>, <2 x i16>* undef
   %loadv4i16 = load <4 x i16>, <4 x i16>* undef
   %loadv2i32 = load <2 x i32>, <2 x i32>* undef
   %loadv4i32 = load <4 x i32>, <4 x i32>* undef
+  %loadnxv2i32 = load <vscale x 2 x i32>, <vscale x 2 x i32>* undef
+  %loadnxv4i32 = load <vscale x 4 x i32>, <vscale x 4 x i32>* undef
 
   %r0 = sext i8 %loadi8 to i16
   %r1 = zext i8 %loadi8 to i16
   %r2 = sext i8 %loadi8 to i32
   %r3 = zext i8 %loadi8 to i32
   %r4 = sext i8 %loadi8 to i64
   %r5 = zext i8 %loadi8 to i64
   %r6 = sext i16 %loadi16 to i32
;
   %v6 = sext <4 x i16> %loadv4i16 to <4 x i32>
   %v7 = zext <4 x i16> %loadv4i16 to <4 x i32>
   %v8 = sext <2 x i16> %loadv2i16 to <2 x i64>
   %v9 = zext <2 x i16> %loadv2i16 to <2 x i64>
   %v10 = sext <2 x i32> %loadv2i32 to <2 x i64>
   %v11 = zext <2 x i32> %loadv2i32 to <2 x i64>
   %v12 = sext <4 x i32> %loadv4i32 to <4 x i64>
   %v13 = zext <4 x i32> %loadv4i32 to <4 x i64>
+  %v14 = sext <vscale x 2 x i32> %loadnxv2i32 to <vscale x 2 x i64>
+  %v15 = zext <vscale x 2 x i32> %loadnxv2i32 to <vscale x 2 x i64>
+  %v16 = sext <vscale x 4 x i32> %loadnxv4i32 to <vscale x 4 x i64>
+  %v17 = zext <vscale x 4 x i32> %loadnxv4i32 to <vscale x 4 x i64>
 
   ret i32 undef
 }
 
 define i32 @store_truncs() {
 ; CHECK-LABEL: 'store_truncs'
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: %r0 = trunc i64 undef to i8
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: store i8 %r0, i8* undef
;
   %r3 = trunc i32 undef to i8
   store i8 %r3, i8* undef
   %r4 = trunc i32 undef to i16
   store i16 %r4, i16* undef
   %r5 = trunc i16 undef to i8
   store i8 %r5, i8* undef
   ret i32 undef
 }
+
+attributes #0 = { "target-features"="+sve" }