llvm/test/CodeGen/AArch64/sve2p1-while-pn-folds.ll

; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
; RUN: llc -mtriple=aarch64-linux-unknown -mattr=+sve2p1 -o - < %s | FileCheck %s -check-prefixes=CHECK,CHECK-SVE
; RUN: llc -mtriple=aarch64-linux-unknown -mattr=+sme2 -force-streaming -o - < %s | FileCheck %s -check-prefixes=CHECK,CHECK-SME

; Tests extracting the first lane from the first segment of a
; predicate-as-counter while folds to a conditional set (CSET) based on the
; "FIRST_ACTIVE" status flags from the while.
;
; %while, %flags = WHILE_*_PRED_COUNTER .. ; predicate-as-counter while
; %first.segment = pext(%while, 0) ; predicate extract of segment 0
; %first.active = extract_elt(%first_segment, 0) ; extract first lane
;
; ->
;
; %while, %flags = WHILE_*_PRED_COUNTER .. ; predicate-as-counter while
; %first.active = cset(%flags, FIRST_ACTIVE)

; Note: FIRST_ACTIVE corresponds to the MI condition code (N == 1), and
; not(FIRST_ACTIVE) corresponds to the PL condition code (N == 0).

define i1 @whilege_first_active(i64 %a, i64 %b) {
; CHECK-LABEL: whilege_first_active:
; CHECK:       // %bb.0:
; CHECK-NEXT:    whilege pn8.b, x0, x1, vlx4
; CHECK-NEXT:    cset w0, mi
; CHECK-NEXT:    ret
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilege.c8(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 16 x i1> @llvm.aarch64.sve.pext.nxv16i1(target("aarch64.svcount") %while, i32 0)
  %bit = extractelement <vscale x 16 x i1> %pext, i64 0
  ret i1 %bit
}

define i1 @whilegt_first_active(i64 %a, i64 %b) {
; CHECK-LABEL: whilegt_first_active:
; CHECK:       // %bb.0:
; CHECK-NEXT:    whilegt pn8.h, x0, x1, vlx4
; CHECK-NEXT:    cset w0, mi
; CHECK-NEXT:    ret
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilegt.c16(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 8 x i1> @llvm.aarch64.sve.pext.nxv8i1(target("aarch64.svcount") %while, i32 0)
  %bit = extractelement <vscale x 8 x i1> %pext, i64 0
  ret i1 %bit
}

define i1 @whilelt_first_active(i64 %a, i64 %b) {
; CHECK-LABEL: whilelt_first_active:
; CHECK:       // %bb.0:
; CHECK-NEXT:    whilelt pn8.s, x0, x1, vlx4
; CHECK-NEXT:    cset w0, mi
; CHECK-NEXT:    ret
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilelt.c32(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 4 x i1> @llvm.aarch64.sve.pext.nxv4i1(target("aarch64.svcount") %while, i32 0)
  %bit = extractelement <vscale x 4 x i1> %pext, i64 0
  ret i1 %bit
}

define i1 @whilele_first_active(i64 %a, i64 %b) {
; CHECK-LABEL: whilele_first_active:
; CHECK:       // %bb.0:
; CHECK-NEXT:    whilele pn8.d, x0, x1, vlx4
; CHECK-NEXT:    cset w0, mi
; CHECK-NEXT:    ret
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilele.c64(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 2 x i1> @llvm.aarch64.sve.pext.nxv2i1(target("aarch64.svcount") %while, i32 0)
  %bit = extractelement <vscale x 2 x i1> %pext, i64 0
  ret i1 %bit
}

define i1 @whilehs_first_active(i64 %a, i64 %b) {
; CHECK-LABEL: whilehs_first_active:
; CHECK:       // %bb.0:
; CHECK-NEXT:    whilehs pn8.b, x0, x1, vlx4
; CHECK-NEXT:    cset w0, mi
; CHECK-NEXT:    ret
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilehs.c8(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 16 x i1> @llvm.aarch64.sve.pext.nxv16i1(target("aarch64.svcount") %while, i32 0)
  %bit = extractelement <vscale x 16 x i1> %pext, i64 0
  ret i1 %bit
}

define i1 @whilehi_first_active(i64 %a, i64 %b) {
; CHECK-LABEL: whilehi_first_active:
; CHECK:       // %bb.0:
; CHECK-NEXT:    whilehi pn8.h, x0, x1, vlx4
; CHECK-NEXT:    cset w0, mi
; CHECK-NEXT:    ret
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilehi.c16(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 8 x i1> @llvm.aarch64.sve.pext.nxv8i1(target("aarch64.svcount") %while, i32 0)
  %bit = extractelement <vscale x 8 x i1> %pext, i64 0
  ret i1 %bit
}

define i1 @whilelo_first_active(i64 %a, i64 %b) {
; CHECK-LABEL: whilelo_first_active:
; CHECK:       // %bb.0:
; CHECK-NEXT:    whilelo pn8.s, x0, x1, vlx4
; CHECK-NEXT:    cset w0, mi
; CHECK-NEXT:    ret
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilelo.c32(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 4 x i1> @llvm.aarch64.sve.pext.nxv4i1(target("aarch64.svcount") %while, i32 0)
  %bit = extractelement <vscale x 4 x i1> %pext, i64 0
  ret i1 %bit
}

define i1 @whilels_first_active(i64 %a, i64 %b) {
; CHECK-LABEL: whilels_first_active:
; CHECK:       // %bb.0:
; CHECK-NEXT:    whilels pn8.d, x0, x1, vlx4
; CHECK-NEXT:    cset w0, mi
; CHECK-NEXT:    ret
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilels.c64(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 2 x i1> @llvm.aarch64.sve.pext.nxv2i1(target("aarch64.svcount") %while, i32 0)
  %bit = extractelement <vscale x 2 x i1> %pext, i64 0
  ret i1 %bit
}

define void @whilege_first_active_branch(i64 %a, i64 %b) {
; CHECK-LABEL: whilege_first_active_branch:
; CHECK:       // %bb.0: // %entry
; CHECK-NEXT:    whilege pn8.b, x0, x1, vlx4
; CHECK-NEXT:    cset w8, mi
; CHECK-NEXT:    cbz w8, .LBB8_2
; CHECK-NEXT:  // %bb.1: // %then
; CHECK-NEXT:    //APP
; CHECK-NEXT:    //NO_APP
; CHECK-NEXT:  .LBB8_2: // %exit
; CHECK-NEXT:    ret
entry:
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilege.c8(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 16 x i1> @llvm.aarch64.sve.pext.nxv16i1(target("aarch64.svcount") %while, i32 0)
  %bit = extractelement <vscale x 16 x i1> %pext, i64 0
  br i1 %bit, label %then, label %exit

then:
  tail call void asm sideeffect "", ""()
  br label %exit

exit:
  ret void
}

define void @whilelo_first_active_branch(i64 %a, i64 %b) {
; CHECK-LABEL: whilelo_first_active_branch:
; CHECK:       // %bb.0: // %entry
; CHECK-NEXT:    whilelo pn8.s, x0, x1, vlx4
; CHECK-NEXT:    cset w8, mi
; CHECK-NEXT:    tbnz w8, #0, .LBB9_2
; CHECK-NEXT:  // %bb.1: // %then
; CHECK-NEXT:    //APP
; CHECK-NEXT:    //NO_APP
; CHECK-NEXT:  .LBB9_2: // %exit
; CHECK-NEXT:    ret
entry:
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilelo.c32(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 4 x i1> @llvm.aarch64.sve.pext.nxv4i1(target("aarch64.svcount") %while, i32 0)
  %bit = extractelement <vscale x 4 x i1> %pext, i64 0
  br i1 %bit, label %exit, label %then

then:
  tail call void asm sideeffect "", ""()
  br label %exit

exit:
  ret void
}

; Negative test: a nonzero pext offset does not match the first-active fold.
define i1 @whilelo_pext_nonzero_offset(i64 %a, i64 %b) {
; CHECK-LABEL: whilelo_pext_nonzero_offset:
; CHECK:       // %bb.0:
; CHECK-NEXT:    whilelo pn8.s, x0, x1, vlx4
; CHECK-NEXT:    pext p0.s, pn8[1]
; CHECK-NEXT:    mov z0.s, p0/z, #1 // =0x1
; CHECK-NEXT:    fmov w8, s0
; CHECK-NEXT:    and w0, w8, #0x1
; CHECK-NEXT:    ret
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilelo.c32(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 4 x i1> @llvm.aarch64.sve.pext.nxv4i1(target("aarch64.svcount") %while, i32 1)
  %bit = extractelement <vscale x 4 x i1> %pext, i64 0
  ret i1 %bit
}

; Negative test: extracting any element other than lane 0 does not match the first-active fold.
define i1 @whilelt_extractelement_nonzero_index(i64 %a, i64 %b) {
; CHECK-SVE-LABEL: whilelt_extractelement_nonzero_index:
; CHECK-SVE:       // %bb.0:
; CHECK-SVE-NEXT:    whilelt pn8.h, x0, x1, vlx4
; CHECK-SVE-NEXT:    pext p0.h, pn8[0]
; CHECK-SVE-NEXT:    mov z0.h, p0/z, #1 // =0x1
; CHECK-SVE-NEXT:    umov w8, v0.h[1]
; CHECK-SVE-NEXT:    and w0, w8, #0x1
; CHECK-SVE-NEXT:    ret
;
; CHECK-SME-LABEL: whilelt_extractelement_nonzero_index:
; CHECK-SME:       // %bb.0:
; CHECK-SME-NEXT:    whilelt pn8.h, x0, x1, vlx4
; CHECK-SME-NEXT:    pext p0.h, pn8[0]
; CHECK-SME-NEXT:    mov z0.h, p0/z, #1 // =0x1
; CHECK-SME-NEXT:    mov z0.h, z0.h[1]
; CHECK-SME-NEXT:    fmov w8, s0
; CHECK-SME-NEXT:    and w0, w8, #0x1
; CHECK-SME-NEXT:    ret
  %while = call target("aarch64.svcount") @llvm.aarch64.sve.whilelt.c16(i64 %a, i64 %b, i32 4)
  %pext = call <vscale x 8 x i1> @llvm.aarch64.sve.pext.nxv8i1(target("aarch64.svcount") %while, i32 0)
  %bit = extractelement <vscale x 8 x i1> %pext, i64 1
  ret i1 %bit
}