test/CodeGen/AArch64/get-active-lane-mask-extract.ll

; 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-SME2 -check-prefix CHECK-SVE2p1
; RUN: llc -mattr=+sve -mattr=+sme2 -force-streaming < %s | FileCheck %s -check-prefix CHECK-SVE2p1-SME2 -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:    whilelo p1.b, w0, w1
; CHECK-SVE-NEXT:    punpklo p0.h, p1.b
; CHECK-SVE-NEXT:    punpkhi p1.h, p1.b
; CHECK-SVE-NEXT:    b use
;
; CHECK-SVE2p1-SME2-LABEL: test_2x8bit_mask_with_32bit_index_and_trip_count:
; CHECK-SVE2p1-SME2:       // %bb.0:
; CHECK-SVE2p1-SME2-NEXT:    mov w8, w1
; CHECK-SVE2p1-SME2-NEXT:    mov w9, w0
; CHECK-SVE2p1-SME2-NEXT:    whilelo { p0.h, p1.h }, x9, x8
; CHECK-SVE2p1-SME2-NEXT:    b use
    %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)
    tail 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:    whilelo p1.b, x0, x1
; CHECK-SVE-NEXT:    punpklo p0.h, p1.b
; CHECK-SVE-NEXT:    punpkhi p1.h, p1.b
; CHECK-SVE-NEXT:    b use
;
; CHECK-SVE2p1-SME2-LABEL: test_2x8bit_mask_with_64bit_index_and_trip_count:
; CHECK-SVE2p1-SME2:       // %bb.0:
; CHECK-SVE2p1-SME2-NEXT:    whilelo { p0.h, p1.h }, x0, x1
; CHECK-SVE2p1-SME2-NEXT:    b use
    %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)
    tail 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:    whilelo p1.d, x0, x1
; CHECK-SVE-NEXT:    punpklo p0.h, p1.b
; CHECK-SVE-NEXT:    punpkhi p1.h, p1.b
; CHECK-SVE-NEXT:    b use
;
; CHECK-SVE2p1-SME2-LABEL: test_edge_case_2x1bit_mask:
; CHECK-SVE2p1-SME2:       // %bb.0:
; CHECK-SVE2p1-SME2-NEXT:    whilelo p1.d, x0, x1
; CHECK-SVE2p1-SME2-NEXT:    punpklo p0.h, p1.b
; CHECK-SVE2p1-SME2-NEXT:    punpkhi p1.h, p1.b
; CHECK-SVE2p1-SME2-NEXT:    b use
    %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)
    tail call void @use(<vscale x 1 x i1> %v0, <vscale x 1 x i1> %v1)
    ret void
}

define void @test_boring_case_2x2bit_mask(i64 %i, i64 %n) #0 {
; CHECK-SVE-LABEL: test_boring_case_2x2bit_mask:
; CHECK-SVE:       // %bb.0:
; CHECK-SVE-NEXT:    whilelo p1.s, x0, x1
; CHECK-SVE-NEXT:    punpklo p0.h, p1.b
; CHECK-SVE-NEXT:    punpkhi p1.h, p1.b
; CHECK-SVE-NEXT:    b use
;
; CHECK-SVE2p1-SME2-LABEL: test_boring_case_2x2bit_mask:
; CHECK-SVE2p1-SME2:       // %bb.0:
; CHECK-SVE2p1-SME2-NEXT:    whilelo { p0.d, p1.d }, x0, x1
; CHECK-SVE2p1-SME2-NEXT:    b use
    %r = call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i64(i64 %i, i64 %n)
    %v0 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv4i1.i64(<vscale x 4 x i1> %r, i64 0)
    %v1 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv4i1.i64(<vscale x 4 x i1> %r, i64 2)
    tail call void @use(<vscale x 2 x i1> %v0, <vscale x 2 x i1> %v1)
    ret void
}

