From 17234b3b05c98ca837b96d4e31b7be0dff3f8734 Mon Sep 17 00:00:00 2001 From: AkshayK Date: Tue, 15 Sep 2026 13:23:31 -0400 Subject: [PATCH] [CodeGen] Add per-exit scratch register clearing Track scratch registers per exit and merge them into register clearing after candidate filtering. Validate target support, register width, reserved and finalized callee-saved registers, return addresses, and live exit operands before emission. Keep stack clearing capability-driven and remove the synthetic scratch-injection option. Exercise the shared consumer with 70 unit cases using real target backends, add an unsupported-stack regression, and register the sources in CMake and GN. Tighten clearing comments and remove private tracker references. Validation: assertions-enabled llc and CodeGenTests built successfully; 391 unit tests and 312 selected lit tests passed. GN source lists were checked; GN was unavailable locally. Co-authored-by: Claude --- .../llvm/CodeGen/TargetFrameLowering.h | 11 + llvm/lib/CodeGen/CMakeLists.txt | 1 + llvm/lib/CodeGen/PrologEpilogInserter.cpp | 147 +++++---- llvm/lib/CodeGen/RegisterClearing.cpp | 97 ++++++ llvm/lib/CodeGen/RegisterClearing.h | 29 ++ .../Target/AArch64/AArch64FrameLowering.cpp | 8 + .../lib/Target/AArch64/AArch64FrameLowering.h | 3 + llvm/lib/Target/RISCV/RISCVFrameLowering.cpp | 5 + llvm/lib/Target/RISCV/RISCVFrameLowering.h | 3 + llvm/lib/Target/X86/X86FrameLowering.cpp | 8 + llvm/lib/Target/X86/X86FrameLowering.h | 3 + llvm/test/CodeGen/ARM/zeroize-fallback.ll | 2 +- .../CodeGen/RISCV/zero-call-used-regs-ra.ll | 1 - .../CodeGen/X86/zeroize-fallback-exits.ll | 1 - .../test/CodeGen/X86/zeroize-fallback-mode.ll | 2 +- .../test/CodeGen/X86/zeroize-fallback-regs.ll | 1 - .../CodeGen/X86/zeroize-stack-no-emission.ll | 19 ++ llvm/test/Verifier/zeroize-stack-musttail.ll | 13 +- llvm/unittests/CodeGen/CMakeLists.txt | 1 + .../CodeGen/RegisterClearingTest.cpp | 303 ++++++++++++++++++ .../gn/secondary/llvm/lib/CodeGen/BUILD.gn | 1 + .../secondary/llvm/unittests/CodeGen/BUILD.gn | 1 + 22 files changed, 577 insertions(+), 83 deletions(-) create mode 100644 llvm/lib/CodeGen/RegisterClearing.cpp create mode 100644 llvm/lib/CodeGen/RegisterClearing.h create mode 100644 llvm/test/CodeGen/X86/zeroize-stack-no-emission.ll create mode 100644 llvm/unittests/CodeGen/RegisterClearingTest.cpp diff --git a/llvm/include/llvm/CodeGen/TargetFrameLowering.h b/llvm/include/llvm/CodeGen/TargetFrameLowering.h index 7da579736bf08e..85144b81c2516d 100644 --- a/llvm/include/llvm/CodeGen/TargetFrameLowering.h +++ b/llvm/include/llvm/CodeGen/TargetFrameLowering.h @@ -219,6 +219,17 @@ class LLVM_ABI TargetFrameLowering { return false; } + /// Whether emitZeroCallUsedRegs can clear Reg as a scratch register. A true + /// answer guarantees that the entire register is cleared, without widening + /// the write into another register (apart from condition flags). PEI checks + /// separately that Reg is allocatable and not needed at the exit. + /// Targets must opt in explicitly; supporting a register-clear mode does not + /// guarantee that every physical register can be cleared. + virtual bool isZeroCallUsedRegsScratchReg(const MachineFunction &MF, + MCRegister Reg) const { + return false; + } + /// emitZeroCallUsedRegs - Zeros out call used registers. Only called on /// targets whose supportsZeroCallUsedRegs returns true. /// diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt index c572128b023c1e..2e2cdbbc6824b9 100644 --- a/llvm/lib/CodeGen/CMakeLists.txt +++ b/llvm/lib/CodeGen/CMakeLists.txt @@ -195,6 +195,7 @@ add_llvm_component_library(LLVMCodeGen RegisterClassInfo.cpp RegisterCoalescer.cpp RegisterPressure.cpp + RegisterClearing.cpp RegisterScavenging.cpp GCEmptyBasicBlocks.cpp Rematerializer.cpp diff --git a/llvm/lib/CodeGen/PrologEpilogInserter.cpp b/llvm/lib/CodeGen/PrologEpilogInserter.cpp index 7ac4c5b9bf2fda..ce828db933cb36 100644 --- a/llvm/lib/CodeGen/PrologEpilogInserter.cpp +++ b/llvm/lib/CodeGen/PrologEpilogInserter.cpp @@ -15,6 +15,7 @@ // //===----------------------------------------------------------------------===// +#include "RegisterClearing.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/BitVector.h" #include "llvm/ADT/STLExtras.h" @@ -92,44 +93,15 @@ namespace { //===----------------------------------------------------------------------===// // The clearing sequence. // -// A function told to destroy what it leaves behind has several things to clear, -// and they are not independent: clearing the frame needs registers, and every -// step writes the flags. So an exit emits one ordered sequence, run at every -// in-scope exit, not a set of independent steps; the order is enforced. -// -// The order is ClearStack, then ClearRegisters, then ClearFlags. Each position -// has a reason that outlives whichever implementation sits behind it: -// -// - ClearStack is first because it needs registers: the value and the address -// it stores through are live in registers derived from the frame, so it -// leaves in registers what it took out of memory, and a register clear in -// front of it would be undone. No target implements it yet; whoever does -// must name those registers so the following register clear covers them. -// -// - ClearRegisters is after every step that needs a register and before every -// step that does not, so it sees the registers' final state. Anything added -// later that computes an address, length or value must go in front of it. -// -// - ClearFlags is last because every other step writes the flags (an x86 -// register clear xors; a looping stack clear sets them from the count), so a -// flag clear placed earlier would be overwritten. -// -// The order is over the emitted code, not over one insertion point. Today every -// step emits at the exit's insertion point, after the epilogue. A step that has -// to run earlier (a frame clear, before the epilogue moves SP and the frame -// stops being addressable) still has to leave every later step behind it in -// program order. -// -// Two whole-sequence invariants a later step must keep: -// -// - It runs only at the exits getEnforceableExit() picks, decided once rather -// than per step. A step that cannot be placed at an in-scope exit is a gap to -// record, not a licence to pick its own sites. -// -// - It does not depend on secret values: which steps run comes from attributes -// and target capabilities, where they run from control-flow shape. Two runs -// of a protected function execute the same sequence. +// Emit ClearStack, ClearRegisters, then ClearFlags at each enforceable exit. +// Stack clearing can leave sensitive data in scratch registers; register +// clearing must follow it, and flag clearing must follow all flag-writing +// steps. A future stack clear may need to run before the epilogue, while the +// frame is addressable, but must preserve this emission order. // +// All steps use the exits selected by getEnforceableExit(). Their selection and +// placement must depend only on attributes, target capabilities, and control +// flow, never on secret values. //===----------------------------------------------------------------------===// /// One step of the clearing sequence. The order is not this enumeration's @@ -283,9 +255,12 @@ class PEIImpl { ClearingDisposition planClearStack(MachineFunction &MF); ClearingDisposition planClearRegisters(MachineFunction &MF, BitVector &CandidateRegsToZero); + ClearingDisposition planClearRegistersForScratch( + MachineFunction &MF, BitVector &CandidateRegsToZero); void emitClearingStep(ClearingStep Step, const ExitClearingPlan &Plan, MachineBasicBlock &MBB, - MachineBasicBlock::iterator InsertPt); + MachineBasicBlock::iterator InsertPt, + BitVector &ScratchRegs); void diagnoseIgnoredZeroizeRequestsOnNakedFunction(MachineFunction &MF); public: @@ -1593,6 +1568,8 @@ void PEIImpl::insertClearingSequences(MachineFunction &MF) { if (!Plan.anyStepEmits() && !PrintClearingSequence) return; + const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo(); + raw_ostream &OS = errs(); if (PrintClearingSequence) OS << "clearing sequence for function '" << MF.getName() << "':\n"; @@ -1623,70 +1600,85 @@ void PEIImpl::insertClearingSequences(MachineFunction &MF) { OS << " " << printMBBReference(MBB) << " " << getExitKindName(*ExitMI) << ":"; + // What the steps in front of the register clear leave in registers. It is + // built as the sequence runs at this exit and read by the register clear + // at this exit; see the comment on scratch registers above. + BitVector ScratchRegs(TRI.getNumRegs()); + for (ClearingStep Step : ClearingSequence) { ClearingDisposition D = Plan.dispositionOf(Step); if (D == ClearingDisposition::Emit) - emitClearingStep(Step, Plan, MBB, InsertPt); + emitClearingStep(Step, Plan, MBB, InsertPt, ScratchRegs); if (PrintClearingSequence) OS << " " << getClearingStepName(Step) << "=" << getClearingDispositionName(D); } - if (PrintClearingSequence) + if (PrintClearingSequence) { + // Only when there are any, so that the line a function without a step + // that declares registers prints is the line it printed before. + if (ScratchRegs.any()) { + OS << " scratch="; + const char *Sep = ""; + for (unsigned Reg : ScratchRegs.set_bits()) { + OS << Sep << TRI.getName(Reg); + Sep = ","; + } + } OS << "\n"; + } } if (PrintClearingSequence) OS << "end clearing sequence for function '" << MF.getName() << "'\n"; } -/// emitClearingStep - Emit one step of the clearing sequence at \p InsertPt. -/// -/// A step that emits nothing today still has its case here, so that the -/// implementation of it lands at the position the order gives it rather than -/// wherever it is convenient. +/// Emit one clearing step at \p InsertPt. Earlier steps add their scratch +/// registers to \p ScratchRegs; the register clear validates and clears them. void PEIImpl::emitClearingStep(ClearingStep Step, const ExitClearingPlan &Plan, MachineBasicBlock &MBB, - MachineBasicBlock::iterator InsertPt) { + MachineBasicBlock::iterator InsertPt, + BitVector &ScratchRegs) { MachineFunction &MF = *MBB.getParent(); - const TargetFrameLowering &TFI = *MF.getSubtarget().getFrameLowering(); const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo(); switch (Step) { case ClearingStep::ClearStack: - // Nothing emits here yet: no target can clear the frame, so planning has - // already refused every request for it and this step never reaches - // emission. It is first in the order because it needs registers to run, - // and the register clear after it is what destroys those. + // A future stack clear must declare its scratch registers here. break; - case ClearingStep::ClearRegisters: - // What to clear is settled here rather than in the plan, because it is the - // exit that decides it: see computeRegsToClearAtExit. - TFI.emitZeroCallUsedRegs( - computeRegsToClearAtExit(Plan.CandidateRegsToZero, MBB, InsertPt, TRI), - MBB, InsertPt, RS); + case ClearingStep::ClearRegisters: { + // Filter the mode's candidates using this exit's register requirements. + BitVector RegsToZero = + computeRegsToClearAtExit(Plan.CandidateRegsToZero, MBB, InsertPt, TRI); + + // Add scratch after mode filtering: these registers were dirtied by the + // clearing sequence and need not have been used by the function. + emitZeroCallUsedRegsWithScratch(std::move(RegsToZero), ScratchRegs, MBB, + InsertPt, RS); break; + } case ClearingStep::ClearFlags: - // Nothing emits here yet. It is last in the order because every step in - // front of it writes the flags, so a flag clear anywhere else is undone by - // what follows it. + // Unimplemented. Flag clearing must follow every flag-writing step. break; } } -/// planClearingSequence - Decide what each step of the sequence does in \p MF. -/// -/// The steps are planned in the order they run, so that a function asking for -/// more than one of them is told about them in that order too. +/// Plan the steps and diagnose unsupported requests in emission order. void PEIImpl::planClearingSequence(MachineFunction &MF, ExitClearingPlan &Plan) { Plan.Stack = planClearStack(MF); Plan.Registers = planClearRegisters(MF, Plan.CandidateRegsToZero); - // Nothing asks for the flags to be cleared and nothing clears them. The step - // is planned all the same, so that the sequence a function runs is described - // by the plan in full rather than in the parts that have an implementation. + + // Stack clearing requires a register clear even without a register attribute + // or with mode "skip", which applies only to the function's own register use. + if (Plan.Stack == ClearingDisposition::Emit && + Plan.Registers == ClearingDisposition::NotRequested) + Plan.Registers = + planClearRegistersForScratch(MF, Plan.CandidateRegsToZero); + + // Flag clearing has neither an attribute nor an implementation yet. Plan.Flags = ClearingDisposition::Unimplemented; } @@ -1711,6 +1703,27 @@ ClearingDisposition PEIImpl::planClearStack(MachineFunction &MF) { return ClearingDisposition::Unimplemented; } +/// Enable register clearing for scratch alone. Leave the mode's candidate set +/// empty; each exit supplies its own scratch declarations. +ClearingDisposition +PEIImpl::planClearRegistersForScratch(MachineFunction &MF, + BitVector &CandidateRegsToZero) { + const Function &F = MF.getFunction(); + const TargetFrameLowering &TFI = *MF.getSubtarget().getFrameLowering(); + const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo(); + + // Stack clearing is incomplete if its scratch registers cannot be cleared. + if (!TFI.supportsZeroCallUsedRegs(MF)) { + F.getContext().diagnose(DiagnosticInfoUnsupported{ + F, "clearing the stack needs the registers it uses to be cleared " + "afterwards, which is not supported by this target"}); + return ClearingDisposition::Unsupported; + } + + CandidateRegsToZero.resize(TRI.getNumRegs()); + return ClearingDisposition::Emit; +} + /// planClearRegisters - Decide what the ClearRegisters step does in \p MF, and /// compute the registers it is allowed to clear. /// diff --git a/llvm/lib/CodeGen/RegisterClearing.cpp b/llvm/lib/CodeGen/RegisterClearing.cpp new file mode 100644 index 00000000000000..effd74aba9891d --- /dev/null +++ b/llvm/lib/CodeGen/RegisterClearing.cpp @@ -0,0 +1,97 @@ +//===- RegisterClearing.cpp --------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "RegisterClearing.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/CodeGen/MachineFunction.h" +#include "llvm/CodeGen/MachineRegisterInfo.h" +#include "llvm/CodeGen/TargetFrameLowering.h" +#include "llvm/CodeGen/TargetRegisterInfo.h" +#include "llvm/CodeGen/TargetSubtargetInfo.h" +#include "llvm/IR/DiagnosticInfo.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/LLVMContext.h" + +using namespace llvm; + +/// Whether clearing \p Regs would overwrite a value required at the exit. +static bool anyRegNeededAtExit(const BitVector &Regs, + const MachineBasicBlock &MBB, + MachineBasicBlock::const_iterator InsertPt, + const TargetRegisterInfo &TRI) { + const MachineFunction &MF = *MBB.getParent(); + + // Some return pseudos do not name the return-address register explicitly. + if (MCRegister RAReg = TRI.getRARegister()) + for (MCRegister Reg : TRI.sub_and_superregs_inclusive(RAReg)) + if (Regs.test(Reg.id())) + return true; + + // Use the finalized list, including custom call-saved registers. These values + // must survive even when the exit does not name them. + for (const MCPhysReg *CSRegs = MF.getRegInfo().getCalleeSavedRegs(); + MCPhysReg CSReg = *CSRegs; ++CSRegs) + for (MCRegister Reg : TRI.sub_and_superregs_inclusive(CSReg)) + if (Regs.test(Reg.id())) + return true; + + // Preserve registers referenced by instructions after the insertion point. + for (const MachineInstr &MI : make_range(InsertPt, MBB.end())) + for (const MachineOperand &MO : MI.operands()) { + if (!MO.isReg() || !MO.getReg()) + continue; + for (MCPhysReg SReg : TRI.sub_and_superregs_inclusive(MO.getReg())) + if (Regs.test(SReg)) + return true; + } + + return false; +} + +/// Validate the scratch contract before handing the declarations to the target. +static bool validateScratchRegs(const BitVector &Regs, + const MachineBasicBlock &MBB, + MachineBasicBlock::const_iterator InsertPt, + const TargetFrameLowering &TFI, + const TargetRegisterInfo &TRI) { + if (Regs.none()) + return true; + + const MachineFunction &MF = *MBB.getParent(); + auto Diagnose = [&](const Twine &Message) { + MF.getFunction().getContext().diagnose(DiagnosticInfoUnsupported{ + MF.getFunction(), Message, DiagnosticLocation(), DS_Error}); + return false; + }; + + BitVector Allocatable = TRI.getAllocatableSet(MF); + for (MCRegister Reg : Regs.set_bits()) + if (!TFI.isZeroCallUsedRegsScratchReg(MF, Reg) || !Allocatable.test(Reg)) + return Diagnose(Twine("register '") + TRI.getName(Reg) + + "' is not supported for scratch register clearing"); + + if (anyRegNeededAtExit(Regs, MBB, InsertPt, TRI)) + return Diagnose("scratch register clearing would overwrite a register " + "needed at the exit"); + return true; +} + +bool llvm::emitZeroCallUsedRegsWithScratch(BitVector RegsToZero, + const BitVector &ScratchRegs, + MachineBasicBlock &MBB, + MachineBasicBlock::iterator InsertPt, + RegScavenger *RS) { + const MachineFunction &MF = *MBB.getParent(); + const TargetFrameLowering &TFI = *MF.getSubtarget().getFrameLowering(); + const TargetRegisterInfo &TRI = *MF.getSubtarget().getRegisterInfo(); + if (!validateScratchRegs(ScratchRegs, MBB, InsertPt, TFI, TRI)) + return false; + RegsToZero |= ScratchRegs; + TFI.emitZeroCallUsedRegs(RegsToZero, MBB, InsertPt, RS); + return true; +} diff --git a/llvm/lib/CodeGen/RegisterClearing.h b/llvm/lib/CodeGen/RegisterClearing.h new file mode 100644 index 00000000000000..a272eac6bb63d0 --- /dev/null +++ b/llvm/lib/CodeGen/RegisterClearing.h @@ -0,0 +1,29 @@ +//===- RegisterClearing.h --------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef LLVM_LIB_CODEGEN_REGISTERCLEARING_H +#define LLVM_LIB_CODEGEN_REGISTERCLEARING_H + +#include "llvm/ADT/BitVector.h" +#include "llvm/CodeGen/MachineBasicBlock.h" +#include "llvm/Support/Compiler.h" + +namespace llvm { +class RegScavenger; + +/// Add validated scratch declarations to the exit's filtered register set and +/// emit the clear. The caller must check supportsZeroCallUsedRegs first. +/// Scratch declarations are per exit and independent of the function's mode. +/// Return false and diagnose an invalid declaration without emitting a clear. +LLVM_ABI bool emitZeroCallUsedRegsWithScratch( + BitVector RegsToZero, const BitVector &ScratchRegs, MachineBasicBlock &MBB, + MachineBasicBlock::iterator InsertPt, RegScavenger *RS); + +} // namespace llvm + +#endif diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp index 300c8722dd2698..078c4ec85e478b 100644 --- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp +++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp @@ -843,6 +843,14 @@ static MCRegister getRegisterOrZero(MCRegister Reg, bool HasSVE) { } } +bool AArch64FrameLowering::isZeroCallUsedRegsScratchReg( + const MachineFunction &MF, MCRegister Reg) const { + // The emitter skips X19-X30 even for calling conventions that do not + // preserve them. Accept only full-width GPRs that it actually clears. + return AArch64::GPR64RegClass.contains(Reg) && + getRegisterOrZero(Reg, /*HasSVE=*/false) == Reg; +} + void AArch64FrameLowering::emitZeroCallUsedRegs( BitVector RegsToZero, MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI, RegScavenger *) const { diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.h b/llvm/lib/Target/AArch64/AArch64FrameLowering.h index 94b3b3d70944c3..348ff21ec6c203 100644 --- a/llvm/lib/Target/AArch64/AArch64FrameLowering.h +++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.h @@ -215,6 +215,9 @@ class AArch64FrameLowering : public TargetFrameLowering { return true; } + bool isZeroCallUsedRegsScratchReg(const MachineFunction &MF, + MCRegister Reg) const override; + /// Emit target zero call-used regs. void emitZeroCallUsedRegs(BitVector RegsToZero, MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI, diff --git a/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp b/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp index 7d383825773f54..918dbf2dfcb1b4 100644 --- a/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp +++ b/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp @@ -1534,6 +1534,11 @@ static MCRegister getLargestFPRegisterOrZero(const RISCVSubtarget &STI, return MCRegister(); } +bool RISCVFrameLowering::isZeroCallUsedRegsScratchReg(const MachineFunction &MF, + MCRegister Reg) const { + return RISCV::GPRRegClass.contains(Reg); +} + void RISCVFrameLowering::emitZeroCallUsedRegs(BitVector RegsToZero, MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI, diff --git a/llvm/lib/Target/RISCV/RISCVFrameLowering.h b/llvm/lib/Target/RISCV/RISCVFrameLowering.h index 4bf8e892ad57ad..21957b53bcdae7 100644 --- a/llvm/lib/Target/RISCV/RISCVFrameLowering.h +++ b/llvm/lib/Target/RISCV/RISCVFrameLowering.h @@ -127,6 +127,9 @@ class RISCVFrameLowering : public TargetFrameLowering { return true; } + bool isZeroCallUsedRegsScratchReg(const MachineFunction &MF, + MCRegister Reg) const override; + /// Emit target zero call-used regs. void emitZeroCallUsedRegs(BitVector RegsToZero, MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI, diff --git a/llvm/lib/Target/X86/X86FrameLowering.cpp b/llvm/lib/Target/X86/X86FrameLowering.cpp index 47053d36ddeb56..325c89a75f7c12 100644 --- a/llvm/lib/Target/X86/X86FrameLowering.cpp +++ b/llvm/lib/Target/X86/X86FrameLowering.cpp @@ -624,6 +624,14 @@ void X86FrameLowering::emitCalleeSavedFrameMoves( } } +bool X86FrameLowering::isZeroCallUsedRegsScratchReg(const MachineFunction &MF, + MCRegister Reg) const { + // XOR32rr also zeros the upper half in 64-bit mode. Require the declaration + // to cover the full write, including sibling subregisters such as AL and AH. + return STI.is64Bit() ? X86::GR64RegClass.contains(Reg) + : X86::GR32RegClass.contains(Reg); +} + void X86FrameLowering::emitZeroCallUsedRegs(BitVector RegsToZero, MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI, diff --git a/llvm/lib/Target/X86/X86FrameLowering.h b/llvm/lib/Target/X86/X86FrameLowering.h index 62318ccd8f1267..f4467bcb438ef6 100644 --- a/llvm/lib/Target/X86/X86FrameLowering.h +++ b/llvm/lib/Target/X86/X86FrameLowering.h @@ -283,6 +283,9 @@ class X86FrameLowering : public TargetFrameLowering { return true; } + bool isZeroCallUsedRegsScratchReg(const MachineFunction &MF, + MCRegister Reg) const override; + /// Emit target zero call-used regs. void emitZeroCallUsedRegs(BitVector RegsToZero, MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI, diff --git a/llvm/test/CodeGen/ARM/zeroize-fallback.ll b/llvm/test/CodeGen/ARM/zeroize-fallback.ll index c87f64802d7408..d3cb7f8047c4e3 100644 --- a/llvm/test/CodeGen/ARM/zeroize-fallback.ll +++ b/llvm/test/CodeGen/ARM/zeroize-fallback.ll @@ -1,5 +1,5 @@ ; Two of the fallbacks are decided before any target is asked, so they show on -; a target that cannot clear anything. trailofbits/vspells-ct-internal-notes#24. +; a target that cannot clear anything. ; RUN: split-file %s %t ; RUN: llc -mtriple=armv7-unknown-linux-gnueabi -verify-machineinstrs -pei-print-clearing-sequence %t/exits.ll -o /dev/null 2>&1 | FileCheck --check-prefix=SEQ %s diff --git a/llvm/test/CodeGen/RISCV/zero-call-used-regs-ra.ll b/llvm/test/CodeGen/RISCV/zero-call-used-regs-ra.ll index 81c4c9061aee24..2c7504c5f1060a 100644 --- a/llvm/test/CodeGen/RISCV/zero-call-used-regs-ra.ll +++ b/llvm/test/CodeGen/RISCV/zero-call-used-regs-ra.ll @@ -1,7 +1,6 @@ ; ghccc has no callee-saved registers, and PseudoRET has no explicit ra operand. ; A call's implicit ra definition must not make the clearer erase the return ; address. Cover used modes and all, which does not depend on the used set. -; Regression for trailofbits/vspells-ct-internal-notes#24. ; RUN: llc -mtriple=riscv32 -mattr=+d -verify-machineinstrs < %s | FileCheck %s ; RUN: llc -mtriple=riscv64 -mattr=+d -verify-machineinstrs < %s | FileCheck %s diff --git a/llvm/test/CodeGen/X86/zeroize-fallback-exits.ll b/llvm/test/CodeGen/X86/zeroize-fallback-exits.ll index 040d02ab129fa5..79787d4b7d72b1 100644 --- a/llvm/test/CodeGen/X86/zeroize-fallback-exits.ll +++ b/llvm/test/CodeGen/X86/zeroize-fallback-exits.ll @@ -1,7 +1,6 @@ ; A block with no successors that ends in an instruction the exit classifier ; cannot account for is in scope: it may leave the function, and a dead ; sequence costs less than an uncleared exit. -; trailofbits/vspells-ct-internal-notes#24. ; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs -pei-print-clearing-sequence %s -o /dev/null 2>&1 | FileCheck --check-prefix=SEQ %s ; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs %s -o - | FileCheck %s diff --git a/llvm/test/CodeGen/X86/zeroize-fallback-mode.ll b/llvm/test/CodeGen/X86/zeroize-fallback-mode.ll index cfa9dde6dca07b..186596ec7a0cc3 100644 --- a/llvm/test/CodeGen/X86/zeroize-fallback-mode.ll +++ b/llvm/test/CodeGen/X86/zeroize-fallback-mode.ll @@ -1,6 +1,6 @@ ; A "zero-call-used-regs" value that names no known mode means the widest one, ; as LangRef fixes for an unrecognized "zeroize-stack" mode. The mode switch -; previously had no default. trailofbits/vspells-ct-internal-notes#24. +; previously had no default. ; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs %s -o - | FileCheck %s diff --git a/llvm/test/CodeGen/X86/zeroize-fallback-regs.ll b/llvm/test/CodeGen/X86/zeroize-fallback-regs.ll index 80ffbff8aa2792..586c0326525ce0 100644 --- a/llvm/test/CodeGen/X86/zeroize-fallback-regs.ll +++ b/llvm/test/CodeGen/X86/zeroize-fallback-regs.ll @@ -1,7 +1,6 @@ ; The "used" modes clear only registers the function touched. Registers an ; instruction touches implicitly were not counted, so an asm clobber or a ; physical-register output kept its contents past the return. -; trailofbits/vspells-ct-internal-notes#24. ; RUN: llc -mtriple=x86_64-unknown-linux-gnu %s -o - | FileCheck %s diff --git a/llvm/test/CodeGen/X86/zeroize-stack-no-emission.ll b/llvm/test/CodeGen/X86/zeroize-stack-no-emission.ll new file mode 100644 index 00000000000000..b07006687704c3 --- /dev/null +++ b/llvm/test/CodeGen/X86/zeroize-stack-no-emission.ll @@ -0,0 +1,19 @@ +; RUN: llc -mtriple=x86_64-linux-gnu -pei-print-clearing-sequence -verify-machineinstrs %s -o /dev/null 2>&1 | FileCheck %s +; +; An unsupported stack clear must not be reported as emitted or force a +; register clear. An explicit register request still runs independently. +; CHECK: warning: {{.*}}in function stack_only {{.*}}"zeroize-stack" is not supported by this target +; CHECK: clearing sequence for function 'stack_only': +; CHECK-NEXT: %bb.0 return: clear-stack=unsupported clear-registers=not-requested clear-flags=unimplemented +; CHECK-NEXT: end clearing sequence for function 'stack_only' +define void @stack_only() "zeroize-stack"="used" { + ret void +} + +; CHECK: warning: {{.*}}in function both {{.*}}"zeroize-stack" is not supported by this target +; CHECK: clearing sequence for function 'both': +; CHECK-NEXT: %bb.0 return: clear-stack=unsupported clear-registers=emitted clear-flags=unimplemented +; CHECK-NEXT: end clearing sequence for function 'both' +define i32 @both(i32 %x) "zeroize-stack"="used" "zero-call-used-regs"="used-gpr" { + ret i32 %x +} diff --git a/llvm/test/Verifier/zeroize-stack-musttail.ll b/llvm/test/Verifier/zeroize-stack-musttail.ll index f8c6a7efd81863..1f3781e5a06e0f 100644 --- a/llvm/test/Verifier/zeroize-stack-musttail.ll +++ b/llvm/test/Verifier/zeroize-stack-musttail.ll @@ -1,14 +1,5 @@ -; An ordinary tail call in a protected function is an optimization, and it is -; suppressed. musttail is not an optimization: the caller is required to be -; replaced by the callee, and no pass is allowed to decide otherwise. A function -; cannot both be replaced at the call and clear its frame after it, so the two -; together describe a function that cannot be generated, and the combination is -; rejected here rather than being honored in one direction without saying so. -; -; The check sits next to the other reasons a musttail call cannot be honored, -; such as inline asm, because it is the same kind of conflict. The frontend -; diagnostic for the same conflict in source is separate: -; trailofbits/vspells-ct-internal-notes#22. +; A musttail call must replace the caller, leaving no point to clear its frame. +; Reject it in protected functions; ordinary tail-call optimization is suppressed. ; RUN: not llvm-as < %s -o /dev/null 2>&1 | FileCheck %s diff --git a/llvm/unittests/CodeGen/CMakeLists.txt b/llvm/unittests/CodeGen/CMakeLists.txt index 709017380fa4e5..75e8ddf231a275 100644 --- a/llvm/unittests/CodeGen/CMakeLists.txt +++ b/llvm/unittests/CodeGen/CMakeLists.txt @@ -40,6 +40,7 @@ add_llvm_unittest(CodeGenTests MIR2VecTest.cpp RegAllocBasicTest.cpp RegAllocScoreTest.cpp + RegisterClearingTest.cpp RegisterTest.cpp PassManagerTest.cpp RematerializerTest.cpp diff --git a/llvm/unittests/CodeGen/RegisterClearingTest.cpp b/llvm/unittests/CodeGen/RegisterClearingTest.cpp new file mode 100644 index 00000000000000..92cb12230c1da2 --- /dev/null +++ b/llvm/unittests/CodeGen/RegisterClearingTest.cpp @@ -0,0 +1,303 @@ +//===- RegisterClearingTest.cpp --------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "../../lib/CodeGen/RegisterClearing.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/CodeGen/MachineFunction.h" +#include "llvm/CodeGen/MachineInstrBuilder.h" +#include "llvm/CodeGen/MachineModuleInfo.h" +#include "llvm/CodeGen/MachineRegisterInfo.h" +#include "llvm/CodeGen/TargetInstrInfo.h" +#include "llvm/CodeGen/TargetRegisterInfo.h" +#include "llvm/CodeGen/TargetSubtargetInfo.h" +#include "llvm/IR/DiagnosticInfo.h" +#include "llvm/IR/DiagnosticPrinter.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Support/TargetSelect.h" +#include "llvm/Target/TargetMachine.h" +#include "gmock/gmock.h" +#include "gtest/gtest.h" + +using namespace llvm; +using testing::HasSubstr; + +namespace { +struct ClearingTarget { + const char *Name; + const char *Triple; + const char *Features; + const char *ReturnOpcode; + const char *ClearOpcode; + const char *Scratch; + const char *Candidate; + const char *ReturnValue; + const char *FullReturnValue; + const char *Reserved; + const char *CalleeSaved; + const char *Unsupported; + const char *Partial; +}; + +class RegisterClearingTest : public testing::TestWithParam { +protected: + LLVMContext Context; + std::unique_ptr TM; + std::unique_ptr Mod; + std::unique_ptr MMI; + MachineFunction *MF = nullptr; + const TargetRegisterInfo *TRI = nullptr; + const TargetInstrInfo *TII = nullptr; + std::string Diagnostics; + unsigned Errors = 0; + + static void SetUpTestSuite() { + InitializeAllTargetInfos(); + InitializeAllTargets(); + InitializeAllTargetMCs(); + } + + void SetUp() override { + Triple TT(GetParam().Triple); + std::string Error; + const Target *T = TargetRegistry::lookupTarget(TT, Error); + if (!T) + GTEST_SKIP() << Error; + TM.reset(T->createTargetMachine(TT, "", GetParam().Features, + TargetOptions(), std::nullopt)); + ASSERT_NE(TM, nullptr); + Mod = std::make_unique("clearing", Context); + Mod->setDataLayout(TM->createDataLayout()); + auto *F = + Function::Create(FunctionType::get(Type::getVoidTy(Context), false), + GlobalValue::ExternalLinkage, "test", *Mod); + ReturnInst::Create(Context, BasicBlock::Create(Context, "entry", F)); + F->addFnAttr("zeroize-stack", "used"); + MMI = std::make_unique(TM.get()); + MF = &MMI->getOrCreateMachineFunction(*F); + MF->getProperties().set(MachineFunctionProperties::Property::NoVRegs); + MF->getRegInfo().freezeReservedRegs(); + TRI = MF->getSubtarget().getRegisterInfo(); + TII = MF->getSubtarget().getInstrInfo(); + Context.setDiagnosticHandlerCallBack( + [](const DiagnosticInfo *DI, void *Opaque) { + auto &Self = *static_cast(Opaque); + raw_string_ostream OS(Self.Diagnostics); + DiagnosticPrinterRawOStream Printer(OS); + DI->print(Printer); + Self.Errors += DI->getSeverity() == DS_Error; + }, + this); + } + + MCRegister reg(StringRef Name) { + for (unsigned I = 1; I != TRI->getNumRegs(); ++I) + if (Name == TRI->getName(I)) + return MCRegister(I); + ADD_FAILURE() << "Unknown register " << Name.str(); + return MCRegister(); + } + + unsigned opcode(StringRef Name) { + for (unsigned I = 0; I != TII->getNumOpcodes(); ++I) + if (Name == TII->getName(I)) + return I; + ADD_FAILURE() << "Unknown opcode " << Name.str(); + return 0; + } + + MachineBasicBlock &addExit() { + auto *MBB = MF->CreateMachineBasicBlock(); + MF->push_back(MBB); + auto Ret = BuildMI(*MBB, MBB->end(), DebugLoc(), + TII->get(opcode(GetParam().ReturnOpcode))); + if (TM->getTargetTriple().isX86()) + Ret.addImm(0); + Ret.addReg(reg(GetParam().ReturnValue), RegState::Implicit); + return *MBB; + } + + BitVector regs(std::initializer_list Registers) { + BitVector Result(TRI->getNumRegs()); + for (MCRegister Reg : Registers) + Result.set(Reg); + return Result; + } + + bool clear(MachineBasicBlock &MBB, const BitVector &Candidates, + const BitVector &Scratch) { + return emitZeroCallUsedRegsWithScratch(Candidates, Scratch, MBB, + MBB.getFirstTerminator(), nullptr); + } + + void expectClears(MachineBasicBlock &MBB, + std::initializer_list Expected) { + EXPECT_EQ(MBB.size(), Expected.size() + 1); + for (MCRegister Reg : Expected) { + unsigned Count = 0; + for (auto I = MBB.begin(), E = MBB.getFirstTerminator(); I != E; ++I) { + EXPECT_EQ(I->getOpcode(), opcode(GetParam().ClearOpcode)); + EXPECT_FALSE(I->modifiesRegister(reg(GetParam().ReturnValue), TRI)); + for (const MachineOperand &MO : I->explicit_operands()) + if (MO.isImm()) + EXPECT_EQ(MO.getImm(), 0); + if (TM->getTargetTriple().isX86()) { + EXPECT_EQ(I->getOperand(0).getReg(), I->getOperand(1).getReg()); + EXPECT_EQ(I->getOperand(0).getReg(), I->getOperand(2).getReg()); + } + Count += I->modifiesRegister(Reg, TRI); + } + EXPECT_EQ(Count, 1u); + } + EXPECT_EQ(Errors, 0u) << Diagnostics; + } + + void expectRejected(MCRegister Reg) { + auto &MBB = addExit(); + unsigned PreviousErrors = Errors; + EXPECT_FALSE(clear(MBB, regs({}), regs({Reg}))); + EXPECT_EQ(MBB.size(), 1u); + EXPECT_EQ(Errors, PreviousErrors + 1); + } +}; + +TEST_P(RegisterClearingTest, ScratchWithoutRegisterAttribute) { + auto &MBB = addExit(); + MCRegister Scratch = reg(GetParam().Scratch); + ASSERT_TRUE(clear(MBB, regs({}), regs({Scratch}))); + expectClears(MBB, {Scratch}); +} + +TEST_P(RegisterClearingTest, ScratchWithSkip) { + MF->getFunction().addFnAttr("zero-call-used-regs", "skip"); + auto &MBB = addExit(); + MCRegister Scratch = reg(GetParam().Scratch); + ASSERT_TRUE(clear(MBB, regs({}), regs({Scratch}))); + expectClears(MBB, {Scratch}); +} + +TEST_P(RegisterClearingTest, ScratchAddedToFilteredCandidates) { + MF->getFunction().addFnAttr("zero-call-used-regs", "used-gpr"); + auto &MBB = addExit(); + MCRegister Scratch = reg(GetParam().Scratch); + MCRegister Candidate = reg(GetParam().Candidate); + ASSERT_TRUE(clear(MBB, regs({Candidate}), regs({Scratch}))); + expectClears(MBB, {Candidate, Scratch}); +} + +TEST_P(RegisterClearingTest, DuplicateRegisterClearedOnce) { + auto &MBB = addExit(); + MCRegister Scratch = reg(GetParam().Scratch); + ASSERT_TRUE(clear(MBB, regs({Scratch}), regs({Scratch}))); + expectClears(MBB, {Scratch}); +} + +TEST_P(RegisterClearingTest, IndependentExitDeclarations) { + auto &First = addExit(); + auto &Second = addExit(); + MCRegister A = reg(GetParam().Scratch); + MCRegister B = reg(GetParam().Candidate); + ASSERT_TRUE(clear(First, regs({}), regs({A}))); + ASSERT_TRUE(clear(Second, regs({}), regs({B}))); + expectClears(First, {A}); + expectClears(Second, {B}); +} + +TEST_P(RegisterClearingTest, EmptyScratchLeavesOnlyCandidates) { + auto &MBB = addExit(); + MCRegister Candidate = reg(GetParam().Candidate); + ASSERT_TRUE(clear(MBB, regs({Candidate}), regs({}))); + expectClears(MBB, {Candidate}); +} + +TEST_P(RegisterClearingTest, ReservedRegisterRejected) { + expectRejected(reg(GetParam().Reserved)); + EXPECT_THAT(Diagnostics, HasSubstr("not supported for scratch")); +} + +TEST_P(RegisterClearingTest, InvalidScratchPreventsCandidateClearing) { + auto &MBB = addExit(); + EXPECT_FALSE(clear(MBB, regs({reg(GetParam().Candidate)}), + regs({reg(GetParam().Unsupported)}))); + EXPECT_EQ(MBB.size(), 1u); + EXPECT_EQ(Errors, 1u); +} + +TEST_P(RegisterClearingTest, LiveReturnRejected) { + expectRejected(reg(GetParam().FullReturnValue)); + EXPECT_THAT(Diagnostics, HasSubstr("needed at the exit")); +} + +TEST_P(RegisterClearingTest, CalleeSavedRegisterRejected) { + expectRejected(reg(GetParam().CalleeSaved)); +} + +TEST_P(RegisterClearingTest, UpdatedCalleeSavedRegisterRejected) { + // Call lowering can extend the target's base CSR list, for example for + // AArch64's +call-saved-x9 feature. + SmallVector CSRs; + for (const MCPhysReg *CSR = MF->getRegInfo().getCalleeSavedRegs(); *CSR; + ++CSR) + CSRs.push_back(*CSR); + MCRegister Scratch = reg(GetParam().Scratch); + CSRs.push_back(Scratch); + MF->getRegInfo().setCalleeSavedRegs(CSRs); + expectRejected(Scratch); + EXPECT_THAT(Diagnostics, HasSubstr("needed at the exit")); +} + +TEST_P(RegisterClearingTest, PartialRegisterRejected) { + // AH is disjoint from the live AL return, but XOR32rr would overwrite AL. + expectRejected(reg(GetParam().Partial)); + EXPECT_THAT(Diagnostics, HasSubstr("not supported for scratch")); +} + +TEST_P(RegisterClearingTest, UnsupportedRegisterRejected) { + expectRejected(reg(GetParam().Unsupported)); + EXPECT_THAT(Diagnostics, HasSubstr("not supported for scratch")); + if (TM->getTargetTriple().isAArch64()) { + // The emitter skips X19-X30 even when the calling convention preserves + // none. + for (auto CC : {CallingConv::PreserveNone, CallingConv::GHC}) { + MF->getFunction().setCallingConv(CC); + for (StringRef Name : {"X19", "X28", "FP"}) + expectRejected(reg(Name)); + } + } +} + +TEST_P(RegisterClearingTest, ReturnAddressRejected) { + if (TM->getTargetTriple().isRISCV()) + // GHC preserves no registers; PseudoRET does not name the return address. + MF->getFunction().setCallingConv(CallingConv::GHC); + expectRejected(TRI->getRARegister()); +} + +INSTANTIATE_TEST_SUITE_P( + Targets, RegisterClearingTest, + testing::Values(ClearingTarget{"X86_64", "x86_64-linux-gnu", "+mmx", + "RET64", "XOR32rr", "R11", "R10", "AL", + "RAX", "RSP", "RBX", "MM0", "AH"}, + ClearingTarget{"X86_32", "i386-linux-gnu", "+mmx", "RET32", + "XOR32rr", "ECX", "EDX", "AL", "EAX", "ESP", + "EBX", "MM0", "AH"}, + ClearingTarget{"AArch64", "aarch64-linux-gnu", "", + "RET_ReallyLR", "MOVZXi", "X9", "X10", "W0", + "X0", "SP", "X19", "D9", "W9"}, + ClearingTarget{"RISCV32", "riscv32-linux-gnu", "+d", + "PseudoRET", "PseudoClearGPR", "X5", "X6", + "X10", "X10", "X2", "X9", "V0", "X5_H"}, + ClearingTarget{"RISCV64", "riscv64-linux-gnu", "+d", + "PseudoRET", "PseudoClearGPR", "X5", "X6", + "X10", "X10", "X2", "X9", "V0", "X5_H"}), + [](const testing::TestParamInfo &Info) { + return Info.param.Name; + }); +} // namespace diff --git a/llvm/utils/gn/secondary/llvm/lib/CodeGen/BUILD.gn b/llvm/utils/gn/secondary/llvm/lib/CodeGen/BUILD.gn index 72c1c0ff844678..beb56dd962c43b 100644 --- a/llvm/utils/gn/secondary/llvm/lib/CodeGen/BUILD.gn +++ b/llvm/utils/gn/secondary/llvm/lib/CodeGen/BUILD.gn @@ -204,6 +204,7 @@ static_library("CodeGen") { "RegisterBank.cpp", "RegisterBankInfo.cpp", "RegisterClassInfo.cpp", + "RegisterClearing.cpp", "RegisterCoalescer.cpp", "RegisterPressure.cpp", "RegisterScavenging.cpp", diff --git a/llvm/utils/gn/secondary/llvm/unittests/CodeGen/BUILD.gn b/llvm/utils/gn/secondary/llvm/unittests/CodeGen/BUILD.gn index 770ba8fdabedde..b0348e9918d43d 100644 --- a/llvm/utils/gn/secondary/llvm/unittests/CodeGen/BUILD.gn +++ b/llvm/utils/gn/secondary/llvm/unittests/CodeGen/BUILD.gn @@ -40,6 +40,7 @@ unittest("CodeGenTests") { "PassManagerTest.cpp", "RegAllocBasicTest.cpp", "RegAllocScoreTest.cpp", + "RegisterClearingTest.cpp", "RegisterTest.cpp", "RematerializerTest.cpp", "ScalableVectorMVTsTest.cpp",