diff --git a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp index f63981b87c1c1..34d74d04c4419 100644 --- a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp +++ b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp @@ -1063,6 +1063,7 @@ AArch64ExpandPseudo::expandCommitZASave(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) { MachineInstr &MI = *MBBI; DebugLoc DL = MI.getDebugLoc(); + [[maybe_unused]] auto *RI = MBB.getParent()->getSubtarget().getRegisterInfo(); // Compare TPIDR2_EL0 against 0. Commit ZA if TPIDR2_EL0 is non-zero. MachineInstrBuilder Branch = @@ -1073,21 +1074,25 @@ AArch64ExpandPseudo::expandCommitZASave(MachineBasicBlock &MBB, MachineInstrBuilder MIB = BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::BL)); // Copy operands (mainly the regmask) from the pseudo. - for (unsigned I = 2; I < MI.getNumOperands(); ++I) + for (unsigned I = 3; I < MI.getNumOperands(); ++I) MIB.add(MI.getOperand(I)); // Clear TPIDR2_EL0. BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::MSR)) .addImm(AArch64SysReg::TPIDR2_EL0) .addReg(AArch64::XZR); bool ZeroZA = MI.getOperand(1).getImm() != 0; + bool ZeroZT0 = MI.getOperand(2).getImm() != 0; if (ZeroZA) { - [[maybe_unused]] auto *TRI = - MBB.getParent()->getSubtarget().getRegisterInfo(); - assert(MI.definesRegister(AArch64::ZAB0, TRI) && "should define ZA!"); + assert(MI.definesRegister(AArch64::ZAB0, RI) && "should define ZA!"); BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::ZERO_M)) .addImm(ZERO_ALL_ZA_MASK) .addDef(AArch64::ZAB0, RegState::ImplicitDefine); } + if (ZeroZT0) { + assert(MI.definesRegister(AArch64::ZT0, RI) && "should define ZT0!"); + BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::ZERO_T)) + .addDef(AArch64::ZT0); + } MI.eraseFromParent(); return &EndBB; diff --git a/llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td b/llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td index 5bb70ee11b06d..737169253ddb3 100644 --- a/llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td +++ b/llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td @@ -108,7 +108,8 @@ def SMEStateAllocPseudo : Pseudo<(outs), (ins), []>, Sched<[]>; def CommitZASavePseudo : Pseudo<(outs), - (ins GPR64:$tpidr2_el0, i1imm:$zero_za, i64imm:$commit_routine, variable_ops), []>, + (ins GPR64:$tpidr2_el0, i1imm:$zero_za, i1imm:$zero_zt0, + i64imm:$commit_routine, variable_ops), []>, Sched<[]>; def AArch64_inout_za_use diff --git a/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp b/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp index 24d30c731b945..2afbec92392f0 100644 --- a/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp +++ b/llvm/lib/Target/AArch64/MachineSMEABIPass.cpp @@ -842,6 +842,7 @@ void MachineSMEABI::emitNewZAPrologue(MachineBasicBlock &MBB, BuildMI(MBB, MBBI, DL, TII->get(AArch64::CommitZASavePseudo)) .addReg(TPIDR2EL0) .addImm(ZeroZA ? 1 : 0) + .addImm(/*ZeroZT0=*/false) .addExternalSymbol(TLI->getLibcallName(RTLIB::SMEABI_TPIDR2_SAVE)) .addRegMask(TRI->SMEABISupportRoutinesCallPreservedMaskFromX0()); if (ZeroZA) diff --git a/llvm/test/CodeGen/AArch64/expand-sme-pseudos.mir b/llvm/test/CodeGen/AArch64/expand-sme-pseudos.mir index 6ca9b9b6cb200..9b745d56c4b7f 100644 --- a/llvm/test/CodeGen/AArch64/expand-sme-pseudos.mir +++ b/llvm/test/CodeGen/AArch64/expand-sme-pseudos.mir @@ -62,7 +62,7 @@ body: | ; CHECK-NEXT: RET undef $lr $x8 = MRS 56965, implicit-def $nzcv - CommitZASavePseudo $x8, 0, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0 + CommitZASavePseudo $x8, 0, 0, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0 RET_ReallyLR @@ -94,7 +94,72 @@ body: | ; CHECK-NEXT: RET undef $lr $x8 = MRS 56965, implicit-def $nzcv - CommitZASavePseudo $x8, 1, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $zab0 + CommitZASavePseudo $x8, 1, 0, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $zab0 + + RET_ReallyLR + +... +--- +# X8 = TPIDR2_EL0 +name: commit_za_save_zero_zt0 +alignment: 4 +tracksRegLiveness: true +body: | + bb.0: + ; CHECK-LABEL: name: commit_za_save_zero_zt0 + ; CHECK: successors: %bb.1(0x40000000), %bb.2(0x40000000) + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: $x8 = MRS 56965, implicit-def $nzcv + ; CHECK-NEXT: CBNZX $x8, %bb.1 + ; CHECK-NEXT: B %bb.2 + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: .1: + ; CHECK-NEXT: successors: %bb.2(0x80000000) + ; CHECK-NEXT: liveins: $x8 + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: BL &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $lr, implicit $sp, implicit-def $zt0 + ; CHECK-NEXT: MSR 56965, $xzr + ; CHECK-NEXT: $zt0 = ZERO_T + ; CHECK-NEXT: B %bb.2 + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: .2: + ; CHECK-NEXT: RET undef $lr + $x8 = MRS 56965, implicit-def $nzcv + + CommitZASavePseudo $x8, 0, 1, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $zt0 + + RET_ReallyLR + +... +--- +# X8 = TPIDR2_EL0 +name: commit_za_save_zero_everything +alignment: 4 +tracksRegLiveness: true +body: | + bb.0: + ; CHECK-LABEL: name: commit_za_save_zero_everything + ; CHECK: successors: %bb.1(0x40000000), %bb.2(0x40000000) + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: $x8 = MRS 56965, implicit-def $nzcv + ; CHECK-NEXT: CBNZX $x8, %bb.1 + ; CHECK-NEXT: B %bb.2 + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: .1: + ; CHECK-NEXT: successors: %bb.2(0x80000000) + ; CHECK-NEXT: liveins: $x8 + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: BL &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $lr, implicit $sp, implicit-def $zab0, implicit-def $zt0 + ; CHECK-NEXT: MSR 56965, $xzr + ; CHECK-NEXT: ZERO_M 255, implicit-def $zab0 + ; CHECK-NEXT: $zt0 = ZERO_T + ; CHECK-NEXT: B %bb.2 + ; CHECK-NEXT: {{ $}} + ; CHECK-NEXT: .2: + ; CHECK-NEXT: RET undef $lr + $x8 = MRS 56965, implicit-def $nzcv + + CommitZASavePseudo $x8, 1, 1, &__arm_tpidr2_save, csr_aarch64_sme_abi_support_routines_preservemost_from_x0, implicit-def $zab0, implicit-def $zt0 RET_ReallyLR