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 "AArch64SMEAttributes.h"
17#include "llvm/ADT/StringRef.h"
18#include "llvm/CodeGen/TargetLowering.h"
19#include "llvm/CodeGen/TargetPassConfig.h"
20#include "llvm/CodeGen/TargetSubtargetInfo.h"
21#include "llvm/IR/IRBuilder.h"
22#include "llvm/IR/Instructions.h"
23#include "llvm/IR/IntrinsicsAArch64.h"
24#include "llvm/IR/LLVMContext.h"
25#include "llvm/IR/Module.h"
26#include "llvm/IR/RuntimeLibcalls.h"
27#include "llvm/Target/TargetMachine.h"
28#include "llvm/Transforms/Utils/Cloning.h"
29
30using namespace llvm;
31
32#define DEBUG_TYPE "aarch64-sme-abi"
33
34namespace {
35struct SMEABI : public FunctionPass {
36 static char ID; // Pass identification, replacement for typeid
37 SMEABI() : FunctionPass(ID) {}
38
39 bool runOnFunction(Function &F) override;
40
41 void getAnalysisUsage(AnalysisUsage &AU) const override {
42 AU.addRequired<TargetPassConfig>();
43 }
44
45private:
46 bool updateNewStateFunctions(Module *M, Function *F, IRBuilder<> &Builder,
47 SMEAttrs FnAttrs, const TargetLowering &TLI);
48};
49} // end anonymous namespace
50
51char SMEABI::ID = 0;
52static const char *name = "SME ABI Pass";
53INITIALIZE_PASS(SMEABI, DEBUG_TYPE, name, false, false)
54
55FunctionPass *llvm::createSMEABIPass() { return new SMEABI(); }
56
57//===----------------------------------------------------------------------===//
58// Utility functions
59//===----------------------------------------------------------------------===//
60
61// Utility function to emit a call to __arm_tpidr2_save and clear TPIDR2_EL0.
62void emitTPIDR2Save(Module *M, IRBuilder<> &Builder, const TargetLowering &TLI,
63 bool ZT0IsUndef = false) {
64 auto &Ctx = M->getContext();
65 auto *TPIDR2SaveTy =
66 FunctionType::get(Result: Builder.getVoidTy(), Params: {}, /*IsVarArgs=*/isVarArg: false);
67 auto Attrs =
68 AttributeList().addFnAttribute(C&: Ctx, Kind: "aarch64_pstate_sm_compatible");
69 RTLIB::Libcall LC = RTLIB::SMEABI_TPIDR2_SAVE;
70 FunctionCallee Callee =
71 M->getOrInsertFunction(Name: TLI.getLibcallName(Call: LC), T: TPIDR2SaveTy, AttributeList: Attrs);
72 CallInst *Call = Builder.CreateCall(Callee);
73
74 // If ZT0 is undefined (i.e. we're at the entry of a "new_zt0" function), mark
75 // that on the __arm_tpidr2_save call. This prevents an unnecessary spill of
76 // ZT0 that can occur before ZA is enabled.
77 if (ZT0IsUndef)
78 Call->addFnAttr(Attr: Attribute::get(Context&: Ctx, Kind: "aarch64_zt0_undef"));
79
80 Call->setCallingConv(TLI.getLibcallCallingConv(Call: LC));
81
82 // A save to TPIDR2 should be followed by clearing TPIDR2_EL0.
83 Function *WriteIntr =
84 Intrinsic::getOrInsertDeclaration(M, id: Intrinsic::aarch64_sme_set_tpidr2);
85 Builder.CreateCall(FTy: WriteIntr->getFunctionType(), Callee: WriteIntr,
86 Args: Builder.getInt64(C: 0));
87}
88
89/// This function generates code at the beginning and end of a function marked
90/// with either `aarch64_new_za` or `aarch64_new_zt0`.
91/// At the beginning of the function, the following code is generated:
92/// - Commit lazy-save if active [Private-ZA Interface*]
93/// - Enable PSTATE.ZA [Private-ZA Interface]
94/// - Zero ZA [Has New ZA State]
95/// - Zero ZT0 [Has New ZT0 State]
96///
97/// * A function with new ZT0 state will not change ZA, so committing the
98/// lazy-save is not strictly necessary. However, the lazy-save mechanism
99/// may be active on entry to the function, with PSTATE.ZA set to 1. If
100/// the new ZT0 function calls a function that does not share ZT0, we will
101/// need to conditionally SMSTOP ZA before the call, setting PSTATE.ZA to 0.
102/// For this reason, it's easier to always commit the lazy-save at the
103/// beginning of the function regardless of whether it has ZA state.
104///
105/// At the end of the function, PSTATE.ZA is disabled if the function has a
106/// Private-ZA Interface. A function is considered to have a Private-ZA
107/// interface if it does not share ZA or ZT0.
108///
109bool SMEABI::updateNewStateFunctions(Module *M, Function *F,
110 IRBuilder<> &Builder, SMEAttrs FnAttrs,
111 const TargetLowering &TLI) {
112 LLVMContext &Context = F->getContext();
113 BasicBlock *OrigBB = &F->getEntryBlock();
114 Builder.SetInsertPoint(&OrigBB->front());
115
116 // Commit any active lazy-saves if this is a Private-ZA function. If the
117 // value read from TPIDR2_EL0 is not null on entry to the function then
118 // the lazy-saving scheme is active and we should call __arm_tpidr2_save
119 // to commit the lazy save.
120 if (FnAttrs.hasPrivateZAInterface()) {
121 // Create the new blocks for reading TPIDR2_EL0 & enabling ZA state.
122 auto *SaveBB = OrigBB->splitBasicBlockBefore(I: OrigBB->begin(), BBName: "save.za");
123 auto *PreludeBB = BasicBlock::Create(Context, Name: "prelude", Parent: F, InsertBefore: SaveBB);
124
125 // Read TPIDR2_EL0 in PreludeBB & branch to SaveBB if not 0.
126 Builder.SetInsertPoint(PreludeBB);
127 Function *TPIDR2Intr =
128 Intrinsic::getOrInsertDeclaration(M, id: Intrinsic::aarch64_sme_get_tpidr2);
129 auto *TPIDR2 = Builder.CreateCall(FTy: TPIDR2Intr->getFunctionType(), Callee: TPIDR2Intr,
130 Args: {}, Name: "tpidr2");
131 auto *Cmp = Builder.CreateCmp(Pred: ICmpInst::ICMP_NE, LHS: TPIDR2,
132 RHS: Builder.getInt64(C: 0), Name: "cmp");
133 Builder.CreateCondBr(Cond: Cmp, True: SaveBB, False: OrigBB);
134
135 // Create a call __arm_tpidr2_save, which commits the lazy save.
136 Builder.SetInsertPoint(&SaveBB->back());
137 emitTPIDR2Save(M, Builder, TLI, /*ZT0IsUndef=*/FnAttrs.isNewZT0());
138
139 // Enable pstate.za at the start of the function.
140 Builder.SetInsertPoint(&OrigBB->front());
141 Function *EnableZAIntr =
142 Intrinsic::getOrInsertDeclaration(M, id: Intrinsic::aarch64_sme_za_enable);
143 Builder.CreateCall(FTy: EnableZAIntr->getFunctionType(), Callee: EnableZAIntr);
144 }
145
146 if (FnAttrs.isNewZA()) {
147 Function *ZeroIntr =
148 Intrinsic::getOrInsertDeclaration(M, id: Intrinsic::aarch64_sme_zero);
149 Builder.CreateCall(FTy: ZeroIntr->getFunctionType(), Callee: ZeroIntr,
150 Args: Builder.getInt32(C: 0xff));
151 }
152
153 if (FnAttrs.isNewZT0()) {
154 Function *ClearZT0Intr =
155 Intrinsic::getOrInsertDeclaration(M, id: Intrinsic::aarch64_sme_zero_zt);
156 Builder.CreateCall(FTy: ClearZT0Intr->getFunctionType(), Callee: ClearZT0Intr,
157 Args: {Builder.getInt32(C: 0)});
158 }
159
160 if (FnAttrs.hasPrivateZAInterface()) {
161 // Before returning, disable pstate.za
162 for (BasicBlock &BB : *F) {
163 Instruction *T = BB.getTerminator();
164 if (!T || !isa<ReturnInst>(Val: T))
165 continue;
166 Builder.SetInsertPoint(T);
167 Function *DisableZAIntr = Intrinsic::getOrInsertDeclaration(
168 M, id: Intrinsic::aarch64_sme_za_disable);
169 Builder.CreateCall(FTy: DisableZAIntr->getFunctionType(), Callee: DisableZAIntr);
170 }
171 }
172
173 F->addFnAttr(Kind: "aarch64_expanded_pstate_za");
174 return true;
175}
176
177bool SMEABI::runOnFunction(Function &F) {
178 Module *M = F.getParent();
179 LLVMContext &Context = F.getContext();
180 IRBuilder<> Builder(Context);
181
182 if (F.isDeclaration() || F.hasFnAttribute(Kind: "aarch64_expanded_pstate_za"))
183 return false;
184
185 const TargetMachine &TM =
186 getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
187 const TargetLowering &TLI = *TM.getSubtargetImpl(F)->getTargetLowering();
188
189 bool Changed = false;
190 SMEAttrs FnAttrs(F);
191 if (FnAttrs.isNewZA() || FnAttrs.isNewZT0())
192 Changed |= updateNewStateFunctions(M, F: &F, Builder, FnAttrs, TLI);
193
194 return Changed;
195}
196