[llvm] [AArch64] Combine getActiveLaneMask with vector_extract (PR #81139)

[David Sherwood via llvm-commits]

================
@@ -0,0 +1,90 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc -mattr=+sve    < %s | FileCheck %s -check-prefix CHECK-SVE
+; RUN: llc -mattr=+sve2p1 < %s | FileCheck %s -check-prefix CHECK-SVE2p1
+; RUN: llc -mattr=+sme2   < %s | FileCheck %s -check-prefix CHECK-SME2
+target triple = "aarch64-linux"
+
+; Test combining of getActiveLaneMask with a pair of extract_vector operations.
+
+define void @f0(i32 %i, i32 %n, ptr %p0, ptr %p1) #0 {
+; CHECK-SVE-LABEL: f0:
+; CHECK-SVE:       // %bb.0:
+; CHECK-SVE-NEXT:    whilelo p0.b, w0, w1
+; CHECK-SVE-NEXT:    ptrue p1.h
+; CHECK-SVE-NEXT:    punpklo p2.h, p0.b
+; CHECK-SVE-NEXT:    punpkhi p0.h, p0.b
+; CHECK-SVE-NEXT:    and p2.b, p2/z, p2.b, p1.b
+; CHECK-SVE-NEXT:    and p0.b, p0/z, p0.b, p1.b
+; CHECK-SVE-NEXT:    str p2, [x2]
+; CHECK-SVE-NEXT:    str p0, [x3]
+; CHECK-SVE-NEXT:    ret
+;
+; CHECK-SVE2p1-LABEL: f0:
+; CHECK-SVE2p1:       // %bb.0:
+; CHECK-SVE2p1-NEXT:    mov w8, w1
+; CHECK-SVE2p1-NEXT:    mov w9, w0
+; CHECK-SVE2p1-NEXT:    whilelo { p0.h, p1.h }, x9, x8
+; CHECK-SVE2p1-NEXT:    str p0, [x2]
+; CHECK-SVE2p1-NEXT:    str p1, [x3]
+; CHECK-SVE2p1-NEXT:    ret
+;
+; CHECK-SME2-LABEL: f0:
+; CHECK-SME2:       // %bb.0:
+; CHECK-SME2-NEXT:    mov w8, w1
+; CHECK-SME2-NEXT:    mov w9, w0
+; CHECK-SME2-NEXT:    whilelo { p0.h, p1.h }, x9, x8
+; CHECK-SME2-NEXT:    str p0, [x2]
+; CHECK-SME2-NEXT:    str p1, [x3]
+; CHECK-SME2-NEXT:    ret
+    %r = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i32(i32 %i, i32 %n)
+    %v0 = call <vscale x 8 x i1> @llvm.vector.extract.nxv8i1.nxv16i1.i64(<vscale x 16 x i1> %r, i64 0)
+    %v1 = call <vscale x 8 x i1> @llvm.vector.extract.nxv8i1.nxv16i1.i64(<vscale x 16 x i1> %r, i64 8)
+    %pg0 = call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv8i1(<vscale x 8 x i1> %v0)
+    %pg1 = call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv8i1(<vscale x 8 x i1> %v1)
+    store <vscale x 16 x i1> %pg0, ptr %p0
+    store <vscale x 16 x i1> %pg1, ptr %p1
+    ret void
+}
+
+define void @f1(i64 %i, i64 %n, ptr %p0, ptr %p1) #0 {
+; CHECK-SVE-LABEL: f1:
+; CHECK-SVE:       // %bb.0:
+; CHECK-SVE-NEXT:    whilelo p0.b, x0, x1
+; CHECK-SVE-NEXT:    ptrue p1.h
+; CHECK-SVE-NEXT:    punpklo p2.h, p0.b
+; CHECK-SVE-NEXT:    punpkhi p0.h, p0.b
+; CHECK-SVE-NEXT:    and p2.b, p2/z, p2.b, p1.b
+; CHECK-SVE-NEXT:    and p0.b, p0/z, p0.b, p1.b
+; CHECK-SVE-NEXT:    str p2, [x2]
+; CHECK-SVE-NEXT:    str p0, [x3]
+; CHECK-SVE-NEXT:    ret
+;
+; CHECK-SVE2p1-LABEL: f1:
+; CHECK-SVE2p1:       // %bb.0:
+; CHECK-SVE2p1-NEXT:    whilelo { p0.h, p1.h }, x0, x1
+; CHECK-SVE2p1-NEXT:    str p0, [x2]
+; CHECK-SVE2p1-NEXT:    str p1, [x3]
+; CHECK-SVE2p1-NEXT:    ret
+;
+; CHECK-SME2-LABEL: f1:
+; CHECK-SME2:       // %bb.0:
+; CHECK-SME2-NEXT:    whilelo { p0.h, p1.h }, x0, x1
+; CHECK-SME2-NEXT:    str p0, [x2]
+; CHECK-SME2-NEXT:    str p1, [x3]
+; CHECK-SME2-NEXT:    ret
+    %r = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 %i, i64 %n)
+    %v0 = call <vscale x 8 x i1> @llvm.vector.extract.nxv8i1.nxv16i1.i64(<vscale x 16 x i1> %r, i64 0)
+    %v1 = call <vscale x 8 x i1> @llvm.vector.extract.nxv8i1.nxv16i1.i64(<vscale x 16 x i1> %r, i64 8)
+    %pg0 = call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv8i1(<vscale x 8 x i1> %v0)
+    %pg1 = call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv8i1(<vscale x 8 x i1> %v1)
+    store <vscale x 16 x i1> %pg0, ptr %p0
+    store <vscale x 16 x i1> %pg1, ptr %p1
+    ret void
+}
+
----------------
david-arm wrote:

I think it would be good to have some more tests, e.g. at least one other predicate type such as nxv4i1.

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