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
SMEABI__anon513b6f060111::SMEABI37 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";
INITIALIZE_PASS_BEGIN(SMEABI,DEBUG_TYPE,name,false,false)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.
emitTPIDR2Save(Module * M,IRBuilder<> & Builder)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.
updateNewZAFunctions(Module * M,Function * F,IRBuilder<> & Builder)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
runOnFunction(Function & F)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