196 lines · cpp
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-exception6//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 typeid37 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 namespace50 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 functions59//===----------------------------------------------------------------------===//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(Builder.getVoidTy(), {}, /*IsVarArgs=*/false);67 auto Attrs =68 AttributeList().addFnAttribute(Ctx, "aarch64_pstate_sm_compatible");69 RTLIB::Libcall LC = RTLIB::SMEABI_TPIDR2_SAVE;70 FunctionCallee Callee =71 M->getOrInsertFunction(TLI.getLibcallName(LC), TPIDR2SaveTy, 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), mark75 // that on the __arm_tpidr2_save call. This prevents an unnecessary spill of76 // ZT0 that can occur before ZA is enabled.77 if (ZT0IsUndef)78 Call->addFnAttr(Attribute::get(Ctx, "aarch64_zt0_undef"));79 80 Call->setCallingConv(TLI.getLibcallCallingConv(LC));81 82 // A save to TPIDR2 should be followed by clearing TPIDR2_EL0.83 Function *WriteIntr =84 Intrinsic::getOrInsertDeclaration(M, Intrinsic::aarch64_sme_set_tpidr2);85 Builder.CreateCall(WriteIntr->getFunctionType(), WriteIntr,86 Builder.getInt64(0));87}88 89/// This function generates code at the beginning and end of a function marked90/// 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 the98/// lazy-save is not strictly necessary. However, the lazy-save mechanism99/// may be active on entry to the function, with PSTATE.ZA set to 1. If100/// the new ZT0 function calls a function that does not share ZT0, we will101/// 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 the103/// 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 a106/// Private-ZA Interface. A function is considered to have a Private-ZA107/// 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 the117 // value read from TPIDR2_EL0 is not null on entry to the function then118 // the lazy-saving scheme is active and we should call __arm_tpidr2_save119 // 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->splitBasicBlock(OrigBB->begin(), "save.za", true);123 auto *PreludeBB = BasicBlock::Create(Context, "prelude", F, 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, Intrinsic::aarch64_sme_get_tpidr2);129 auto *TPIDR2 = Builder.CreateCall(TPIDR2Intr->getFunctionType(), TPIDR2Intr,130 {}, "tpidr2");131 auto *Cmp = Builder.CreateCmp(ICmpInst::ICMP_NE, TPIDR2,132 Builder.getInt64(0), "cmp");133 Builder.CreateCondBr(Cmp, SaveBB, 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, Intrinsic::aarch64_sme_za_enable);143 Builder.CreateCall(EnableZAIntr->getFunctionType(), EnableZAIntr);144 }145 146 if (FnAttrs.isNewZA()) {147 Function *ZeroIntr =148 Intrinsic::getOrInsertDeclaration(M, Intrinsic::aarch64_sme_zero);149 Builder.CreateCall(ZeroIntr->getFunctionType(), ZeroIntr,150 Builder.getInt32(0xff));151 }152 153 if (FnAttrs.isNewZT0()) {154 Function *ClearZT0Intr =155 Intrinsic::getOrInsertDeclaration(M, Intrinsic::aarch64_sme_zero_zt);156 Builder.CreateCall(ClearZT0Intr->getFunctionType(), ClearZT0Intr,157 {Builder.getInt32(0)});158 }159 160 if (FnAttrs.hasPrivateZAInterface()) {161 // Before returning, disable pstate.za162 for (BasicBlock &BB : *F) {163 Instruction *T = BB.getTerminator();164 if (!T || !isa<ReturnInst>(T))165 continue;166 Builder.SetInsertPoint(T);167 Function *DisableZAIntr = Intrinsic::getOrInsertDeclaration(168 M, Intrinsic::aarch64_sme_za_disable);169 Builder.CreateCall(DisableZAIntr->getFunctionType(), DisableZAIntr);170 }171 }172 173 F->addFnAttr("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("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, Builder, FnAttrs, TLI);193 194 return Changed;195}196