|
| 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 | + void getAnalysisUsage(AnalysisUsage &AU) const override; |
| 43 | + |
| 44 | +private: |
| 45 | + bool updateNewZAFunctions(Module *M, Function *F, IRBuilder<> &Builder); |
| 46 | +}; |
| 47 | +} // end anonymous namespace |
| 48 | + |
| 49 | +void SMEABI::getAnalysisUsage(AnalysisUsage &AU) const { AU.setPreservesCFG(); } |
| 50 | + |
| 51 | +char SMEABI::ID = 0; |
| 52 | +static const char *name = "SME ABI Pass"; |
| 53 | +INITIALIZE_PASS_BEGIN(SMEABI, DEBUG_TYPE, name, false, false) |
| 54 | +INITIALIZE_PASS_END(SMEABI, DEBUG_TYPE, name, false, false) |
| 55 | + |
| 56 | +FunctionPass *llvm::createSMEABIPass() { return new SMEABI(); } |
| 57 | + |
| 58 | +//===----------------------------------------------------------------------===// |
| 59 | +// Utility functions |
| 60 | +//===----------------------------------------------------------------------===// |
| 61 | + |
| 62 | +// Utility function to emit a call to __arm_tpidr2_save and clear TPIDR2_EL0. |
| 63 | +void emitTPIDR2Save(Module *M, IRBuilder<> &Builder) { |
| 64 | + auto *TPIDR2SaveTy = |
| 65 | + FunctionType::get(Builder.getVoidTy(), {}, /*IsVarArgs=*/false); |
| 66 | + |
| 67 | + auto Attrs = |
| 68 | + AttributeList::get(M->getContext(), 0, {"aarch64_pstate_sm_compatible"}); |
| 69 | + FunctionCallee Callee = |
| 70 | + M->getOrInsertFunction("__arm_tpidr2_save", TPIDR2SaveTy, Attrs); |
| 71 | + Builder.CreateCall(Callee); |
| 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->getBasicBlockList()) { |
| 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 | +} |
0 commit comments