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(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///
100bool 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
167bool 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