[llvm] [AArch64] Extend usage of `XAR` instruction for fixed-length operations (PR #139460)

[Rajveer Singh Bharadwaj via llvm-commits]

https://github.com/Rajveer100 updated https://github.com/llvm/llvm-project/pull/139460

>From 60a2ff0e9e2a09aa23f836398613266244e2ff64 Mon Sep 17 00:00:00 2001
From: Rajveer <rajveer.developer at icloud.com>
Date: Sun, 11 May 2025 22:27:50 +0530
Subject: [PATCH] [AArch64] Extend usage of `XAR` instruction for fixed-length
 operations

Resolves #139229

In #137162, support for `v2i64` was implemented for vector rotate
transformation, although types like `v4i32`, `v8i16` and `v16i8`
do not have Neon SHA3, we can use SVE operations if sve2-sha3
is available.
---
 .../Target/AArch64/AArch64ISelDAGToDAG.cpp    | 27 ++++++++++++++-----
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
index 34f6db9374cb5..2dd2b306a3497 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -4632,20 +4632,31 @@ bool AArch64DAGToDAGISel::trySelectXAR(SDNode *N) {
   SDValue Imm = CurDAG->getTargetConstant(
       ShAmt, DL, N0.getOperand(1).getValueType(), false);
 
-  if (ShAmt + HsAmt != 64)
+  if (ShAmt + HsAmt != VT.getScalarSizeInBits())
     return false;
 
   if (!IsXOROperand) {
     SDValue Zero = CurDAG->getTargetConstant(0, DL, MVT::i64);
     SDNode *MOV =
-        CurDAG->getMachineNode(AArch64::MOVIv2d_ns, DL, MVT::v2i64, Zero);
+        CurDAG->getMachineNode(AArch64::MOVIv2d_ns, DL, MVT::nxv4i32, Zero);
     SDValue MOVIV = SDValue(MOV, 0);
-    R1 = N1->getOperand(0);
-    R2 = MOVIV;
+
+    SDValue Undef = SDValue(
+        CurDAG->getMachineNode(TargetOpcode::IMPLICIT_DEF, DL, MVT::nxv4i32),
+        0);
+    SDValue ZSub = CurDAG->getTargetConstant(AArch64::zsub, DL, MVT::i32);
+
+    R1 =
+        SDValue(CurDAG->getMachineNode(AArch64::INSERT_SUBREG, DL, MVT::nxv4i32,
+                                       Undef, N1->getOperand(0), ZSub),
+                0);
+    R2 = SDValue(CurDAG->getMachineNode(AArch64::INSERT_SUBREG, DL,
+                                        MVT::nxv4i32, Undef, MOVIV, ZSub),
+                 0);
   }
 
   // If the input is a v1i64, widen to a v2i64 to use XAR.
-  assert((VT == MVT::v1i64 || VT == MVT::v2i64) && "Unexpected XAR type!");
+  // assert((VT == MVT::v1i64 || VT == MVT::v2i64) && "Unexpected XAR type!");
   if (VT == MVT::v1i64) {
     EVT SVT = MVT::v2i64;
     SDValue Undef =
@@ -4661,12 +4672,16 @@ bool AArch64DAGToDAGISel::trySelectXAR(SDNode *N) {
   }
 
   SDValue Ops[] = {R1, R2, Imm};
-  SDNode *XAR = CurDAG->getMachineNode(AArch64::XAR, DL, MVT::v2i64, Ops);
+  SDNode *XAR = CurDAG->getMachineNode(AArch64::XAR_ZZZI_S, DL, VT, Ops);
 
   if (VT == MVT::v1i64) {
     SDValue DSub = CurDAG->getTargetConstant(AArch64::dsub, DL, MVT::i32);
     XAR = CurDAG->getMachineNode(AArch64::EXTRACT_SUBREG, DL, VT,
                                  SDValue(XAR, 0), DSub);
+  } else {
+    SDValue ZSub = CurDAG->getTargetConstant(AArch64::zsub, DL, MVT::i32);
+    XAR = CurDAG->getMachineNode(AArch64::EXTRACT_SUBREG, DL, VT,
+                                 SDValue(XAR, 0), ZSub);
   }
   ReplaceNode(N, XAR);
   return true;