define void @test_legal_4x2bit_mask(i64 %i, i64 %n) #0 {
; CHECK-SVE-LABEL: test_legal_4x2bit_mask:
; CHECK-SVE:       // %bb.0:
; CHECK-SVE-NEXT:    whilelo p0.h, x0, x1
; CHECK-SVE-NEXT:    punpkhi p1.h, p0.b
; CHECK-SVE-NEXT:    punpklo p4.h, p0.b
; CHECK-SVE-NEXT:    punpkhi p3.h, p1.b
; CHECK-SVE-NEXT:    punpklo p2.h, p1.b
; CHECK-SVE-NEXT:    punpklo p0.h, p4.b
; CHECK-SVE-NEXT:    punpkhi p1.h, p4.b
; CHECK-SVE-NEXT:    b use
;
; CHECK-SVE2p1-SME2-LABEL: test_legal_4x2bit_mask:
; CHECK-SVE2p1-SME2:       // %bb.0:
; CHECK-SVE2p1-SME2-NEXT:    cntw x8
; CHECK-SVE2p1-SME2-NEXT:    adds x8, x0, x8
; CHECK-SVE2p1-SME2-NEXT:    csinv x8, x8, xzr, lo
; CHECK-SVE2p1-SME2-NEXT:    whilelo { p0.d, p1.d }, x0, x1
; CHECK-SVE2p1-SME2-NEXT:    whilelo { p2.d, p3.d }, x8, x1
; CHECK-SVE2p1-SME2-NEXT:    b use
  %r = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %i, i64 %n)
  %v0 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 6)
  %v1 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 4)
  %v2 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 2)
  %v3 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 0)
  tail call void @use(<vscale x 2 x i1> %v3, <vscale x 2 x i1> %v2, <vscale x 2 x i1> %v1, <vscale x 2 x i1> %v0)
  ret void
}

; Negative test where the extract types are correct but we are not extracting all parts of the mask
; Note: We could still create a whilelo_x2 for the first two extracts, but we don't expect this case often yet.
define void @test_partial_extract_correct_types(i64 %i, i64 %n) #0 {
; CHECK-SVE-LABEL: test_partial_extract_correct_types:
; CHECK-SVE:       // %bb.0:
; CHECK-SVE-NEXT:    whilelo p0.h, x0, x1
; CHECK-SVE-NEXT:    punpklo p1.h, p0.b
; CHECK-SVE-NEXT:    punpkhi p2.h, p0.b
; CHECK-SVE-NEXT:    punpklo p0.h, p1.b
; CHECK-SVE-NEXT:    punpkhi p1.h, p1.b
; CHECK-SVE-NEXT:    punpkhi p2.h, p2.b
; CHECK-SVE-NEXT:    b use
;
; CHECK-SVE2p1-SME2-LABEL: test_partial_extract_correct_types:
; CHECK-SVE2p1-SME2:       // %bb.0:
; CHECK-SVE2p1-SME2-NEXT:    whilelo p0.h, x0, x1
; CHECK-SVE2p1-SME2-NEXT:    punpklo p1.h, p0.b
; CHECK-SVE2p1-SME2-NEXT:    punpkhi p2.h, p0.b
; CHECK-SVE2p1-SME2-NEXT:    punpklo p0.h, p1.b
; CHECK-SVE2p1-SME2-NEXT:    punpkhi p1.h, p1.b
; CHECK-SVE2p1-SME2-NEXT:    punpkhi p2.h, p2.b
; CHECK-SVE2p1-SME2-NEXT:    b use
  %r = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %i, i64 %n)
  %v0 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 0)
  %v1 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 2)
  %v2 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 6)
  tail call void @use(<vscale x 2 x i1> %v0, <vscale x 2 x i1> %v1, <vscale x 2 x i1> %v2)
  ret void
}

; Negative test for when not extracting exactly two halves of the source vector
define void @test_partial_extract(i64 %i, i64 %n) #0 {
; CHECK-SVE-LABEL: test_partial_extract:
; CHECK-SVE:       // %bb.0:
; CHECK-SVE-NEXT:    whilelo p0.h, x0, x1
; CHECK-SVE-NEXT:    punpklo p1.h, p0.b
; CHECK-SVE-NEXT:    punpkhi p2.h, p0.b
; CHECK-SVE-NEXT:    punpklo p0.h, p1.b
; CHECK-SVE-NEXT:    punpklo p1.h, p2.b
; CHECK-SVE-NEXT:    b use
;
; CHECK-SVE2p1-SME2-LABEL: test_partial_extract:
; CHECK-SVE2p1-SME2:       // %bb.0:
; CHECK-SVE2p1-SME2-NEXT:    whilelo p0.h, x0, x1
; CHECK-SVE2p1-SME2-NEXT:    punpklo p1.h, p0.b
; CHECK-SVE2p1-SME2-NEXT:    punpkhi p2.h, p0.b
; CHECK-SVE2p1-SME2-NEXT:    punpklo p0.h, p1.b
; CHECK-SVE2p1-SME2-NEXT:    punpklo p1.h, p2.b
; CHECK-SVE2p1-SME2-NEXT:    b use
    %r = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %i, i64 %n)
    %v0 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 0)
    %v1 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 4)
    tail call void @use(<vscale x 2 x i1> %v0, <vscale x 2 x i1> %v1)
    ret void
}

