; NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
; RUN: llc -mtriple=aarch64 -aarch64-streaming-hazard-size=0 -mattr=+sme,+sve -aarch64-new-sme-abi -stop-before=aarch64-machine-sme-abi -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK-COMMON,CHECK-BEFORE-SMEABI
; RUN: llc -mtriple=aarch64 -aarch64-streaming-hazard-size=0 -mattr=+sme,+sve -aarch64-new-sme-abi -stop-after=aarch64-machine-sme-abi -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK-COMMON,CHECK-AFTER-SMEABI

declare void @private_za_callee()
declare void @shared_za_callee() "aarch64_inout_za"
declare i64 @shared_za_callee_many_args(i64, i64, i64, i64, i64, i64, i64, i64, i64, i64) "aarch64_inout_za"

; Tests a "RequiresZASavePseudo" is placed after the "ADJCALLSTACKDOWN" for private ZA callees.
define void @requires_za_save() nounwind "aarch64_inout_za" {
  ; CHECK-BEFORE-SMEABI-LABEL: name: requires_za_save
  ; CHECK-BEFORE-SMEABI: bb.0 (%ir-block.0):
  ; CHECK-BEFORE-SMEABI-NEXT:   ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-BEFORE-SMEABI-NEXT:   RequiresZASavePseudo
  ; CHECK-BEFORE-SMEABI-NEXT:   BL @private_za_callee, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
  ; CHECK-BEFORE-SMEABI-NEXT:   ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-BEFORE-SMEABI-NEXT:   RET_ReallyLR
  ;
  ; CHECK-AFTER-SMEABI-LABEL: name: requires_za_save
  ; CHECK-AFTER-SMEABI: bb.0 (%ir-block.0):
  ; CHECK-AFTER-SMEABI-NEXT:   [[RDSVLI_XI:%[0-9]+]]:gpr64 = RDSVLI_XI 1, implicit $vg
  ; CHECK-AFTER-SMEABI-NEXT:   [[COPY:%[0-9]+]]:gpr64 = COPY $sp
  ; CHECK-AFTER-SMEABI-NEXT:   [[MSUBXrrr:%[0-9]+]]:gpr64 = MSUBXrrr [[RDSVLI_XI]], [[RDSVLI_XI]], [[COPY]]
  ; CHECK-AFTER-SMEABI-NEXT:   $sp = COPY [[MSUBXrrr]]
  ; CHECK-AFTER-SMEABI-NEXT:   STPXi [[MSUBXrrr]], [[RDSVLI_XI]], %stack.0, 0
  ; CHECK-AFTER-SMEABI-NEXT:   [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri %stack.0, 0, 0
  ; CHECK-AFTER-SMEABI-NEXT:   [[COPY1:%[0-9]+]]:gpr64 = COPY [[ADDXri]]
  ; CHECK-AFTER-SMEABI-NEXT:   MSR 56965, [[COPY1]]
  ; CHECK-AFTER-SMEABI-NEXT:   ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-AFTER-SMEABI-NEXT:   RequiresZASavePseudo
  ; CHECK-AFTER-SMEABI-NEXT:   BL @private_za_callee, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
  ; CHECK-AFTER-SMEABI-NEXT:   ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-AFTER-SMEABI-NEXT:   MSRpstatesvcrImm1 2, 1, implicit-def $nzcv
  ; CHECK-AFTER-SMEABI-NEXT:   [[MRS:%[0-9]+]]:gpr64 = MRS 56965, implicit-def $nzcv
  ; CHECK-AFTER-SMEABI-NEXT:   $x0 = ADDXri %stack.0, 0, 0
  ; CHECK-AFTER-SMEABI-NEXT:   RestoreZAPseudo [[MRS]], $x0, &__arm_tpidr2_restore, csr_aarch64_sme_abi_support_routines_preservemost_from_x0
  ; CHECK-AFTER-SMEABI-NEXT:   MSR 56965, $xzr
  ; CHECK-AFTER-SMEABI-NEXT:   RET_ReallyLR
  call void @private_za_callee()
  ret void
}

; Tests ZA state markers like "RequiresZASavePseudo" are placed before any streaming mode changes.
define void @requires_za_save_streaming_mode_change() nounwind "aarch64_inout_za" "aarch64_pstate_sm_enabled" {
  ; CHECK-BEFORE-SMEABI-LABEL: name: requires_za_save_streaming_mode_change
  ; CHECK-BEFORE-SMEABI: bb.0 (%ir-block.0):
  ; CHECK-BEFORE-SMEABI-NEXT:   ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-BEFORE-SMEABI-NEXT:   RequiresZASavePseudo
  ; CHECK-BEFORE-SMEABI-NEXT:   MSRpstatesvcrImm1 1, 0, csr_aarch64_smstartstop, implicit-def dead $nzcv, implicit-def $sp, implicit $vg, implicit-def $vg, implicit-def $fpmr
  ; CHECK-BEFORE-SMEABI-NEXT:   BL @private_za_callee, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
  ; CHECK-BEFORE-SMEABI-NEXT:   ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-BEFORE-SMEABI-NEXT:   MSRpstatesvcrImm1 1, 1, csr_aarch64_smstartstop, implicit-def dead $nzcv, implicit $vg, implicit-def $vg, implicit-def $fpmr
  ; CHECK-BEFORE-SMEABI-NEXT:   RET_ReallyLR
  ;
  ; CHECK-AFTER-SMEABI-LABEL: name: requires_za_save_streaming_mode_change
  ; CHECK-AFTER-SMEABI: bb.0 (%ir-block.0):
  ; CHECK-AFTER-SMEABI-NEXT:   [[RDSVLI_XI:%[0-9]+]]:gpr64 = RDSVLI_XI 1, implicit $vg
  ; CHECK-AFTER-SMEABI-NEXT:   [[COPY:%[0-9]+]]:gpr64 = COPY $sp
  ; CHECK-AFTER-SMEABI-NEXT:   [[MSUBXrrr:%[0-9]+]]:gpr64 = MSUBXrrr [[RDSVLI_XI]], [[RDSVLI_XI]], [[COPY]]
  ; CHECK-AFTER-SMEABI-NEXT:   $sp = COPY [[MSUBXrrr]]
  ; CHECK-AFTER-SMEABI-NEXT:   STPXi [[MSUBXrrr]], [[RDSVLI_XI]], %stack.0, 0
  ; CHECK-AFTER-SMEABI-NEXT:   [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri %stack.0, 0, 0
  ; CHECK-AFTER-SMEABI-NEXT:   [[COPY1:%[0-9]+]]:gpr64 = COPY [[ADDXri]]
  ; CHECK-AFTER-SMEABI-NEXT:   MSR 56965, [[COPY1]]
  ; CHECK-AFTER-SMEABI-NEXT:   ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-AFTER-SMEABI-NEXT:   RequiresZASavePseudo
  ; CHECK-AFTER-SMEABI-NEXT:   MSRpstatesvcrImm1 1, 0, csr_aarch64_smstartstop, implicit-def dead $nzcv, implicit-def $sp, implicit $vg, implicit-def $vg, implicit-def $fpmr
  ; CHECK-AFTER-SMEABI-NEXT:   BL @private_za_callee, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
  ; CHECK-AFTER-SMEABI-NEXT:   ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-AFTER-SMEABI-NEXT:   MSRpstatesvcrImm1 1, 1, csr_aarch64_smstartstop, implicit-def dead $nzcv, implicit $vg, implicit-def $vg, implicit-def $fpmr
  ; CHECK-AFTER-SMEABI-NEXT:   MSRpstatesvcrImm1 2, 1, implicit-def $nzcv
  ; CHECK-AFTER-SMEABI-NEXT:   [[MRS:%[0-9]+]]:gpr64 = MRS 56965, implicit-def $nzcv
  ; CHECK-AFTER-SMEABI-NEXT:   $x0 = ADDXri %stack.0, 0, 0
  ; CHECK-AFTER-SMEABI-NEXT:   RestoreZAPseudo [[MRS]], $x0, &__arm_tpidr2_restore, csr_aarch64_sme_abi_support_routines_preservemost_from_x0
  ; CHECK-AFTER-SMEABI-NEXT:   MSR 56965, $xzr
  ; CHECK-AFTER-SMEABI-NEXT:   RET_ReallyLR
  call void @private_za_callee()
  ret void
}

; Tests "InOutZAUsePseudo" is placed after the "ADJCALLSTACKDOWN" for shared ZA callees.
define void @inout_za_call() nounwind "aarch64_inout_za" {
  ; CHECK-COMMON-LABEL: name: inout_za_call
  ; CHECK-COMMON: bb.0 (%ir-block.0):
  ; CHECK-COMMON-NEXT:   ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-COMMON-NEXT:   InOutZAUsePseudo
  ; CHECK-COMMON-NEXT:   BL @shared_za_callee, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp
  ; CHECK-COMMON-NEXT:   ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-COMMON-NEXT:   RET_ReallyLR
  call void @shared_za_callee()
  ret void
}

; Tests ZA state markers like "InOutZAUsePseudo" are placed before any stack argument setup.
define void @many_args_inout_za_call(i64 %0, i64 %1, i64 %2, i64 %3, i64 %4, i64 %5, i64 %6, i64 %7, i64 %8, i64 %9) "aarch64_inout_za" {
  ; CHECK-COMMON-LABEL: name: many_args_inout_za_call
  ; CHECK-COMMON: bb.0 (%ir-block.10):
  ; CHECK-COMMON-NEXT:   liveins: $x0, $x1, $x2, $x3, $x4, $x5, $x6, $x7
  ; CHECK-COMMON-NEXT: {{  $}}
  ; CHECK-COMMON-NEXT:   [[COPY:%[0-9]+]]:gpr64 = COPY $x7
  ; CHECK-COMMON-NEXT:   [[COPY1:%[0-9]+]]:gpr64 = COPY $x6
  ; CHECK-COMMON-NEXT:   [[COPY2:%[0-9]+]]:gpr64 = COPY $x5
  ; CHECK-COMMON-NEXT:   [[COPY3:%[0-9]+]]:gpr64 = COPY $x4
  ; CHECK-COMMON-NEXT:   [[COPY4:%[0-9]+]]:gpr64 = COPY $x3
  ; CHECK-COMMON-NEXT:   [[COPY5:%[0-9]+]]:gpr64 = COPY $x2
  ; CHECK-COMMON-NEXT:   [[COPY6:%[0-9]+]]:gpr64 = COPY $x1
  ; CHECK-COMMON-NEXT:   [[COPY7:%[0-9]+]]:gpr64 = COPY $x0
  ; CHECK-COMMON-NEXT:   [[LDRXui:%[0-9]+]]:gpr64 = LDRXui %fixed-stack.1, 0 :: (load (s64) from %fixed-stack.1, align 16)
  ; CHECK-COMMON-NEXT:   [[LDRXui1:%[0-9]+]]:gpr64 = LDRXui %fixed-stack.0, 0 :: (load (s64) from %fixed-stack.0)
  ; CHECK-COMMON-NEXT:   ADJCALLSTACKDOWN 16, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-COMMON-NEXT:   InOutZAUsePseudo
  ; CHECK-COMMON-NEXT:   [[COPY8:%[0-9]+]]:gpr64sp = COPY $sp
  ; CHECK-COMMON-NEXT:   STRXui killed [[LDRXui1]], [[COPY8]], 1 :: (store (s64) into stack + 8)
  ; CHECK-COMMON-NEXT:   STRXui killed [[LDRXui]], [[COPY8]], 0 :: (store (s64) into stack)
  ; CHECK-COMMON-NEXT:   $x0 = COPY [[COPY7]]
  ; CHECK-COMMON-NEXT:   $x1 = COPY [[COPY6]]
  ; CHECK-COMMON-NEXT:   $x2 = COPY [[COPY5]]
  ; CHECK-COMMON-NEXT:   $x3 = COPY [[COPY4]]
  ; CHECK-COMMON-NEXT:   $x4 = COPY [[COPY3]]
  ; CHECK-COMMON-NEXT:   $x5 = COPY [[COPY2]]
  ; CHECK-COMMON-NEXT:   $x6 = COPY [[COPY1]]
  ; CHECK-COMMON-NEXT:   $x7 = COPY [[COPY]]
  ; CHECK-COMMON-NEXT:   BL @shared_za_callee_many_args, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit $x0, implicit $x1, implicit $x2, implicit $x3, implicit $x4, implicit $x5, implicit $x6, implicit $x7, implicit-def $sp, implicit-def $x0
  ; CHECK-COMMON-NEXT:   ADJCALLSTACKUP 16, 0, implicit-def dead $sp, implicit $sp
  ; CHECK-COMMON-NEXT:   [[COPY9:%[0-9]+]]:gpr64all = COPY $x0
  ; CHECK-COMMON-NEXT:   RET_ReallyLR
  %ret = call i64 @shared_za_callee_many_args(i64 %0, i64 %1, i64 %2, i64 %3, i64 %4, i64 %5, i64 %6, i64 %7, i64 %8, i64 %9)
  ret void
}