[llvm] [AArch64] [DAG] combine vector_interleave to zip/zip2 (PR #231066)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: Chen Zheng (chenzheng1030)

<details>
<summary>Changes</summary>

Add a ZIP/ZIP2 combine pattern for scalable vector types.

NEON has similiar opportunities but since the pattern is not the same, so NEON pattern will be handled differently.

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


2 Files Affected:

- (modified) llvm/lib/Target/AArch64/AArch64ISelLowering.cpp (+26) 
- (added) llvm/test/CodeGen/AArch64/interleave-combine-zip.ll (+38) 


``````````diff
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index b86a0c439bfa3..c39f35659cd8c 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -22739,6 +22739,32 @@ static SDValue performConcatVectorsCombine(SDNode *N,
     return DAG.getNode(AArch64ISD::TRN1, DL, VT, Op0MoreElems, Op1MoreElems);
   }
 
+  // Interleaving the low (or high) halves of two packed vectors is a single
+  // ZIP1 (or ZIP2) of the whole vectors:
+  //   (concat_vectors (vector_interleave (extract_subvector A, I),
+  //                                      (extract_subvector B, I)))
+  //   -> (zip1 A, B) if I == 0, (zip2 A, B) if I == half the elements.
+  // Below pattern is only for scalable vectors as fixed vectors will use
+  // vector_shuffle instead of vector_interleave.
+  if (VT.isScalableVector() && N->getNumOperands() == 2 &&
+      N0Opc == ISD::VECTOR_INTERLEAVE && N0.getNode() == N1.getNode() &&
+      N0->getNumOperands() == 2 && N0.getResNo() == 0 && N1.getResNo() == 1) {
+    SDValue Lo = N0->getOperand(0), Hi = N0->getOperand(1);
+    if (Lo.getOpcode() == ISD::EXTRACT_SUBVECTOR &&
+        Hi.getOpcode() == ISD::EXTRACT_SUBVECTOR &&
+        Lo.getOperand(0).getValueType() == VT &&
+        Hi.getOperand(0).getValueType() == VT &&
+        Lo.getConstantOperandVal(1) == Hi.getConstantOperandVal(1)) {
+      uint64_t Idx = Lo.getConstantOperandVal(1);
+      if (Idx == 0)
+        return DAG.getNode(AArch64ISD::ZIP1, DL, VT, Lo.getOperand(0),
+                           Hi.getOperand(0));
+      if (Idx == VT.getVectorMinNumElements() / 2)
+        return DAG.getNode(AArch64ISD::ZIP2, DL, VT, Lo.getOperand(0),
+                           Hi.getOperand(0));
+    }
+  }
+
   if (VT.isScalableVector())
     return SDValue();
 
diff --git a/llvm/test/CodeGen/AArch64/interleave-combine-zip.ll b/llvm/test/CodeGen/AArch64/interleave-combine-zip.ll
new file mode 100644
index 0000000000000..08a148d207f0f
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/interleave-combine-zip.ll
@@ -0,0 +1,38 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc < %s -mtriple=aarch64-linux-gnu -mattr=+sve -verify-machineinstrs | FileCheck %s
+
+define <vscale x 8 x bfloat> @interleave2_lo_halves_nxv8bf16(<vscale x 8 x bfloat> %a, <vscale x 8 x bfloat> %b) {
+; CHECK-LABEL: interleave2_lo_halves_nxv8bf16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    zip1 z0.h, z0.h, z1.h
+; CHECK-NEXT:    ret
+  %x = call <vscale x 4 x bfloat> @llvm.vector.extract.nxv4bf16.nxv8bf16(<vscale x 8 x bfloat> %a, i64 0)
+  %y = call <vscale x 4 x bfloat> @llvm.vector.extract.nxv4bf16.nxv8bf16(<vscale x 8 x bfloat> %b, i64 0)
+  %r = call <vscale x 8 x bfloat> @llvm.vector.interleave2.nxv8bf16(<vscale x 4 x bfloat> %x, <vscale x 4 x bfloat> %y)
+  ret <vscale x 8 x bfloat> %r
+}
+
+define <vscale x 4 x float> @interleave2_hi_halves_nxv4f32(<vscale x 4 x float> %a, <vscale x 4 x float> %b) {
+; CHECK-LABEL: interleave2_hi_halves_nxv4f32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    zip2 z0.s, z0.s, z1.s
+; CHECK-NEXT:    ret
+  %x = call <vscale x 2 x float> @llvm.vector.extract.nxv2f32.nxv4f32(<vscale x 4 x float> %a, i64 2)
+  %y = call <vscale x 2 x float> @llvm.vector.extract.nxv2f32.nxv4f32(<vscale x 4 x float> %b, i64 2)
+  %r = call <vscale x 4 x float> @llvm.vector.interleave2.nxv4f32(<vscale x 2 x float> %x, <vscale x 2 x float> %y)
+  ret <vscale x 4 x float> %r
+}
+
+define <vscale x 8 x bfloat> @interleave2_mixed_halves_nxv8bf16(<vscale x 8 x bfloat> %a, <vscale x 8 x bfloat> %b) {
+; CHECK-LABEL: interleave2_mixed_halves_nxv8bf16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpklo z0.s, z0.h
+; CHECK-NEXT:    uunpkhi z1.s, z1.h
+; CHECK-NEXT:    trn1 z0.h, z0.h, z1.h
+; CHECK-NEXT:    ret
+  %x = call <vscale x 4 x bfloat> @llvm.vector.extract.nxv4bf16.nxv8bf16(<vscale x 8 x bfloat> %a, i64 0)
+  %y = call <vscale x 4 x bfloat> @llvm.vector.extract.nxv4bf16.nxv8bf16(<vscale x 8 x bfloat> %b, i64 4)
+  %r = call <vscale x 8 x bfloat> @llvm.vector.interleave2.nxv8bf16(<vscale x 4 x bfloat> %x, <vscale x 4 x bfloat> %y)
+  ret <vscale x 8 x bfloat> %r
+}
+

``````````

</details>


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