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

[Momchil Velikov via llvm-commits]

================
@@ -0,0 +1,395 @@
+; 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 @test_2x8bit_mask_with_32bit_index_and_trip_count(i32 %i, i32 %n) #0 {
+; CHECK-SVE-LABEL: test_2x8bit_mask_with_32bit_index_and_trip_count:
+; CHECK-SVE:       // %bb.0:
+; CHECK-SVE-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-SVE-NEXT:    whilelo p1.b, w0, w1
+; CHECK-SVE-NEXT:    punpklo p0.h, p1.b
+; CHECK-SVE-NEXT:    punpkhi p1.h, p1.b
+; CHECK-SVE-NEXT:    bl use
+; CHECK-SVE-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-SVE-NEXT:    ret
+;
+; CHECK-SVE2p1-LABEL: test_2x8bit_mask_with_32bit_index_and_trip_count:
+; CHECK-SVE2p1:       // %bb.0:
+; CHECK-SVE2p1-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; 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:    bl use
+; CHECK-SVE2p1-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-SVE2p1-NEXT:    ret
+;
+; CHECK-SME2-LABEL: test_2x8bit_mask_with_32bit_index_and_trip_count:
+; CHECK-SME2:       // %bb.0:
+; CHECK-SME2-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; 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:    bl use
+; CHECK-SME2-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; 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)
+    call void @use(<vscale x 8 x i1> %v0, <vscale x 8 x i1> %v1)
+    ret void
+}
+
+define void @test_2x8bit_mask_with_64bit_index_and_trip_count(i64 %i, i64 %n) #0 {
+; CHECK-SVE-LABEL: test_2x8bit_mask_with_64bit_index_and_trip_count:
+; CHECK-SVE:       // %bb.0:
+; CHECK-SVE-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-SVE-NEXT:    whilelo p1.b, x0, x1
+; CHECK-SVE-NEXT:    punpklo p0.h, p1.b
+; CHECK-SVE-NEXT:    punpkhi p1.h, p1.b
+; CHECK-SVE-NEXT:    bl use
+; CHECK-SVE-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-SVE-NEXT:    ret
+;
+; CHECK-SVE2p1-LABEL: test_2x8bit_mask_with_64bit_index_and_trip_count:
+; CHECK-SVE2p1:       // %bb.0:
+; CHECK-SVE2p1-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-SVE2p1-NEXT:    whilelo { p0.h, p1.h }, x0, x1
+; CHECK-SVE2p1-NEXT:    bl use
+; CHECK-SVE2p1-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-SVE2p1-NEXT:    ret
+;
+; CHECK-SME2-LABEL: test_2x8bit_mask_with_64bit_index_and_trip_count:
+; CHECK-SME2:       // %bb.0:
+; CHECK-SME2-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-SME2-NEXT:    whilelo { p0.h, p1.h }, x0, x1
+; CHECK-SME2-NEXT:    bl use
+; CHECK-SME2-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; 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)
+    call void @use(<vscale x 8 x i1> %v0, <vscale x 8 x i1> %v1)
+    ret void
+}
+
+define void @test_edge_case_2x1bit_mask(i64 %i, i64 %n) #0 {
+; CHECK-SVE-LABEL: test_edge_case_2x1bit_mask:
+; CHECK-SVE:       // %bb.0:
+; CHECK-SVE-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-SVE-NEXT:    whilelo p1.d, x0, x1
+; CHECK-SVE-NEXT:    punpklo p0.h, p1.b
+; CHECK-SVE-NEXT:    punpkhi p1.h, p1.b
+; CHECK-SVE-NEXT:    bl use
+; CHECK-SVE-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-SVE-NEXT:    ret
+;
+; CHECK-SVE2p1-LABEL: test_edge_case_2x1bit_mask:
+; CHECK-SVE2p1:       // %bb.0:
+; CHECK-SVE2p1-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-SVE2p1-NEXT:    whilelo p1.d, x0, x1
+; CHECK-SVE2p1-NEXT:    punpklo p0.h, p1.b
+; CHECK-SVE2p1-NEXT:    punpkhi p1.h, p1.b
+; CHECK-SVE2p1-NEXT:    bl use
+; CHECK-SVE2p1-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-SVE2p1-NEXT:    ret
+;
+; CHECK-SME2-LABEL: test_edge_case_2x1bit_mask:
+; CHECK-SME2:       // %bb.0:
+; CHECK-SME2-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-SME2-NEXT:    whilelo p1.d, x0, x1
+; CHECK-SME2-NEXT:    punpklo p0.h, p1.b
+; CHECK-SME2-NEXT:    punpkhi p1.h, p1.b
+; CHECK-SME2-NEXT:    bl use
+; CHECK-SME2-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-SME2-NEXT:    ret
+    %r = call <vscale x 2 x i1> @llvm.get.active.lane.mask.nxv2i1.i64(i64 %i, i64 %n)
+    %v0 = call <vscale x 1 x i1> @llvm.vector.extract.nxv1i1.nxv2i1.i64(<vscale x 2 x i1> %r, i64 0)
+    %v1 = call <vscale x 1 x i1> @llvm.vector.extract.nxv1i1.nxv2i1.i64(<vscale x 2 x i1> %r, i64 1)
+    call void @use(<vscale x 1 x i1> %v0, <vscale x 1 x i1> %v1)
+    ret void
+}
+
+define void @test_edge_case_2x16bit_mask(i64 %i, i64 %n) #0 {
----------------
momchil-velikov wrote:

Done

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