[llvm] [AArch64] Improve cost model for legal subvec insert/extract (PR #81135)

[Sander de Smalen via llvm-commits]

================
@@ -568,6 +568,64 @@ AArch64TTIImpl::getIntrinsicInstrCost(const IntrinsicCostAttributes &ICA,
     }
     return Cost;
   }
+  case Intrinsic::vector_extract: {
+    // If both the vector argument and the return type are legal types and the
+    // index is 0, then this should be a no-op or simple operation; return a
+    // relatively low cost.
+
+    // If arguments aren't actually supplied, then we cannot determine the
+    // value of the index. We also want to skip this on predicate types.
+    if (ICA.getArgs().size() != 2 ||
+        ICA.getReturnType()->getScalarType()->isIntegerTy(1))
+      break;
+    LLVMContext &C = RetTy->getContext();
+    EVT MRTy = getTLI()->getValueType(DL, RetTy);
+    EVT MPTy = getTLI()->getValueType(DL, ICA.getArgTypes()[0]);
+    // Skip this if either the return type or the vector argument are unpacked
+    // SVE types; they may get lowered to stack stores and loads.
+    if ((MRTy.isScalableVector() &&
+         MRTy.getSizeInBits().getKnownMinValue() != AArch64::SVEBitsPerBlock) ||
+        (MPTy.isScalableVector() &&
+         MPTy.getSizeInBits().getKnownMinValue() != AArch64::SVEBitsPerBlock))
+      break;
+    TargetLoweringBase::LegalizeKind RLK = getTLI()->getTypeConversion(C, MRTy);
+    TargetLoweringBase::LegalizeKind PLK = getTLI()->getTypeConversion(C, MPTy);
+    const ConstantInt *Idx = dyn_cast<ConstantInt>(ICA.getArgs()[1]);
+    if (RLK.first == TargetLoweringBase::TypeLegal &&
+        PLK.first == TargetLoweringBase::TypeLegal && Idx &&
+        Idx->getZExtValue() == 0)
+      return TTI::TCC_Free;
+    break;
+  }
+  case Intrinsic::vector_insert: {
+    // If both the vector and subvector arguments are legal types and the index
+    // is 0, then this should be a no-op or simple operation; return a
+    // relatively low cost.
+
+    // If arguments aren't actually supplied, then we cannot determine the
+    // value of the index. We also want to skip this on predicate types.
+    if (ICA.getArgs().size() != 3 ||
+        ICA.getReturnType()->getScalarType()->isIntegerTy(1))
+      break;
+    LLVMContext &C = RetTy->getContext();
+    EVT MTy0 = getTLI()->getValueType(DL, ICA.getArgTypes()[0]);
+    EVT MTy1 = getTLI()->getValueType(DL, ICA.getArgTypes()[1]);
----------------
sdesmalen-arm wrote:

nit:
```suggestion
    EVT VecVT = getTLI()->getValueType(DL, ICA.getArgTypes()[0]);
    EVT SubVecVT = getTLI()->getValueType(DL, ICA.getArgTypes()[1]);
```

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