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 37 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"; 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 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. 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 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