[llvm] [AArch64][CostModel] Consider some nxv1 operations as legal (PR #214471)

[Paul Walker via llvm-commits]

=?utf-8?q?Gaëtan?= Bossu <gaetan.bossu at arm.com>,
=?utf-8?q?Gaëtan?= Bossu <gaetan.bossu at arm.com>,
=?utf-8?q?Gaëtan?= Bossu <gaetan.bossu at arm.com>
Message-ID:
In-Reply-To: <llvm.org/llvm/llvm-project/pull/214471 at github.com>


================
@@ -5677,17 +5654,23 @@ InstructionCost AArch64TTIImpl::getMemoryOpCost(unsigned Opcode, Type *Ty,
   if (!LT.first.isValid())
     return InstructionCost::getInvalid();
 
-  // The code-generator is currently not able to handle scalable vectors
-  // of <vscale x 1 x eltty> yet, so return an invalid cost to avoid selecting
-  // it. This change will be removed when code-generation for these types is
-  // sufficiently reliable.
-  // We also only support full register predicate loads and stores.
-  if (auto *VTy = dyn_cast<ScalableVectorType>(Ty))
-    if (VTy->getElementCount() == ElementCount::getScalable(1) ||
-        (VTy->getElementType()->isIntegerTy(1) &&
-         !VTy->getElementCount().isKnownMultipleOf(
-             ElementCount::getScalable(16))))
+  if (auto *VTy = dyn_cast<ScalableVectorType>(Ty)) {
+    // <vscale x 1 x eltty> operations require crafting a new mask.
+    if (VTy->getElementCount() == ElementCount::getScalable(1)) {
+      Intrinsic::ID IID = Opcode == Instruction::Load ? Intrinsic::masked_load
+                                                      : Intrinsic::masked_store;
+      return getMaskedMemoryOpCost(
+                 MemIntrinsicCostAttributes(IID, Ty, Alignment, AddressSpace),
+                 CostKind) +
+             1;
+    }
+
+    // We only support full register predicate loads and stores.
+    if (VTy->getElementType()->isIntegerTy(1) &&
+        !VTy->getElementCount().isKnownMultipleOf(
+            ElementCount::getScalable(16)))
       return InstructionCost::getInvalid();
----------------
paulwalker-arm wrote:

For the avoidance of doubt, please move this before the `ElementCount::getScalable(1)` block.

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