; Negative test for when extracting a fixed-length vector.
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.s, x0, x1
; CHECK-SVE-NEXT:    cset w8, mi
; CHECK-SVE-NEXT:    mov z1.s, p0/z, #1 // =0x1
; CHECK-SVE-NEXT:    fmov s0, w8
; 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:    // 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.s, x0, x1
; CHECK-SVE2p1-NEXT:    cset w8, mi
; CHECK-SVE2p1-NEXT:    mov z1.s, p0/z, #1 // =0x1
; CHECK-SVE2p1-NEXT:    fmov s0, w8
; 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:    // 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.s, x0, x1
; CHECK-SME2-NEXT:    cset w8, mi
; 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 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i64(i64 %i, i64 %n)
    %v0 = call <2 x i1> @llvm.vector.extract.v2i1.nxv4i1.i64(<vscale x 4 x i1> %r, i64 0)
    %v1 = call <2 x i1> @llvm.vector.extract.v2i1.nxv4i1.i64(<vscale x 4 x i1> %r, i64 2)
    tail call void @use(<2 x i1> %v0, <2 x i1> %v1)
    ret void
}

; Negative test where the number of extracts is right, but they cannot be combined because
; there is not an extract for each part
define void @test_4x2bit_duplicate_mask(i64 %i, i64 %n) #0 {
; CHECK-SVE-LABEL: test_4x2bit_duplicate_mask:
; CHECK-SVE:       // %bb.0:
; CHECK-SVE-NEXT:    whilelo p0.h, x0, x1
; CHECK-SVE-NEXT:    punpklo p1.h, p0.b
; CHECK-SVE-NEXT:    punpkhi p3.h, p0.b
; CHECK-SVE-NEXT:    punpkhi p0.h, p1.b
; CHECK-SVE-NEXT:    punpklo p2.h, p3.b
; CHECK-SVE-NEXT:    punpkhi p3.h, p3.b
; CHECK-SVE-NEXT:    mov p1.b, p0.b
; CHECK-SVE-NEXT:    b use
;
; CHECK-SVE2p1-SME2-LABEL: test_4x2bit_duplicate_mask:
; CHECK-SVE2p1-SME2:       // %bb.0:
; CHECK-SVE2p1-SME2-NEXT:    whilelo p0.h, x0, x1
; CHECK-SVE2p1-SME2-NEXT:    punpklo p1.h, p0.b
; CHECK-SVE2p1-SME2-NEXT:    punpkhi p3.h, p0.b
; CHECK-SVE2p1-SME2-NEXT:    punpkhi p0.h, p1.b
; CHECK-SVE2p1-SME2-NEXT:    punpklo p2.h, p3.b
; CHECK-SVE2p1-SME2-NEXT:    punpkhi p3.h, p3.b
; CHECK-SVE2p1-SME2-NEXT:    mov p1.b, p0.b
; CHECK-SVE2p1-SME2-NEXT:    b use
  %r = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %i, i64 %n)
  %v0 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 2)
  %v1 = call <vscale x 2 x i1> @llvm.vector.extract.nxv2i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 2)
  %v2 = call <vscale x 2 x i1> @llvm.vector.extract.nxv4i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 4)
  %v3 = call <vscale x 2 x i1> @llvm.vector.extract.nxv4i1.nxv8i1.i64(<vscale x 8 x i1> %r, i64 6)
  tail call void @use(<vscale x 2 x i1> %v0, <vscale x 2 x i1> %v1, <vscale x 2 x i1> %v2, <vscale x 2 x i1> %v3)
  ret void
}

; Illegal Types

