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(Builder.getVoidTy(), {}, /*IsVarArgs=*/false);
63   auto Attrs =
64       AttributeList()
65           .addFnAttribute(M->getContext(), "aarch64_pstate_sm_compatible")
66           .addFnAttribute(M->getContext(), "aarch64_pstate_za_preserved");
67   FunctionCallee Callee =
68       M->getOrInsertFunction("__arm_tpidr2_save", TPIDR2SaveTy, Attrs);
69   CallInst *Call = Builder.CreateCall(Callee);
70   Call->setCallingConv(
71       CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost_From_X0);
72 
73   // A save to TPIDR2 should be followed by clearing TPIDR2_EL0.
74   Function *WriteIntr =
75       Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_set_tpidr2);
76   Builder.CreateCall(WriteIntr->getFunctionType(), WriteIntr,
77                      Builder.getInt64(0));
78 }
79 
80 /// This function generates code at the beginning and end of a function marked
81 /// with either `aarch64_pstate_za_new` or `aarch64_new_zt0`.
82 /// At the beginning of the function, the following code is generated:
83 ///  - Commit lazy-save if active   [Private-ZA Interface*]
84 ///  - Enable PSTATE.ZA             [Private-ZA Interface]
85 ///  - Zero ZA                      [Has New ZA State]
86 ///  - Zero ZT0                     [Has New ZT0 State]
87 ///
88 /// * A function with new ZT0 state will not change ZA, so committing the
89 /// lazy-save is not strictly necessary. However, the lazy-save mechanism
90 /// may be active on entry to the function, with PSTATE.ZA set to 1. If
91 /// the new ZT0 function calls a function that does not share ZT0, we will
92 /// need to conditionally SMSTOP ZA before the call, setting PSTATE.ZA to 0.
93 /// For this reason, it's easier to always commit the lazy-save at the
94 /// beginning of the function regardless of whether it has ZA state.
95 ///
96 /// At the end of the function, PSTATE.ZA is disabled if the function has a
97 /// Private-ZA Interface. A function is considered to have a Private-ZA
98 /// interface if it does not share ZA or ZT0.
99 ///
100 bool SMEABI::updateNewStateFunctions(Module *M, Function *F,
101                                      IRBuilder<> &Builder, SMEAttrs FnAttrs) {
102   LLVMContext &Context = F->getContext();
103   BasicBlock *OrigBB = &F->getEntryBlock();
104   Builder.SetInsertPoint(&OrigBB->front());
105 
106   // Commit any active lazy-saves if this is a Private-ZA function. If the
107   // value read from TPIDR2_EL0 is not null on entry to the function then
108   // the lazy-saving scheme is active and we should call __arm_tpidr2_save
109   // to commit the lazy save.
110   if (FnAttrs.hasPrivateZAInterface()) {
111     // Create the new blocks for reading TPIDR2_EL0 & enabling ZA state.
112     auto *SaveBB = OrigBB->splitBasicBlock(OrigBB->begin(), "save.za", true);
113     auto *PreludeBB = BasicBlock::Create(Context, "prelude", F, SaveBB);
114 
115     // Read TPIDR2_EL0 in PreludeBB & branch to SaveBB if not 0.
116     Builder.SetInsertPoint(PreludeBB);
117     Function *TPIDR2Intr =
118         Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_get_tpidr2);
119     auto *TPIDR2 = Builder.CreateCall(TPIDR2Intr->getFunctionType(), TPIDR2Intr,
120                                       {}, "tpidr2");
121     auto *Cmp = Builder.CreateCmp(ICmpInst::ICMP_NE, TPIDR2,
122                                   Builder.getInt64(0), "cmp");
123     Builder.CreateCondBr(Cmp, SaveBB, OrigBB);
124 
125     // Create a call __arm_tpidr2_save, which commits the lazy save.
126     Builder.SetInsertPoint(&SaveBB->back());
127     emitTPIDR2Save(M, Builder);
128 
129     // Enable pstate.za at the start of the function.
130     Builder.SetInsertPoint(&OrigBB->front());
131     Function *EnableZAIntr =
132         Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_za_enable);
133     Builder.CreateCall(EnableZAIntr->getFunctionType(), EnableZAIntr);
134   }
135 
136   if (FnAttrs.hasNewZABody()) {
137     Function *ZeroIntr =
138         Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_zero);
139     Builder.CreateCall(ZeroIntr->getFunctionType(), ZeroIntr,
140                        Builder.getInt32(0xff));
141   }
142 
143   if (FnAttrs.isNewZT0()) {
144     Function *ClearZT0Intr =
145         Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_zero_zt);
146     Builder.CreateCall(ClearZT0Intr->getFunctionType(), ClearZT0Intr,
147                        {Builder.getInt32(0)});
148   }
149 
150   if (FnAttrs.hasPrivateZAInterface()) {
151     // Before returning, disable pstate.za
152     for (BasicBlock &BB : *F) {
153       Instruction *T = BB.getTerminator();
154       if (!T || !isa<ReturnInst>(T))
155         continue;
156       Builder.SetInsertPoint(T);
157       Function *DisableZAIntr =
158           Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_za_disable);
159       Builder.CreateCall(DisableZAIntr->getFunctionType(), DisableZAIntr);
160     }
161   }
162 
163   F->addFnAttr("aarch64_expanded_pstate_za");
164   return true;
165 }
166 
167 bool SMEABI::runOnFunction(Function &F) {
168   Module *M = F.getParent();
169   LLVMContext &Context = F.getContext();
170   IRBuilder<> Builder(Context);
171 
172   if (F.isDeclaration() || F.hasFnAttribute("aarch64_expanded_pstate_za"))
173     return false;
174 
175   bool Changed = false;
176   SMEAttrs FnAttrs(F);
177   if (FnAttrs.hasNewZABody() || FnAttrs.isNewZT0())
178     Changed |= updateNewStateFunctions(M, &F, Builder, FnAttrs);
179 
180   return Changed;
181 }
182