llvm/lib/Target/AArch64/SMEABIPass.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-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 = AttributeList().addFnAttribute(M->getContext(),
  64                                               "aarch64_pstate_sm_compatible");
  65   FunctionCallee Callee =
  66       M->getOrInsertFunction("__arm_tpidr2_save", TPIDR2SaveTy, Attrs);
  67   CallInst *Call = Builder.CreateCall(Callee);
  68   Call->setCallingConv(
  69       CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost_From_X0);
  70
  71   // A save to TPIDR2 should be followed by clearing TPIDR2_EL0.
  72   Function *WriteIntr =
  73       Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_set_tpidr2);
  74   Builder.CreateCall(WriteIntr->getFunctionType(), WriteIntr,
  75                      Builder.getInt64(0));
  76 }
  77
  78 /// This function generates code at the beginning and end of a function marked
  79 /// with either `aarch64_new_za` or `aarch64_new_zt0`.
  80 /// At the beginning of the function, the following code is generated:
  81 ///  - Commit lazy-save if active   [Private-ZA Interface*]
  82 ///  - Enable PSTATE.ZA             [Private-ZA Interface]
  83 ///  - Zero ZA                      [Has New ZA State]
  84 ///  - Zero ZT0                     [Has New ZT0 State]
  85 ///
  86 /// * A function with new ZT0 state will not change ZA, so committing the
  87 /// lazy-save is not strictly necessary. However, the lazy-save mechanism
  88 /// may be active on entry to the function, with PSTATE.ZA set to 1. If
  89 /// the new ZT0 function calls a function that does not share ZT0, we will
  90 /// need to conditionally SMSTOP ZA before the call, setting PSTATE.ZA to 0.
  91 /// For this reason, it's easier to always commit the lazy-save at the
  92 /// beginning of the function regardless of whether it has ZA state.
  93 ///
  94 /// At the end of the function, PSTATE.ZA is disabled if the function has a
  95 /// Private-ZA Interface. A function is considered to have a Private-ZA
  96 /// interface if it does not share ZA or ZT0.
  97 ///
  98 bool SMEABI::updateNewStateFunctions(Module *M, Function *F,
  99                                      IRBuilder<> &Builder, SMEAttrs FnAttrs) {
 100   LLVMContext &Context = F->getContext();
 101   BasicBlock *OrigBB = &F->getEntryBlock();
 102   Builder.SetInsertPoint(&OrigBB->front());
 103
 104   // Commit any active lazy-saves if this is a Private-ZA function. If the
 105   // value read from TPIDR2_EL0 is not null on entry to the function then
 106   // the lazy-saving scheme is active and we should call __arm_tpidr2_save
 107   // to commit the lazy save.
 108   if (FnAttrs.hasPrivateZAInterface()) {
 109     // Create the new blocks for reading TPIDR2_EL0 & enabling ZA state.
 110     auto *SaveBB = OrigBB->splitBasicBlock(OrigBB->begin(), "save.za", true);
 111     auto *PreludeBB = BasicBlock::Create(Context, "prelude", F, SaveBB);
 112
 113     // Read TPIDR2_EL0 in PreludeBB & branch to SaveBB if not 0.
 114     Builder.SetInsertPoint(PreludeBB);
 115     Function *TPIDR2Intr =
 116         Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_get_tpidr2);
 117     auto *TPIDR2 = Builder.CreateCall(TPIDR2Intr->getFunctionType(), TPIDR2Intr,
 118                                       {}, "tpidr2");
 119     auto *Cmp = Builder.CreateCmp(ICmpInst::ICMP_NE, TPIDR2,
 120                                   Builder.getInt64(0), "cmp");
 121     Builder.CreateCondBr(Cmp, SaveBB, OrigBB);
 122
 123     // Create a call __arm_tpidr2_save, which commits the lazy save.
 124     Builder.SetInsertPoint(&SaveBB->back());
 125     emitTPIDR2Save(M, Builder);
 126
 127     // Enable pstate.za at the start of the function.
 128     Builder.SetInsertPoint(&OrigBB->front());
 129     Function *EnableZAIntr =
 130         Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_za_enable);
 131     Builder.CreateCall(EnableZAIntr->getFunctionType(), EnableZAIntr);
 132   }
 133
 134   if (FnAttrs.isNewZA()) {
 135     Function *ZeroIntr =
 136         Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_zero);
 137     Builder.CreateCall(ZeroIntr->getFunctionType(), ZeroIntr,
 138                        Builder.getInt32(0xff));
 139   }
 140
 141   if (FnAttrs.isNewZT0()) {
 142     Function *ClearZT0Intr =
 143         Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_zero_zt);
 144     Builder.CreateCall(ClearZT0Intr->getFunctionType(), ClearZT0Intr,
 145                        {Builder.getInt32(0)});
 146   }
 147
 148   if (FnAttrs.hasPrivateZAInterface()) {
 149     // Before returning, disable pstate.za
 150     for (BasicBlock &BB : *F) {
 151       Instruction *T = BB.getTerminator();
 152       if (!T || !isa<ReturnInst>(T))
 153         continue;
 154       Builder.SetInsertPoint(T);
 155       Function *DisableZAIntr =
 156           Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_za_disable);
 157       Builder.CreateCall(DisableZAIntr->getFunctionType(), DisableZAIntr);
 158     }
 159   }
 160
 161   F->addFnAttr("aarch64_expanded_pstate_za");
 162   return true;
 163 }
 164
 165 bool SMEABI::runOnFunction(Function &F) {
 166   Module *M = F.getParent();
 167   LLVMContext &Context = F.getContext();
 168   IRBuilder<> Builder(Context);
 169
 170   if (F.isDeclaration() || F.hasFnAttribute("aarch64_expanded_pstate_za"))
 171     return false;
 172
 173   bool Changed = false;
 174   SMEAttrs FnAttrs(F);
 175   if (FnAttrs.isNewZA() || FnAttrs.isNewZT0())
 176     Changed |= updateNewStateFunctions(M, &F, Builder, FnAttrs);
 177
 178   return Changed;
 179 }