[AArch64] Fix incorrect use of MVT::getVectorNumElements in AArch64TTIImpl::getVectorInstrCost

[david-arm]

Summary

If we are inserting into or extracting from a scalable vector we do
not know the number of elements at runtime, so we can only let the
index wrap for fixed-length vectors.

Tests added here:

Analysis/CostModel/AArch64/sve-insert-extract.ll

[c-rhodes]

LGTM. I think you meant compile-time not runtime in the commit message?

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

InstructionCost AArch64TTIImpl::getVectorInstrCost(unsigned Opcode, Type *Val,
   if (Index != -1U) {
     // Legalize the type.
     std::pair<InstructionCost, MVT> LT = TLI->getTypeLegalizationCost(DL, Val);
 
     // This type is legalized to a scalar type.
     if (!LT.second.isVector())
       return 0;
 
-    // The type may be split. Normalize the index to the new type.
+    // The type may be split. For fixed-width vectors we can normalize the
+    // index to the new type.
+    if (LT.second.isFixedLengthVector()) {
       unsigned Width = LT.second.getVectorNumElements();
       Index = Index % Width;
+    }
 
     // The element at index zero is already inside the vector.
     if (Index == 0)
       return 0;
   }
 
   // All other insert/extracts cost this much.
   return ST->getVectorInsertExtractBaseCost();

llvm/test/Analysis/CostModel/AArch64/sve-insert-extract.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -cost-model -analyze -S < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+
+define void @ins_el0() #0 {
+; CHECK: Cost Model: Found an estimated cost of 0 for instruction:   %v0 = insertelement <vscale x 16 x i8> zeroinitializer, i8 0, i64 0
+; CHECK: Cost Model: Found an estimated cost of 0 for instruction:   %v1 = insertelement <vscale x 8 x i16> zeroinitializer, i16 0, i64 0
+; CHECK: Cost Model: Found an estimated cost of 0 for instruction:   %v2 = insertelement <vscale x 4 x i32> zeroinitializer, i32 0, i64 0
+; CHECK: Cost Model: Found an estimated cost of 0 for instruction:   %v3 = insertelement <vscale x 2 x i64> zeroinitializer, i64 0, i64 0
+  %v0 = insertelement <vscale x 16 x i8> zeroinitializer, i8 0, i64 0
+  %v1 = insertelement <vscale x 8 x i16> zeroinitializer, i16 0, i64 0
+  %v2 = insertelement <vscale x 4 x i32> zeroinitializer, i32 0, i64 0
+  %v3 = insertelement <vscale x 2 x i64> zeroinitializer, i64 0, i64 0
+  ret void
+}
+
+define void @ins_el1() #0 {
+; CHECK: Cost Model: Found an estimated cost of 3 for instruction:   %v0 = insertelement <vscale x 16 x i8> zeroinitializer, i8 0, i64 1
+; CHECK: Cost Model: Found an estimated cost of 3 for instruction:   %v1 = insertelement <vscale x 8 x i16> zeroinitializer, i16 0, i64 1
+; CHECK: Cost Model: Found an estimated cost of 3 for instruction:   %v2 = insertelement <vscale x 4 x i32> zeroinitializer, i32 0, i64 1
+; CHECK: Cost Model: Found an estimated cost of 3 for instruction:   %v3 = insertelement <vscale x 2 x i64> zeroinitializer, i64 0, i64 1
+  %v0 = insertelement <vscale x 16 x i8> zeroinitializer, i8 0, i64 1
+  %v1 = insertelement <vscale x 8 x i16> zeroinitializer, i16 0, i64 1
+  %v2 = insertelement <vscale x 4 x i32> zeroinitializer, i32 0, i64 1
+  %v3 = insertelement <vscale x 2 x i64> zeroinitializer, i64 0, i64 1
+  ret void
+}
+
+
+define void @ext_el0() #0 {
+; CHECK: Cost Model: Found an estimated cost of 0 for instruction:   %v0 = extractelement <vscale x 16 x i8> zeroinitializer, i64 0
+; CHECK: Cost Model: Found an estimated cost of 0 for instruction:   %v1 = extractelement <vscale x 8 x i16> zeroinitializer, i64 0
+; CHECK: Cost Model: Found an estimated cost of 0 for instruction:   %v2 = extractelement <vscale x 4 x i32> zeroinitializer, i64 0
+; CHECK: Cost Model: Found an estimated cost of 0 for instruction:   %v3 = extractelement <vscale x 2 x i64> zeroinitializer, i64 0
+  %v0 = extractelement <vscale x 16 x i8> zeroinitializer, i64 0
+  %v1 = extractelement <vscale x 8 x i16> zeroinitializer, i64 0
+  %v2 = extractelement <vscale x 4 x i32> zeroinitializer, i64 0
+  %v3 = extractelement <vscale x 2 x i64> zeroinitializer, i64 0
+  ret void
+}
+
+define void @ext_el1() #0 {
+; CHECK: Cost Model: Found an estimated cost of 3 for instruction:   %v0 = extractelement <vscale x 16 x i8> zeroinitializer, i64 1
+; CHECK: Cost Model: Found an estimated cost of 3 for instruction:   %v1 = extractelement <vscale x 8 x i16> zeroinitializer, i64 1
+; CHECK: Cost Model: Found an estimated cost of 3 for instruction:   %v2 = extractelement <vscale x 4 x i32> zeroinitializer, i64 1
+; CHECK: Cost Model: Found an estimated cost of 3 for instruction:   %v3 = extractelement <vscale x 2 x i64> zeroinitializer, i64 1
+  %v0 = extractelement <vscale x 16 x i8> zeroinitializer, i64 1
+  %v1 = extractelement <vscale x 8 x i16> zeroinitializer, i64 1
+  %v2 = extractelement <vscale x 4 x i32> zeroinitializer, i64 1
+  %v3 = extractelement <vscale x 2 x i64> zeroinitializer, i64 1
+  ret void
+}
+
+
+attributes #0 = { "target-features"="+sve" vscale_range(1, 16) }

