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 29 using namespace llvm; 30 31 #define DEBUG_TYPE "aarch64-sme-abi" 32 33 namespace { 34 struct 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 42 private: 43 bool updateNewStateFunctions(Module *M, Function *F, IRBuilder<> &Builder, 44 SMEAttrs FnAttrs); 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 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 /// 100 bool 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 167 bool 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