[llvm] [AArch64] Extend performActiveLaneMaskCombine for more than two extracts (PR #146725)

[David Sherwood via llvm-commits]

================
@@ -116,57 +175,80 @@ define void @test_partial_extract(i64 %i, i64 %n) #0 {
 define void @test_fixed_extract(i64 %i, i64 %n) #0 {
 ; CHECK-SVE-LABEL: test_fixed_extract:
 ; CHECK-SVE:       // %bb.0:
-; CHECK-SVE-NEXT:    whilelo p0.h, x0, x1
+; CHECK-SVE-NEXT:    whilelo p0.s, x0, x1
 ; CHECK-SVE-NEXT:    cset w8, mi
-; CHECK-SVE-NEXT:    mov z0.h, p0/z, #1 // =0x1
-; CHECK-SVE-NEXT:    umov w9, v0.h[4]
-; CHECK-SVE-NEXT:    umov w10, v0.h[1]
-; CHECK-SVE-NEXT:    umov w11, v0.h[5]
+; CHECK-SVE-NEXT:    mov z1.s, p0/z, #1 // =0x1
 ; CHECK-SVE-NEXT:    fmov s0, w8
-; CHECK-SVE-NEXT:    fmov s1, w9
-; CHECK-SVE-NEXT:    mov v0.s[1], w10
+; CHECK-SVE-NEXT:    mov v0.s[1], v1.s[1]
+; CHECK-SVE-NEXT:    ext z1.b, z1.b, z1.b, #8
 ; CHECK-SVE-NEXT:    // kill: def $d0 killed $d0 killed $q0
-; CHECK-SVE-NEXT:    mov v1.s[1], w11
-; CHECK-SVE-NEXT:    // kill: def $d1 killed $d1 killed $q1
+; CHECK-SVE-NEXT:    // kill: def $d1 killed $d1 killed $z1
 ; CHECK-SVE-NEXT:    b use
 ;
 ; CHECK-SVE2p1-LABEL: test_fixed_extract:
 ; CHECK-SVE2p1:       // %bb.0:
-; CHECK-SVE2p1-NEXT:    whilelo p0.h, x0, x1
+; CHECK-SVE2p1-NEXT:    whilelo p0.s, x0, x1
 ; CHECK-SVE2p1-NEXT:    cset w8, mi
-; CHECK-SVE2p1-NEXT:    mov z0.h, p0/z, #1 // =0x1
-; CHECK-SVE2p1-NEXT:    umov w9, v0.h[4]
-; CHECK-SVE2p1-NEXT:    umov w10, v0.h[1]
-; CHECK-SVE2p1-NEXT:    umov w11, v0.h[5]
+; CHECK-SVE2p1-NEXT:    mov z1.s, p0/z, #1 // =0x1
 ; CHECK-SVE2p1-NEXT:    fmov s0, w8
-; CHECK-SVE2p1-NEXT:    fmov s1, w9
-; CHECK-SVE2p1-NEXT:    mov v0.s[1], w10
+; CHECK-SVE2p1-NEXT:    mov v0.s[1], v1.s[1]
+; CHECK-SVE2p1-NEXT:    ext z1.b, z1.b, z1.b, #8
 ; CHECK-SVE2p1-NEXT:    // kill: def $d0 killed $d0 killed $q0
-; CHECK-SVE2p1-NEXT:    mov v1.s[1], w11
-; CHECK-SVE2p1-NEXT:    // kill: def $d1 killed $d1 killed $q1
+; CHECK-SVE2p1-NEXT:    // kill: def $d1 killed $d1 killed $z1
 ; CHECK-SVE2p1-NEXT:    b use
 ;
 ; CHECK-SME2-LABEL: test_fixed_extract:
 ; CHECK-SME2:       // %bb.0:
-; CHECK-SME2-NEXT:    whilelo p0.h, x0, x1
+; CHECK-SME2-NEXT:    whilelo p0.s, x0, x1
 ; CHECK-SME2-NEXT:    cset w8, mi
-; CHECK-SME2-NEXT:    mov z0.h, p0/z, #1 // =0x1
-; CHECK-SME2-NEXT:    mov z1.h, z0.h[1]
-; CHECK-SME2-NEXT:    mov z2.h, z0.h[5]
-; CHECK-SME2-NEXT:    mov z3.h, z0.h[4]
-; CHECK-SME2-NEXT:    fmov s0, w8
-; CHECK-SME2-NEXT:    zip1 z0.s, z0.s, z1.s
-; CHECK-SME2-NEXT:    zip1 z1.s, z3.s, z2.s
-; CHECK-SME2-NEXT:    // kill: def $d0 killed $d0 killed $z0
+; CHECK-SME2-NEXT:    mov z1.s, p0/z, #1 // =0x1
+; CHECK-SME2-NEXT:    fmov s2, w8
+; CHECK-SME2-NEXT:    mov z0.s, z1.s[1]
+; CHECK-SME2-NEXT:    ext z1.b, z1.b, z1.b, #8
 ; CHECK-SME2-NEXT:    // kill: def $d1 killed $d1 killed $z1
+; CHECK-SME2-NEXT:    zip1 z0.s, z2.s, z0.s
+; CHECK-SME2-NEXT:    // kill: def $d0 killed $d0 killed $z0
 ; CHECK-SME2-NEXT:    b use
-    %r = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %i, i64 %n)
-    %v0 = call <2 x i1> @llvm.vector.extract.v2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 0)
-    %v1 = call <2 x i1> @llvm.vector.extract.v2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 4)
+    %r = call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i64(i64 %i, i64 %n)
+    %v0 = call <2 x i1> @llvm.vector.extract.v2i1.nxv8i1.i64(<vscale x 4 x i1> %r, i64 0)
----------------
david-arm wrote:

nit: I think these should be `@llvm.vector.extract.v2i1.nxv4i1.i64`

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