[llvm] [AArch64][SME] Remove immediate argument restriction for svldr and svstr (PR #68565)

[Sander de Smalen via llvm-commits]

================
@@ -1741,6 +1742,54 @@ void AArch64DAGToDAGISel::SelectCVTIntrinsic(SDNode *N, unsigned NumVecs,
   CurDAG->RemoveDeadNode(N);
 }
 
+void AArch64DAGToDAGISel::SelectSMELdrStrZA(SDNode *N, bool IsLoad) {
+  // Lower an SME LDR/STR ZA intrinsic to LDR_ZA_PSEUDO or STR_ZA.
+  // If the vector select parameter is an immediate in the range 0-15 then we
+  // can emit it directly into the instruction as it's a legal operand.
+  // Otherwise we must emit 0 as the vector select operand and modify the base
+  // register instead.
+  SDLoc DL(N);
+
+  SDValue VecNum = N->getOperand(4), Base = N->getOperand(3),
+          TileSlice = N->getOperand(2);
+  int Imm = -1;
+  if (auto ImmNode = dyn_cast<ConstantSDNode>(VecNum))
+    Imm = ImmNode->getZExtValue();
+
+  if (Imm >= 0 && Imm <= 15) {
+    // 0-15 is a legal immediate so just pass it directly as a TargetConstant
+    VecNum = CurDAG->getTargetConstant(Imm, DL, MVT::i32);
+  } else {
+    // Get the vector length that will be multiplied by vnum
+    auto SVL = SDValue(
+        CurDAG->getMachineNode(AArch64::RDSVLI_XI, DL, MVT::i64,
+                               CurDAG->getTargetConstant(1, DL, MVT::i32)),
+        0);
+
+    // Multiply SVL and vnum then add it to the base register
+    if (VecNum.getValueType() == MVT::i32)
+      VecNum = Widen(CurDAG, VecNum);
+    SDValue AddOps[] = {SVL, VecNum, Base};
+    auto Add = SDValue(
+        CurDAG->getMachineNode(AArch64::MADDXrrr, DL, MVT::i64, AddOps), 0);
+
+    // The base register has been modified to take vnum into account so just
----------------
sdesmalen-arm wrote:

The (unscaled) `vnum` needs to be added the tileslice as well.

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