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

[Paul Walker 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 {
----------------
paulwalker-arm wrote:

Does this test offer any value?  It contains an illegal type so there's nothing in this patch that'll have any material effect.  For small tests this would be ok but this just repeats the massive sequence of instructions that's already covered by `active_lane_mask.ll`.

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