[AArch64] Extend usage of XAR instruction for fixed-length operations

[Rajveer100]

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.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Rajveer Singh Bharadwaj (Rajveer100)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/139460.diff

1 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp (+22-2)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
index 96fa85179d023..bb059928e33a3 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -4632,18 +4632,38 @@ 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;
 
+  bool UseSVE2Instr = false;
   if (!IsXOROperand) {
+    if (VT.getVectorElementType() != MVT::i64 && Subtarget->hasSVE2())
+      UseSVE2Instr = true;
+
     SDValue Zero = CurDAG->getTargetConstant(0, DL, MVT::i64);
     SDNode *MOV = CurDAG->getMachineNode(AArch64::MOVIv2d_ns, DL, VT, Zero);
     SDValue MOVIV = SDValue(MOV, 0);
+
     R1 = N1->getOperand(0);
-    R2 = MOVIV;
+    if (UseSVE2Instr) {
+      SDValue ZSub = CurDAG->getTargetConstant(AArch64::zsub, DL, MVT::i32);
+      SDNode *SubRegToReg = CurDAG->getMachineNode(AArch64::SUBREG_TO_REG, DL,
+                                                   VT, Zero, MOVIV, ZSub);
+      R2 = SDValue(SubRegToReg, 0);
+    } else {
+      R2 = MOVIV;
+    }
   }
 
   SDValue Ops[] = {R1, R2, Imm};
+  if (UseSVE2Instr) {
+    if (auto Opc = SelectOpcodeFromVT<SelectTypeKind::Int>(
+            VT, {AArch64::XAR_ZZZI_B, AArch64::XAR_ZZZI_H, AArch64::XAR_ZZZI_S,
+                 AArch64::XAR_ZZZI_D})) {
+      CurDAG->SelectNodeTo(N, Opc, VT, Ops);
+      return true;
+    }
+  }
   CurDAG->SelectNodeTo(N, AArch64::XAR, N0.getValueType(), Ops);
 
   return true;

[Rajveer100]

@davemgreen
Let me know if this is in the right direction. Also, I am probably not using the right VT here causing an assertion.

[davemgreen]

I think it will need a INSERT_SUBREG IMPLICIT_DEF, A, zsub for the input and a EXTRACT_SUBREG xar, zsub to make sure the result is kept as the right type for the result.

It can apply to both the rotr(xor(a, b))->xar(a,b) and the rotr(a)->xar(a,0) versions (so it might be easier to expand R1 and R2.

[Rajveer100]

I have pushed changes, let me know if this was the intended direction.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[Rajveer100]

@davemgreen
Everything works well now.

[Rajveer100]

Pushed changes :)

[davemgreen]

Thanks - nice work. This looks good (it gets a bit fiddly with all the combos), but I think the smaller types need to be expanding to vectors with the same element sizes to make sure the rotates operate correctly.

[Rajveer100]

Is it possible to convert v2i32 to nxv4i32 in one step, or do we need to first convert it to v4i32 (same for other types as well)?

[Rajveer100]

Yep, I was right! All working well now.

Edit: Code clean up is what remains :)

Edit 2: Clean up done.

[davemgreen]

Is it possible to convert v2i32 to nxv4i32 in one step, or do we need to first convert it to v4i32 (same for other types as well)?

I was wondering if that would work or not. I think it might be possibly, but going through two steps like you have sounds fine to me.

[davemgreen]

LGTM if the type of the MOVIv2d_ns is changed to v2i64. Thanks

[Rajveer100]

I guess you can go ahead and merge it for me, since it might take some more time!

Edit: Got access :)