Skip to content

Commit abfef3b

Browse files
MacDueMacDue
committed
[AArch64][SME] Avoid ZA save state changes in loops in MachineSMEABIPass
This patch uses the MachineLoopInfo to give blocks within loops a higher weight when choosing the bundle ZA state. MachineLoopInfo does not find loop trip counts, so this uses an arbitrary weight (default 10), which can be configured with the `-aarch64-sme-abi-loop-edge-weight` flag. This makes the MachineSMEABIPass pass more likely to pick a bundle state that matches the loop's entry/exit state, which avoids state changes in the loop (which we assume will happen more than once). This does require some extra analysis, so this is only enabled at -O1 and above. Change-Id: If318c809d2f7cc1fca144fbe424ba2a2ca7fb19f
1 parent de812ce commit abfef3b

File tree

5 files changed

+160
-27
lines changed

5 files changed

+160
-27
lines changed

‎llvm/lib/Target/AArch64/AArch64.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ FunctionPass *createAArch64CleanupLocalDynamicTLSPass();
6060
FunctionPass *createAArch64CollectLOHPass();
6161
FunctionPass *createSMEABIPass();
6262
FunctionPass *createSMEPeepholeOptPass();
63-
FunctionPass *createMachineSMEABIPass();
63+
FunctionPass *createMachineSMEABIPass(CodeGenOptLevel);
6464
ModulePass *createSVEIntrinsicOptsPass();
6565
InstructionSelector *
6666
createAArch64InstructionSelector(const AArch64TargetMachine &,

‎llvm/lib/Target/AArch64/AArch64TargetMachine.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -791,8 +791,8 @@ bool AArch64PassConfig::addGlobalInstructionSelect() {
791791
}
792792

793793
void AArch64PassConfig::addMachineSSAOptimization() {
794-
if (EnableNewSMEABILowering && TM->getOptLevel() != CodeGenOptLevel::None)
795-
addPass(createMachineSMEABIPass());
794+
if (TM->getOptLevel() != CodeGenOptLevel::None && EnableNewSMEABILowering)
795+
addPass(createMachineSMEABIPass(TM->getOptLevel()));
796796

797797
if (TM->getOptLevel() != CodeGenOptLevel::None && EnableSMEPeepholeOpt)
798798
addPass(createSMEPeepholeOptPass());
@@ -825,7 +825,7 @@ bool AArch64PassConfig::addILPOpts() {
825825

826826
void AArch64PassConfig::addPreRegAlloc() {
827827
if (TM->getOptLevel() == CodeGenOptLevel::None && EnableNewSMEABILowering)
828-
addPass(createMachineSMEABIPass());
828+
addPass(createMachineSMEABIPass(CodeGenOptLevel::None));
829829

830830
// Change dead register definitions to refer to the zero register.
831831
if (TM->getOptLevel() != CodeGenOptLevel::None &&

‎llvm/lib/Target/AArch64/MachineSMEABIPass.cpp‎

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -63,13 +63,20 @@
6363
#include "llvm/CodeGen/LivePhysRegs.h"
6464
#include "llvm/CodeGen/MachineBasicBlock.h"
6565
#include "llvm/CodeGen/MachineFunctionPass.h"
66+
#include "llvm/CodeGen/MachineLoopInfo.h"
6667
#include "llvm/CodeGen/MachineRegisterInfo.h"
6768
#include "llvm/CodeGen/TargetRegisterInfo.h"
6869

6970
using namespace llvm;
7071

7172
#define DEBUG_TYPE "aarch64-machine-sme-abi"
7273

74+
static cl::opt<int>
75+
LoopEdgeWeight("aarch64-sme-abi-loop-edge-weight", cl::ReallyHidden,
76+
cl::init(10),
77+
cl::desc("Edge weight for basic blocks witin loops (used "
78+
"for placing ZA saves/restores)"));
79+
7380
namespace {
7481

7582
enum ZAState {
@@ -176,7 +183,8 @@ getZAStateBeforeInst(const TargetRegisterInfo &TRI, MachineInstr &MI,
176183
struct MachineSMEABI : public MachineFunctionPass {
177184
inline static char ID = 0;
178185

179-
MachineSMEABI() : MachineFunctionPass(ID) {}
186+
MachineSMEABI(CodeGenOptLevel OptLevel = CodeGenOptLevel::Default)
187+
: MachineFunctionPass(ID), OptLevel(OptLevel) {}
180188

181189
bool runOnMachineFunction(MachineFunction &MF) override;
182190

@@ -185,6 +193,9 @@ struct MachineSMEABI : public MachineFunctionPass {
185193
void getAnalysisUsage(AnalysisUsage &AU) const override {
186194
AU.setPreservesCFG();
187195
AU.addRequired<EdgeBundlesWrapperLegacy>();
196+
// Only analyse loops at -01 and above.
197+
if (OptLevel != CodeGenOptLevel::None)
198+
AU.addRequired<MachineLoopInfoWrapperPass>();
188199
AU.addPreservedID(MachineLoopInfoID);
189200
AU.addPreservedID(MachineDominatorsID);
190201
MachineFunctionPass::getAnalysisUsage(AU);
@@ -280,6 +291,8 @@ struct MachineSMEABI : public MachineFunctionPass {
280291
LiveRegs PhysLiveRegsAtExit = LiveRegs::None;
281292
};
282293

294+
CodeGenOptLevel OptLevel = CodeGenOptLevel::Default;
295+
283296
// All pass state that must be cleared between functions.
284297
struct PassState {
285298
SmallVector<BlockInfo> Blocks;
@@ -297,6 +310,7 @@ struct MachineSMEABI : public MachineFunctionPass {
297310
const AArch64RegisterInfo *TRI = nullptr;
298311
const TargetInstrInfo *TII = nullptr;
299312
MachineRegisterInfo *MRI = nullptr;
313+
MachineLoopInfo *MLI = nullptr;
300314
};
301315

302316
void MachineSMEABI::collectNeededZAStates(SMEAttrs SMEFnAttrs) {
@@ -387,18 +401,23 @@ void MachineSMEABI::assignBundleZAStates() {
387401
LLVM_DEBUG(dbgs() << " (no state preference)\n");
388402
continue;
389403
}
404+
bool IsLoop = MLI && MLI->getLoopFor(MF->getBlockNumbered(BlockID));
390405
bool InEdge = Bundles->getBundle(BlockID, /*Out=*/false) == I;
391406
bool OutEdge = Bundles->getBundle(BlockID, /*Out=*/true) == I;
407+
int EdgeWeight = IsLoop ? LoopEdgeWeight : 1;
408+
if (IsLoop)
409+
LLVM_DEBUG(dbgs() << " IsLoop");
392410

411+
LLVM_DEBUG(dbgs() << " (EdgeWeight: " << EdgeWeight << ')');
393412
ZAState DesiredIncomingState = Block.Insts.front().NeededState;
394413
if (InEdge && isLegalEdgeBundleZAState(DesiredIncomingState)) {
395-
EdgeStateCounts[DesiredIncomingState]++;
414+
EdgeStateCounts[DesiredIncomingState] += EdgeWeight;
396415
LLVM_DEBUG(dbgs() << " DesiredIncomingState: "
397416
<< getZAStateString(DesiredIncomingState));
398417
}
399418
ZAState DesiredOutgoingState = Block.Insts.back().NeededState;
400419
if (OutEdge && isLegalEdgeBundleZAState(DesiredOutgoingState)) {
401-
EdgeStateCounts[DesiredOutgoingState]++;
420+
EdgeStateCounts[DesiredOutgoingState] += EdgeWeight;
402421
LLVM_DEBUG(dbgs() << " DesiredOutgoingState: "
403422
<< getZAStateString(DesiredOutgoingState));
404423
}
@@ -823,6 +842,8 @@ bool MachineSMEABI::runOnMachineFunction(MachineFunction &MF) {
823842
TII = Subtarget->getInstrInfo();
824843
TRI = Subtarget->getRegisterInfo();
825844
MRI = &MF.getRegInfo();
845+
if (OptLevel != CodeGenOptLevel::None)
846+
MLI = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
826847

827848
bool IsAgnosticZA = SMEFnAttrs.hasAgnosticZAInterface();
828849

@@ -851,4 +872,6 @@ bool MachineSMEABI::runOnMachineFunction(MachineFunction &MF) {
851872
return true;
852873
}
853874

854-
FunctionPass *llvm::createMachineSMEABIPass() { return new MachineSMEABI(); }
875+
FunctionPass *llvm::createMachineSMEABIPass(CodeGenOptLevel OptLevel) {
876+
return new MachineSMEABI(OptLevel);
877+
}

‎llvm/test/CodeGen/AArch64/sme-lazy-save-in-loop.ll‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
2+
; RUN: llc -O0 -mtriple=aarch64-linux-gnu -mattr=+sme -aarch64-new-sme-abi < %s | FileCheck %s --check-prefix=CHECK-O0
3+
; RUN: llc -O1 -mtriple=aarch64-linux-gnu -mattr=+sme -aarch64-new-sme-abi < %s | FileCheck %s --check-prefix=CHECK-O1
4+
5+
declare void @private_za_call()
6+
declare void @shared_za_call() "aarch64_inout_za"
7+
8+
; This test checks that at -O0 we don't attempt to optimize lazy save state
9+
; changes in loops, and that -O1 (and above) we attempt to push state changes
10+
; out of loops.
11+
12+
define void @private_za_loop_active_entry_and_exit(i32 %n) "aarch64_inout_za" nounwind {
13+
; CHECK-O0-LABEL: private_za_loop_active_entry_and_exit:
14+
; CHECK-O0: // %bb.0: // %entry
15+
; CHECK-O0-NEXT: stp x29, x30, [sp, #-16]! // 16-byte Folded Spill
16+
; CHECK-O0-NEXT: mov x29, sp
17+
; CHECK-O0-NEXT: sub sp, sp, #32
18+
; CHECK-O0-NEXT: rdsvl x9, #1
19+
; CHECK-O0-NEXT: mov x8, sp
20+
; CHECK-O0-NEXT: msub x8, x9, x9, x8
21+
; CHECK-O0-NEXT: mov sp, x8
22+
; CHECK-O0-NEXT: stp x8, x9, [x29, #-16]
23+
; CHECK-O0-NEXT: stur w0, [x29, #-24] // 4-byte Folded Spill
24+
; CHECK-O0-NEXT: bl shared_za_call
25+
; CHECK-O0-NEXT: ldur w0, [x29, #-24] // 4-byte Folded Reload
26+
; CHECK-O0-NEXT: mov w8, wzr
27+
; CHECK-O0-NEXT: subs w9, w0, #1
28+
; CHECK-O0-NEXT: stur w8, [x29, #-20] // 4-byte Folded Spill
29+
; CHECK-O0-NEXT: b.lt .LBB0_4
30+
; CHECK-O0-NEXT: b .LBB0_1
31+
; CHECK-O0-NEXT: .LBB0_1: // %loop
32+
; CHECK-O0-NEXT: // =>This Inner Loop Header: Depth=1
33+
; CHECK-O0-NEXT: ldur w8, [x29, #-20] // 4-byte Folded Reload
34+
; CHECK-O0-NEXT: stur w8, [x29, #-28] // 4-byte Folded Spill
35+
; CHECK-O0-NEXT: sub x8, x29, #16
36+
; CHECK-O0-NEXT: msr TPIDR2_EL0, x8
37+
; CHECK-O0-NEXT: bl private_za_call
38+
; CHECK-O0-NEXT: ldur w8, [x29, #-28] // 4-byte Folded Reload
39+
; CHECK-O0-NEXT: ldur w10, [x29, #-24] // 4-byte Folded Reload
40+
; CHECK-O0-NEXT: add w9, w8, #1
41+
; CHECK-O0-NEXT: mov w8, w9
42+
; CHECK-O0-NEXT: subs w9, w9, w10
43+
; CHECK-O0-NEXT: mrs x9, NZCV
44+
; CHECK-O0-NEXT: smstart za
45+
; CHECK-O0-NEXT: mrs x10, TPIDR2_EL0
46+
; CHECK-O0-NEXT: sub x0, x29, #16
47+
; CHECK-O0-NEXT: cbz x10, .LBB0_2
48+
; CHECK-O0-NEXT: b .LBB0_3
49+
; CHECK-O0-NEXT: .LBB0_2: // %loop
50+
; CHECK-O0-NEXT: // in Loop: Header=BB0_1 Depth=1
51+
; CHECK-O0-NEXT: bl __arm_tpidr2_restore
52+
; CHECK-O0-NEXT: b .LBB0_3
53+
; CHECK-O0-NEXT: .LBB0_3: // %loop
54+
; CHECK-O0-NEXT: // in Loop: Header=BB0_1 Depth=1
55+
; CHECK-O0-NEXT: msr TPIDR2_EL0, xzr
56+
; CHECK-O0-NEXT: msr NZCV, x9
57+
; CHECK-O0-NEXT: stur w8, [x29, #-20] // 4-byte Folded Spill
58+
; CHECK-O0-NEXT: b.ne .LBB0_1
59+
; CHECK-O0-NEXT: b .LBB0_4
60+
; CHECK-O0-NEXT: .LBB0_4: // %exit
61+
; CHECK-O0-NEXT: mov sp, x29
62+
; CHECK-O0-NEXT: ldp x29, x30, [sp], #16 // 16-byte Folded Reload
63+
; CHECK-O0-NEXT: b shared_za_call
64+
;
65+
; CHECK-O1-LABEL: private_za_loop_active_entry_and_exit:
66+
; CHECK-O1: // %bb.0: // %entry
67+
; CHECK-O1-NEXT: stp x29, x30, [sp, #-32]! // 16-byte Folded Spill
68+
; CHECK-O1-NEXT: str x19, [sp, #16] // 8-byte Folded Spill
69+
; CHECK-O1-NEXT: mov x29, sp
70+
; CHECK-O1-NEXT: sub sp, sp, #16
71+
; CHECK-O1-NEXT: rdsvl x8, #1
72+
; CHECK-O1-NEXT: mov x9, sp
73+
; CHECK-O1-NEXT: msub x9, x8, x8, x9
74+
; CHECK-O1-NEXT: mov sp, x9
75+
; CHECK-O1-NEXT: mov w19, w0
76+
; CHECK-O1-NEXT: stp x9, x8, [x29, #-16]
77+
; CHECK-O1-NEXT: bl shared_za_call
78+
; CHECK-O1-NEXT: cmp w19, #1
79+
; CHECK-O1-NEXT: sub x8, x29, #16
80+
; CHECK-O1-NEXT: msr TPIDR2_EL0, x8
81+
; CHECK-O1-NEXT: b.lt .LBB0_2
82+
; CHECK-O1-NEXT: .LBB0_1: // %loop
83+
; CHECK-O1-NEXT: // =>This Inner Loop Header: Depth=1
84+
; CHECK-O1-NEXT: bl private_za_call
85+
; CHECK-O1-NEXT: subs w19, w19, #1
86+
; CHECK-O1-NEXT: b.ne .LBB0_1
87+
; CHECK-O1-NEXT: .LBB0_2: // %exit
88+
; CHECK-O1-NEXT: smstart za
89+
; CHECK-O1-NEXT: mrs x8, TPIDR2_EL0
90+
; CHECK-O1-NEXT: sub x0, x29, #16
91+
; CHECK-O1-NEXT: cbnz x8, .LBB0_4
92+
; CHECK-O1-NEXT: // %bb.3: // %exit
93+
; CHECK-O1-NEXT: bl __arm_tpidr2_restore
94+
; CHECK-O1-NEXT: .LBB0_4: // %exit
95+
; CHECK-O1-NEXT: msr TPIDR2_EL0, xzr
96+
; CHECK-O1-NEXT: mov sp, x29
97+
; CHECK-O1-NEXT: ldr x19, [sp, #16] // 8-byte Folded Reload
98+
; CHECK-O1-NEXT: ldp x29, x30, [sp], #32 // 16-byte Folded Reload
99+
; CHECK-O1-NEXT: b shared_za_call
100+
entry:
101+
%cmpgt = icmp sgt i32 %n, 0
102+
tail call void @shared_za_call()
103+
br i1 %cmpgt, label %loop, label %exit
104+
105+
loop:
106+
%iv = phi i32 [ %next_iv, %loop ], [ 0, %entry ]
107+
tail call void @private_za_call()
108+
%next_iv = add nuw nsw i32 %iv, 1
109+
%cmpeq = icmp eq i32 %next_iv, %n
110+
br i1 %cmpeq, label %exit, label %loop
111+
112+
exit:
113+
tail call void @shared_za_call()
114+
ret void
115+
}

‎llvm/test/CodeGen/AArch64/sme-za-control-flow.ll‎

Lines changed: 14 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ exit:
9696
ret void
9797
}
9898

99-
; FIXME: In the new lowering we could weight edges to avoid doing the lazy save in the loop.
99+
; This tests that with the new lowering we push state changes out of loops (at -O1 and above).
100100
define void @private_za_loop_active_entry_and_exit(i32 %n) "aarch64_inout_za" nounwind {
101101
; CHECK-LABEL: private_za_loop_active_entry_and_exit:
102102
; CHECK: // %bb.0: // %entry
@@ -142,7 +142,7 @@ define void @private_za_loop_active_entry_and_exit(i32 %n) "aarch64_inout_za" no
142142
; CHECK-NEWLOWERING-LABEL: private_za_loop_active_entry_and_exit:
143143
; CHECK-NEWLOWERING: // %bb.0: // %entry
144144
; CHECK-NEWLOWERING-NEXT: stp x29, x30, [sp, #-32]! // 16-byte Folded Spill
145-
; CHECK-NEWLOWERING-NEXT: stp x20, x19, [sp, #16] // 16-byte Folded Spill
145+
; CHECK-NEWLOWERING-NEXT: str x19, [sp, #16] // 8-byte Folded Spill
146146
; CHECK-NEWLOWERING-NEXT: mov x29, sp
147147
; CHECK-NEWLOWERING-NEXT: sub sp, sp, #16
148148
; CHECK-NEWLOWERING-NEXT: rdsvl x8, #1
@@ -153,30 +153,25 @@ define void @private_za_loop_active_entry_and_exit(i32 %n) "aarch64_inout_za" no
153153
; CHECK-NEWLOWERING-NEXT: stp x9, x8, [x29, #-16]
154154
; CHECK-NEWLOWERING-NEXT: bl shared_za_call
155155
; CHECK-NEWLOWERING-NEXT: cmp w19, #1
156-
; CHECK-NEWLOWERING-NEXT: b.lt .LBB1_5
157-
; CHECK-NEWLOWERING-NEXT: // %bb.1: // %loop.preheader
158-
; CHECK-NEWLOWERING-NEXT: sub x20, x29, #16
159-
; CHECK-NEWLOWERING-NEXT: b .LBB1_3
160-
; CHECK-NEWLOWERING-NEXT: .LBB1_2: // %loop
161-
; CHECK-NEWLOWERING-NEXT: // in Loop: Header=BB1_3 Depth=1
162-
; CHECK-NEWLOWERING-NEXT: msr TPIDR2_EL0, xzr
163-
; CHECK-NEWLOWERING-NEXT: cbz w19, .LBB1_5
164-
; CHECK-NEWLOWERING-NEXT: .LBB1_3: // %loop
156+
; CHECK-NEWLOWERING-NEXT: sub x8, x29, #16
157+
; CHECK-NEWLOWERING-NEXT: msr TPIDR2_EL0, x8
158+
; CHECK-NEWLOWERING-NEXT: b.lt .LBB1_2
159+
; CHECK-NEWLOWERING-NEXT: .LBB1_1: // %loop
165160
; CHECK-NEWLOWERING-NEXT: // =>This Inner Loop Header: Depth=1
166-
; CHECK-NEWLOWERING-NEXT: msr TPIDR2_EL0, x20
167161
; CHECK-NEWLOWERING-NEXT: bl private_za_call
168-
; CHECK-NEWLOWERING-NEXT: sub w19, w19, #1
162+
; CHECK-NEWLOWERING-NEXT: subs w19, w19, #1
163+
; CHECK-NEWLOWERING-NEXT: b.ne .LBB1_1
164+
; CHECK-NEWLOWERING-NEXT: .LBB1_2: // %exit
169165
; CHECK-NEWLOWERING-NEXT: smstart za
170166
; CHECK-NEWLOWERING-NEXT: mrs x8, TPIDR2_EL0
171167
; CHECK-NEWLOWERING-NEXT: sub x0, x29, #16
172-
; CHECK-NEWLOWERING-NEXT: cbnz x8, .LBB1_2
173-
; CHECK-NEWLOWERING-NEXT: // %bb.4: // %loop
174-
; CHECK-NEWLOWERING-NEXT: // in Loop: Header=BB1_3 Depth=1
168+
; CHECK-NEWLOWERING-NEXT: cbnz x8, .LBB1_4
169+
; CHECK-NEWLOWERING-NEXT: // %bb.3: // %exit
175170
; CHECK-NEWLOWERING-NEXT: bl __arm_tpidr2_restore
176-
; CHECK-NEWLOWERING-NEXT: b .LBB1_2
177-
; CHECK-NEWLOWERING-NEXT: .LBB1_5: // %exit
171+
; CHECK-NEWLOWERING-NEXT: .LBB1_4: // %exit
172+
; CHECK-NEWLOWERING-NEXT: msr TPIDR2_EL0, xzr
178173
; CHECK-NEWLOWERING-NEXT: mov sp, x29
179-
; CHECK-NEWLOWERING-NEXT: ldp x20, x19, [sp, #16] // 16-byte Folded Reload
174+
; CHECK-NEWLOWERING-NEXT: ldr x19, [sp, #16] // 8-byte Folded Reload
180175
; CHECK-NEWLOWERING-NEXT: ldp x29, x30, [sp], #32 // 16-byte Folded Reload
181176
; CHECK-NEWLOWERING-NEXT: b shared_za_call
182177
entry:

0 commit comments

Comments
 (0)