1 | //===--------- SMEABI - SME ABI-------------------------------------------===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | // |
9 | // This pass implements parts of the the SME ABI, such as: |
10 | // * Using the lazy-save mechanism before enabling the use of ZA. |
11 | // * Setting up the lazy-save mechanism around invokes. |
12 | // |
13 | //===----------------------------------------------------------------------===// |
14 | |
15 | #include "AArch64.h" |
16 | #include "Utils/AArch64BaseInfo.h" |
17 | #include "Utils/AArch64SMEAttributes.h" |
18 | #include "llvm/ADT/StringRef.h" |
19 | #include "llvm/IR/Constants.h" |
20 | #include "llvm/IR/IRBuilder.h" |
21 | #include "llvm/IR/Instructions.h" |
22 | #include "llvm/IR/IntrinsicInst.h" |
23 | #include "llvm/IR/IntrinsicsAArch64.h" |
24 | #include "llvm/IR/LLVMContext.h" |
25 | #include "llvm/InitializePasses.h" |
26 | #include "llvm/Support/Debug.h" |
27 | #include "llvm/Transforms/Utils/Cloning.h" |
28 | |
29 | using namespace llvm; |
30 | |
31 | #define DEBUG_TYPE "aarch64-sme-abi" |
32 | |
33 | namespace { |
34 | struct SMEABI : public FunctionPass { |
35 | static char ID; // Pass identification, replacement for typeid |
36 | SMEABI() : FunctionPass(ID) { |
37 | initializeSMEABIPass(*PassRegistry::getPassRegistry()); |
38 | } |
39 | |
40 | bool runOnFunction(Function &F) override; |
41 | |
42 | private: |
43 | bool updateNewStateFunctions(Module *M, Function *F, IRBuilder<> &Builder, |
44 | SMEAttrs FnAttrs); |
45 | }; |
46 | } // end anonymous namespace |
47 | |
48 | char SMEABI::ID = 0; |
49 | static const char *name = "SME ABI Pass" ; |
50 | INITIALIZE_PASS_BEGIN(SMEABI, DEBUG_TYPE, name, false, false) |
51 | INITIALIZE_PASS_END(SMEABI, DEBUG_TYPE, name, false, false) |
52 | |
53 | FunctionPass *llvm::createSMEABIPass() { return new SMEABI(); } |
54 | |
55 | //===----------------------------------------------------------------------===// |
56 | // Utility functions |
57 | //===----------------------------------------------------------------------===// |
58 | |
59 | // Utility function to emit a call to __arm_tpidr2_save and clear TPIDR2_EL0. |
60 | void emitTPIDR2Save(Module *M, IRBuilder<> &Builder) { |
61 | auto *TPIDR2SaveTy = |
62 | FunctionType::get(Result: Builder.getVoidTy(), Params: {}, /*IsVarArgs=*/isVarArg: false); |
63 | auto Attrs = AttributeList().addFnAttribute(C&: M->getContext(), |
64 | Kind: "aarch64_pstate_sm_compatible" ); |
65 | FunctionCallee Callee = |
66 | M->getOrInsertFunction(Name: "__arm_tpidr2_save" , T: TPIDR2SaveTy, AttributeList: Attrs); |
67 | CallInst *Call = Builder.CreateCall(Callee); |
68 | Call->setCallingConv( |
69 | CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost_From_X0); |
70 | |
71 | // A save to TPIDR2 should be followed by clearing TPIDR2_EL0. |
72 | Function *WriteIntr = |
73 | Intrinsic::getDeclaration(M, Intrinsic::id: aarch64_sme_set_tpidr2); |
74 | Builder.CreateCall(FTy: WriteIntr->getFunctionType(), Callee: WriteIntr, |
75 | Args: Builder.getInt64(C: 0)); |
76 | } |
77 | |
78 | /// This function generates code at the beginning and end of a function marked |
79 | /// with either `aarch64_new_za` or `aarch64_new_zt0`. |
80 | /// At the beginning of the function, the following code is generated: |
81 | /// - Commit lazy-save if active [Private-ZA Interface*] |
82 | /// - Enable PSTATE.ZA [Private-ZA Interface] |
83 | /// - Zero ZA [Has New ZA State] |
84 | /// - Zero ZT0 [Has New ZT0 State] |
85 | /// |
86 | /// * A function with new ZT0 state will not change ZA, so committing the |
87 | /// lazy-save is not strictly necessary. However, the lazy-save mechanism |
88 | /// may be active on entry to the function, with PSTATE.ZA set to 1. If |
89 | /// the new ZT0 function calls a function that does not share ZT0, we will |
90 | /// need to conditionally SMSTOP ZA before the call, setting PSTATE.ZA to 0. |
91 | /// For this reason, it's easier to always commit the lazy-save at the |
92 | /// beginning of the function regardless of whether it has ZA state. |
93 | /// |
94 | /// At the end of the function, PSTATE.ZA is disabled if the function has a |
95 | /// Private-ZA Interface. A function is considered to have a Private-ZA |
96 | /// interface if it does not share ZA or ZT0. |
97 | /// |
98 | bool SMEABI::updateNewStateFunctions(Module *M, Function *F, |
99 | IRBuilder<> &Builder, SMEAttrs FnAttrs) { |
100 | LLVMContext &Context = F->getContext(); |
101 | BasicBlock *OrigBB = &F->getEntryBlock(); |
102 | Builder.SetInsertPoint(&OrigBB->front()); |
103 | |
104 | // Commit any active lazy-saves if this is a Private-ZA function. If the |
105 | // value read from TPIDR2_EL0 is not null on entry to the function then |
106 | // the lazy-saving scheme is active and we should call __arm_tpidr2_save |
107 | // to commit the lazy save. |
108 | if (FnAttrs.hasPrivateZAInterface()) { |
109 | // Create the new blocks for reading TPIDR2_EL0 & enabling ZA state. |
110 | auto *SaveBB = OrigBB->splitBasicBlock(I: OrigBB->begin(), BBName: "save.za" , Before: true); |
111 | auto *PreludeBB = BasicBlock::Create(Context, Name: "prelude" , Parent: F, InsertBefore: SaveBB); |
112 | |
113 | // Read TPIDR2_EL0 in PreludeBB & branch to SaveBB if not 0. |
114 | Builder.SetInsertPoint(PreludeBB); |
115 | Function *TPIDR2Intr = |
116 | Intrinsic::getDeclaration(M, Intrinsic::id: aarch64_sme_get_tpidr2); |
117 | auto *TPIDR2 = Builder.CreateCall(FTy: TPIDR2Intr->getFunctionType(), Callee: TPIDR2Intr, |
118 | Args: {}, Name: "tpidr2" ); |
119 | auto *Cmp = Builder.CreateCmp(Pred: ICmpInst::ICMP_NE, LHS: TPIDR2, |
120 | RHS: Builder.getInt64(C: 0), Name: "cmp" ); |
121 | Builder.CreateCondBr(Cmp, SaveBB, OrigBB); |
122 | |
123 | // Create a call __arm_tpidr2_save, which commits the lazy save. |
124 | Builder.SetInsertPoint(&SaveBB->back()); |
125 | emitTPIDR2Save(M, Builder); |
126 | |
127 | // Enable pstate.za at the start of the function. |
128 | Builder.SetInsertPoint(&OrigBB->front()); |
129 | Function *EnableZAIntr = |
130 | Intrinsic::getDeclaration(M, Intrinsic::id: aarch64_sme_za_enable); |
131 | Builder.CreateCall(FTy: EnableZAIntr->getFunctionType(), Callee: EnableZAIntr); |
132 | } |
133 | |
134 | if (FnAttrs.isNewZA()) { |
135 | Function *ZeroIntr = |
136 | Intrinsic::getDeclaration(M, Intrinsic::id: aarch64_sme_zero); |
137 | Builder.CreateCall(FTy: ZeroIntr->getFunctionType(), Callee: ZeroIntr, |
138 | Args: Builder.getInt32(C: 0xff)); |
139 | } |
140 | |
141 | if (FnAttrs.isNewZT0()) { |
142 | Function *ClearZT0Intr = |
143 | Intrinsic::getDeclaration(M, Intrinsic::id: aarch64_sme_zero_zt); |
144 | Builder.CreateCall(FTy: ClearZT0Intr->getFunctionType(), Callee: ClearZT0Intr, |
145 | Args: {Builder.getInt32(C: 0)}); |
146 | } |
147 | |
148 | if (FnAttrs.hasPrivateZAInterface()) { |
149 | // Before returning, disable pstate.za |
150 | for (BasicBlock &BB : *F) { |
151 | Instruction *T = BB.getTerminator(); |
152 | if (!T || !isa<ReturnInst>(Val: T)) |
153 | continue; |
154 | Builder.SetInsertPoint(T); |
155 | Function *DisableZAIntr = |
156 | Intrinsic::getDeclaration(M, Intrinsic::id: aarch64_sme_za_disable); |
157 | Builder.CreateCall(FTy: DisableZAIntr->getFunctionType(), Callee: DisableZAIntr); |
158 | } |
159 | } |
160 | |
161 | F->addFnAttr(Kind: "aarch64_expanded_pstate_za" ); |
162 | return true; |
163 | } |
164 | |
165 | bool SMEABI::runOnFunction(Function &F) { |
166 | Module *M = F.getParent(); |
167 | LLVMContext &Context = F.getContext(); |
168 | IRBuilder<> Builder(Context); |
169 | |
170 | if (F.isDeclaration() || F.hasFnAttribute(Kind: "aarch64_expanded_pstate_za" )) |
171 | return false; |
172 | |
173 | bool Changed = false; |
174 | SMEAttrs FnAttrs(F); |
175 | if (FnAttrs.isNewZA() || FnAttrs.isNewZT0()) |
176 | Changed |= updateNewStateFunctions(M, F: &F, Builder, FnAttrs); |
177 | |
178 | return Changed; |
179 | } |
180 | |