define void @test_2x16bit_mask_with_32bit_index_and_trip_count(i32 %i, i32 %n) #0 {
; CHECK-SVE-LABEL: test_2x16bit_mask_with_32bit_index_and_trip_count:
; CHECK-SVE:       // %bb.0:
; CHECK-SVE-NEXT:    rdvl x8, #1
; CHECK-SVE-NEXT:    adds w8, w0, w8
; CHECK-SVE-NEXT:    csinv w8, w8, wzr, lo
; CHECK-SVE-NEXT:    whilelo p0.b, w0, w1
; CHECK-SVE-NEXT:    whilelo p1.b, w8, w1
; CHECK-SVE-NEXT:    b use
;
; CHECK-SVE2p1-SME2-LABEL: test_2x16bit_mask_with_32bit_index_and_trip_count:
; CHECK-SVE2p1-SME2:       // %bb.0:
; CHECK-SVE2p1-SME2-NEXT:    mov w8, w1
; CHECK-SVE2p1-SME2-NEXT:    mov w9, w0
; CHECK-SVE2p1-SME2-NEXT:    whilelo { p0.b, p1.b }, x9, x8
; CHECK-SVE2p1-SME2-NEXT:    b use
  %r = call <vscale x 32 x i1> @llvm.get.active.lane.mask.nxv32i1.i32(i32 %i, i32 %n)
  %v0 = call <vscale x 16 x i1> @llvm.vector.extract.nxv16i1.nxv32i1.i64(<vscale x 32 x i1> %r, i64 0)
  %v1 = call <vscale x 16 x i1> @llvm.vector.extract.nxv16i1.nxv32i1.i64(<vscale x 32 x i1> %r, i64 16)
  tail call void @use(<vscale x 16 x i1> %v0, <vscale x 16 x i1> %v1)
  ret void
}

define void @test_2x32bit_mask_with_32bit_index_and_trip_count(i32 %i, i32 %n) #0 {
; CHECK-SVE-LABEL: test_2x32bit_mask_with_32bit_index_and_trip_count:
; CHECK-SVE:       // %bb.0:
; CHECK-SVE-NEXT:    rdvl x8, #2
; CHECK-SVE-NEXT:    rdvl x9, #1
; CHECK-SVE-NEXT:    adds w8, w0, w8
; CHECK-SVE-NEXT:    csinv w8, w8, wzr, lo
; CHECK-SVE-NEXT:    adds w10, w8, w9
; CHECK-SVE-NEXT:    csinv w10, w10, wzr, lo
; CHECK-SVE-NEXT:    whilelo p3.b, w10, w1
; CHECK-SVE-NEXT:    adds w9, w0, w9
; CHECK-SVE-NEXT:    csinv w9, w9, wzr, lo
; CHECK-SVE-NEXT:    whilelo p0.b, w0, w1
; CHECK-SVE-NEXT:    whilelo p1.b, w9, w1
; CHECK-SVE-NEXT:    whilelo p2.b, w8, w1
; CHECK-SVE-NEXT:    b use
;
; CHECK-SVE2p1-SME2-LABEL: test_2x32bit_mask_with_32bit_index_and_trip_count:
; CHECK-SVE2p1-SME2:       // %bb.0:
; CHECK-SVE2p1-SME2-NEXT:    rdvl x8, #2
; CHECK-SVE2p1-SME2-NEXT:    mov w9, w1
; CHECK-SVE2p1-SME2-NEXT:    mov w10, w0
; CHECK-SVE2p1-SME2-NEXT:    adds w8, w0, w8
; CHECK-SVE2p1-SME2-NEXT:    csinv w8, w8, wzr, lo
; CHECK-SVE2p1-SME2-NEXT:    whilelo { p0.b, p1.b }, x10, x9
; CHECK-SVE2p1-SME2-NEXT:    whilelo { p2.b, p3.b }, x8, x9
; CHECK-SVE2p1-SME2-NEXT:    b use
  %r = call <vscale x 64 x i1> @llvm.get.active.lane.mask.nxv64i1.i32(i32 %i, i32 %n)
  %v0 = call <vscale x 16 x i1> @llvm.vector.extract.nxv16i1.nxv64i1.i64(<vscale x 64 x i1> %r, i64 0)
  %v1 = call <vscale x 16 x i1> @llvm.vector.extract.nxv16i1.nxv64i1.i64(<vscale x 64 x i1> %r, i64 16)
  %v2 = call <vscale x 16 x i1> @llvm.vector.extract.nxv16i1.nxv64i1.i64(<vscale x 64 x i1> %r, i64 32)
  %v3 = call <vscale x 16 x i1> @llvm.vector.extract.nxv16i1.nxv64i1.i64(<vscale x 64 x i1> %r, i64 48)
  tail call void @use(<vscale x 16 x i1> %v0, <vscale x 16 x i1> %v1, <vscale x 16 x i1> %v2, <vscale x 16 x i1> %v3)
  ret void
}

declare void @use(...)

attributes #0 = { nounwind }