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
29using namespace llvm;
30
31#define DEBUG_TYPE "aarch64-sme-abi"
32
33namespace {
34struct 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
42private:
43 bool updateNewStateFunctions(Module *M, Function *F, IRBuilder<> &Builder,
44 SMEAttrs FnAttrs);
45};
46} // end anonymous namespace
47
48char SMEABI::ID = 0;
49static const char *name = "SME ABI Pass";
50INITIALIZE_PASS_BEGIN(SMEABI, DEBUG_TYPE, name, false, false)
51INITIALIZE_PASS_END(SMEABI, DEBUG_TYPE, name, false, false)
52
53FunctionPass *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.
60void 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///
98bool 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
165bool 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

source code of llvm/lib/Target/AArch64/SMEABIPass.cpp