llvm/test/Analysis/CostModel/AArch64/sve-intrinsics.ll

 ; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py
 ; RUN: opt < %s -analyze -cost-model -S -mtriple=aarch64--linux-gnu -mattr=+sve | FileCheck %s
 
 define void @vector_insert_extract(<vscale x 4 x i32> %v0, <vscale x 16 x i32> %v1, <16 x i32> %v2) {
 ; CHECK-LABEL: 'vector_insert_extract'
-; CHECK-NEXT:  Cost Model: Found an estimated cost of 72 for instruction: %extract_fixed_from_scalable = call <16 x i32> @llvm.experimental.vector.extract.v16i32.nxv4i32(<vscale x 4 x i32> %v0, i64 0)
-; CHECK-NEXT:  Cost Model: Found an estimated cost of 72 for instruction: %insert_fixed_into_scalable = call <vscale x 4 x i32> @llvm.experimental.vector.insert.nxv4i32.v16i32(<vscale x 4 x i32> %v0, <16 x i32> %v2, i64 0)
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 81 for instruction: %extract_fixed_from_scalable = call <16 x i32> @llvm.experimental.vector.extract.v16i32.nxv4i32(<vscale x 4 x i32> %v0, i64 0)
+; CHECK-NEXT:  Cost Model: Found an estimated cost of 81 for instruction: %insert_fixed_into_scalable = call <vscale x 4 x i32> @llvm.experimental.vector.insert.nxv4i32.v16i32(<vscale x 4 x i32> %v0, <16 x i32> %v2, i64 0)
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %extract_scalable_from_scalable = call <vscale x 4 x i32> @llvm.experimental.vector.extract.nxv4i32.nxv16i32(<vscale x 16 x i32> %v1, i64 0)
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 1 for instruction: %insert_scalable_into_scalable = call <vscale x 16 x i32> @llvm.experimental.vector.insert.nxv16i32.nxv4i32(<vscale x 16 x i32> %v1, <vscale x 4 x i32> %v0, i64 0)
 ; CHECK-NEXT:  Cost Model: Found an estimated cost of 0 for instruction: ret void
 ;
   %extract_fixed_from_scalable = call <16 x i32> @llvm.experimental.vector.extract.v16i32.nxv4i32(<vscale x 4 x i32> %v0, i64 0)
   %insert_fixed_into_scalable = call <vscale x 4 x i32> @llvm.experimental.vector.insert.nxv4i32.v16i32(<vscale x 4 x i32> %v0, <16 x i32> %v2, i64 0)
   %extract_scalable_from_scalable = call <vscale x 4 x i32> @llvm.experimental.vector.extract.nxv4i32.nxv16i32(<vscale x 16 x i32> %v1, i64 0)
   %insert_scalable_into_scalable = call <vscale x 16 x i32> @llvm.experimental.vector.insert.nxv16i32.nxv4i32(<vscale x 16 x i32> %v1, <vscale x 4 x i32> %v0, i64 0)