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/SmallVector.h"
19 #include "llvm/ADT/StringRef.h"
20 #include "llvm/IR/Constants.h"
21 #include "llvm/IR/IRBuilder.h"
22 #include "llvm/IR/Instructions.h"
23 #include "llvm/IR/IntrinsicInst.h"
24 #include "llvm/IR/IntrinsicsAArch64.h"
25 #include "llvm/IR/LLVMContext.h"
26 #include "llvm/InitializePasses.h"
27 #include "llvm/Support/Debug.h"
28 #include "llvm/Transforms/Utils/Cloning.h"
29 
30 using namespace llvm;
31 
32 #define DEBUG_TYPE "aarch64-sme-abi"
33 
34 namespace {
35 struct SMEABI : public FunctionPass {
36   static char ID; // Pass identification, replacement for typeid
37   SMEABI() : FunctionPass(ID) {
38     initializeSMEABIPass(*PassRegistry::getPassRegistry());
39   }
40 
41   bool runOnFunction(Function &F) override;
42 
43 private:
44   bool updateNewZAFunctions(Module *M, Function *F, IRBuilder<> &Builder);
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 to commit a lazy save at the beginning of a
81 /// function marked with `aarch64_pstate_za_new`. If the value read from
82 /// TPIDR2_EL0 is not null on entry to the function then the lazy-saving scheme
83 /// is active and we should call __arm_tpidr2_save to commit the lazy save.
84 /// Additionally, PSTATE.ZA should be enabled at the beginning of the function
85 /// and disabled before returning.
86 bool SMEABI::updateNewZAFunctions(Module *M, Function *F,
87                                   IRBuilder<> &Builder) {
88   LLVMContext &Context = F->getContext();
89   BasicBlock *OrigBB = &F->getEntryBlock();
90 
91   // Create the new blocks for reading TPIDR2_EL0 & enabling ZA state.
92   auto *SaveBB = OrigBB->splitBasicBlock(OrigBB->begin(), "save.za", true);
93   auto *PreludeBB = BasicBlock::Create(Context, "prelude", F, SaveBB);
94 
95   // Read TPIDR2_EL0 in PreludeBB & branch to SaveBB if not 0.
96   Builder.SetInsertPoint(PreludeBB);
97   Function *TPIDR2Intr =
98       Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_get_tpidr2);
99   auto *TPIDR2 = Builder.CreateCall(TPIDR2Intr->getFunctionType(), TPIDR2Intr,
100                                     {}, "tpidr2");
101   auto *Cmp =
102       Builder.CreateCmp(ICmpInst::ICMP_NE, TPIDR2, Builder.getInt64(0), "cmp");
103   Builder.CreateCondBr(Cmp, SaveBB, OrigBB);
104 
105   // Create a call __arm_tpidr2_save, which commits the lazy save.
106   Builder.SetInsertPoint(&SaveBB->back());
107   emitTPIDR2Save(M, Builder);
108 
109   // Enable pstate.za at the start of the function.
110   Builder.SetInsertPoint(&OrigBB->front());
111   Function *EnableZAIntr =
112       Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_za_enable);
113   Builder.CreateCall(EnableZAIntr->getFunctionType(), EnableZAIntr);
114 
115   // Before returning, disable pstate.za
116   for (BasicBlock &BB : *F) {
117     Instruction *T = BB.getTerminator();
118     if (!T || !isa<ReturnInst>(T))
119       continue;
120     Builder.SetInsertPoint(T);
121     Function *DisableZAIntr =
122         Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_za_disable);
123     Builder.CreateCall(DisableZAIntr->getFunctionType(), DisableZAIntr);
124   }
125 
126   F->addFnAttr("aarch64_expanded_pstate_za");
127   return true;
128 }
129 
130 bool SMEABI::runOnFunction(Function &F) {
131   Module *M = F.getParent();
132   LLVMContext &Context = F.getContext();
133   IRBuilder<> Builder(Context);
134 
135   if (F.isDeclaration() || F.hasFnAttribute("aarch64_expanded_pstate_za"))
136     return false;
137 
138   bool Changed = false;
139   SMEAttrs FnAttrs(F);
140   if (FnAttrs.hasNewZAInterface())
141     Changed |= updateNewZAFunctions(M, &F, Builder);
142 
143   return Changed;
144 }
145