[AArch64][SME] Add codegen pass to handle ZA state in arm_new_za functions.

[david-arm]

Summary

The new pass implements the following:

• Inserts code at the start of an arm_new_za function to commit a lazy-save when the lazy-save mechanism is active.
• Adds a smstart intrinsic at the start of the function.
• Adds a smstop intrinsic at the end of the function.

Patch co-authored by kmclaughlin.

[aemerson]

Using an IR pass to implement ABI doesn't seem right to me. This could be done at codegen time, and having it at the IR level places implicit constraints on later IR passes that aren't enforced.

Even the IR pass was the right choice, I don't see anything here that necessitates a Module pass instead of the usual FunctionPass.

[shruthiashwath]

llvm/lib/Target/AArch64/SMEABIPass.cpp

137

I am not sure if I did look properly but I could not find this attribute in the documents, can you please help me understand this attribute "aarch64_expanded_pstate_za", or redirect me to the document

[david-arm]

In D133896#3802726, @aemerson wrote:

Using an IR pass to implement ABI doesn't seem right to me. This could be done at codegen time, and having it at the IR level places implicit constraints on later IR passes that aren't enforced.

Even the IR pass was the right choice, I don't see anything here that necessitates a Module pass instead of the usual FunctionPass.

Hi @aemerson, I think the reasons why we chose to do this in IR are:

1. Partly due to legacy reasons, since for an earlier version of the ABI it made sense to implement aspects of the ABI in an IR pass. However, I guess there is now a less compelling reason to do so.
2. I think it's probably neater to do this in IR since we can do both the function prologue and epilogue at the same time, otherwise in codegen I imagine we have to do it separately in different lowering functions? In an IR pass we can deal with all aspects of the attributes at the same time in the same function. IIRC it's easier to create/split blocks in IR too?

I can look into how much effort it would take to implement this in codegen, however I'm not sure what you mean by having it at the IR level places implicit constraints on later IR passes that aren't enforced? It's worth saying that we will need an SME ABI IR pass like when it comes to exception handling because that's too complicated to do at the codegen level.

[aemerson]

In D133896#3805074, @david-arm wrote:

In D133896#3802726, @aemerson wrote:

Using an IR pass to implement ABI doesn't seem right to me. This could be done at codegen time, and having it at the IR level places implicit constraints on later IR passes that aren't enforced.

Even the IR pass was the right choice, I don't see anything here that necessitates a Module pass instead of the usual FunctionPass.

Hi @aemerson, I think the reasons why we chose to do this in IR are:

1. Partly due to legacy reasons, since for an earlier version of the ABI it made sense to implement aspects of the ABI in an IR pass. However, I guess there is now a less compelling reason to do so.
2. I think it's probably neater to do this in IR since we can do both the function prologue and epilogue at the same time, otherwise in codegen I imagine we have to do it separately in different lowering functions? In an IR pass we can deal with all aspects of the attributes at the same time in the same function. IIRC it's easier to create/split blocks in IR too?

I can look into how much effort it would take to implement this in codegen, however I'm not sure what you mean by having it at the IR level places implicit constraints on later IR passes that aren't enforced? It's worth saying that we will need an SME ABI IR pass like when it comes to exception handling because that's too complicated to do at the codegen level.

My concern was that by decorating the IR for ABI reasons we also rely on any future IR passes to not accidentally break it by placing instructions in the wrong place. Having said that, I suppose you technically have the same issue if you do it in a MachineFunction pass.

So in that respect my objection to an IR pass is gone, but I think this should be moved as late as possible in the pipeline before codegen runs. The current place it runs seems to be before a bunch of other passes.

[david-arm]

Hi @aemerson and @shruthiashwath, while @sdesmalen is away I'm trying to deal with comments on his behalf. However, Phabricator doesn't let me update his patch, and instead forked off a new one here - D134437. Could you leave any comments on D134437 instead of this one from now on? Thanks!

[david-arm]

• Commandeered patch as per @aemerson's suggestion!

[shruthiashwath]

llvm/docs/AArch64SME.rst

432 ↗

(On Diff #462196)

How do you plan to use this further?
As the pass already handles arm_new_za, is there any specific functionality that needs to handled after the SelectionDAG ?

[david-arm]

llvm/docs/AArch64SME.rst

432 ↗

(On Diff #462196)

This attribute is added purely as a marker to say the function has been processed. We did this mainly as a precaution in case for some reason the SME ABI pass got invoked twice in a row. Perhaps it's being overly cautious. :)

[david-arm]

gentle ping

[bryanpkc]

llvm/lib/Target/AArch64/SMEABIPass.cpp

2

NIT: There should be one fewer space between "SME" and "ABI", and there should be one space after "ABI".

[aemerson]

LGTM.

[c-rhodes]

llvm/test/CodeGen/AArch64/sme-new-za-function.ll

3

Is there a reason for specifying the pass twice?

[sdesmalen]

llvm/test/CodeGen/AArch64/sme-new-za-function.ll

3

We added it to test that running the pass more than once leads to the same IR.

[c-rhodes]

llvm/test/CodeGen/AArch64/sme-new-za-function.ll

3

We added it to test that running the pass more than once leads to the same IR.

Ah ok! Thanks for clarifying, thought it might have been a mistake.

llvm/lib/Target/AArch64/AArch64.h

 FunctionPass *createFalkorHWPFFixPass();
 FunctionPass *createFalkorMarkStridedAccessesPass();
 FunctionPass *createAArch64BranchTargetsPass();
 FunctionPass *createAArch64MIPeepholeOptPass();
 
 FunctionPass *createAArch64CleanupLocalDynamicTLSPass();
 
 FunctionPass *createAArch64CollectLOHPass();
+ModulePass *createSMEABIPass();
 ModulePass *createSVEIntrinsicOptsPass();
 InstructionSelector *
 createAArch64InstructionSelector(const AArch64TargetMachine &,
                                  AArch64Subtarget &, AArch64RegisterBankInfo &);
 FunctionPass *createAArch64O0PreLegalizerCombiner();
 FunctionPass *createAArch64PreLegalizerCombiner();
 FunctionPass *createAArch64PostLegalizerCombiner(bool IsOptNone);
 FunctionPass *createAArch64PostLegalizerLowering();

 void initializeAArch64PostLegalizerLoweringPass(PassRegistry &);
 void initializeAArch64PostSelectOptimizePass(PassRegistry &);
 void initializeAArch64PromoteConstantPass(PassRegistry&);
 void initializeAArch64RedundantCopyEliminationPass(PassRegistry&);
 void initializeAArch64StorePairSuppressPass(PassRegistry&);
 void initializeFalkorHWPFFixPass(PassRegistry&);
 void initializeFalkorMarkStridedAccessesLegacyPass(PassRegistry&);
 void initializeLDTLSCleanupPass(PassRegistry&);
+void initializeSMEABIPass(PassRegistry &);
 void initializeSVEIntrinsicOptsPass(PassRegistry&);
 void initializeAArch64StackTaggingPass(PassRegistry&);
 void initializeAArch64StackTaggingPreRAPass(PassRegistry&);
 } // end namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64TargetMachine.cpp

extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAArch64Target() {
   initializeAArch64PostLegalizerLoweringPass(*PR);
   initializeAArch64PostSelectOptimizePass(*PR);
   initializeAArch64PromoteConstantPass(*PR);
   initializeAArch64RedundantCopyEliminationPass(*PR);
   initializeAArch64StorePairSuppressPass(*PR);
   initializeFalkorHWPFFixPass(*PR);
   initializeFalkorMarkStridedAccessesLegacyPass(*PR);
   initializeLDTLSCleanupPass(*PR);
+  initializeSMEABIPass(*PR);
   initializeSVEIntrinsicOptsPass(*PR);
   initializeAArch64SpeculationHardeningPass(*PR);
   initializeAArch64SLSHardeningPass(*PR);
   initializeAArch64StackTaggingPass(*PR);
   initializeAArch64StackTaggingPreRAPass(*PR);
   initializeAArch64LowerHomogeneousPrologEpilogPass(*PR);
 }
 
void AArch64PassConfig::addIRPasses() {
   // Always expand atomic operations, we don't deal with atomicrmw or cmpxchg
   // ourselves.
   addPass(createAtomicExpandPass());
 
   // Expand any SVE vector library calls that we can't code generate directly.
   if (EnableSVEIntrinsicOpts && TM->getOptLevel() == CodeGenOpt::Aggressive)
     addPass(createSVEIntrinsicOptsPass());
 
+  // Expand any functions marked with SME attributes which require special
+  // changes for the calling convention or that require the lazy-saving
+  // mechanism specified in the SME ABI.
+  addPass(createSMEABIPass());
+
   // Cmpxchg instructions are often used with a subsequent comparison to
   // determine whether it succeeded. We can exploit existing control-flow in
   // ldrex/strex loops to simplify this, but it needs tidying up.
   if (TM->getOptLevel() != CodeGenOpt::None && EnableAtomicTidy)
     addPass(createCFGSimplificationPass(SimplifyCFGOptions()
                                             .forwardSwitchCondToPhi(true)
                                             .convertSwitchRangeToICmp(true)
                                             .convertSwitchToLookupTable(true)

llvm/lib/Target/AArch64/CMakeLists.txt

add_llvm_target(AArch64CodeGen
   AArch64SpeculationHardening.cpp
   AArch64StackTagging.cpp
   AArch64StackTaggingPreRA.cpp
   AArch64StorePairSuppress.cpp
   AArch64Subtarget.cpp
   AArch64TargetMachine.cpp
   AArch64TargetObjectFile.cpp
   AArch64TargetTransformInfo.cpp
+  SMEABIPass.cpp
   SVEIntrinsicOpts.cpp
   AArch64SIMDInstrOpt.cpp
 
   DEPENDS
   intrinsics_gen
 
   LINK_COMPONENTS
   AArch64Desc

llvm/lib/Target/AArch64/SMEABIPass.cpp

+//===--------- SMEABI - SME  ABI-------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass implements parts of the the SME ABI, such as:
+// * Using the lazy-save mechanism before enabling the use of ZA.
+// * Setting up the lazy-save mechanism around invokes.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AArch64.h"
+#include "Utils/AArch64BaseInfo.h"
+#include "Utils/AArch64SMEAttributes.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/IntrinsicsAArch64.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "aarch64-sme-abi"
+
+namespace {
+struct SMEABI : public ModulePass {
+  static char ID; // Pass identification, replacement for typeid
+  SMEABI() : ModulePass(ID) {
+    initializeSMEABIPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnModule(Module &M) override;
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+
+private:
+  bool updateNewZAFunctions(Module *M, Function *F, IRBuilder<> &Builder);
+};
+} // end anonymous namespace
+
+void SMEABI::getAnalysisUsage(AnalysisUsage &AU) const { AU.setPreservesCFG(); }
+
+char SMEABI::ID = 0;
+static const char *name = "SME ABI Pass";
+INITIALIZE_PASS_BEGIN(SMEABI, DEBUG_TYPE, name, false, false)
+INITIALIZE_PASS_END(SMEABI, DEBUG_TYPE, name, false, false)
+
+ModulePass *llvm::createSMEABIPass() { return new SMEABI(); }
+
+//===----------------------------------------------------------------------===//
+// Utility functions
+//===----------------------------------------------------------------------===//
+
+// Utility function to emit a call to __arm_tpidr2_save and clear TPIDR2_EL0.
+void emitTPIDR2Save(Module *M, IRBuilder<> &Builder) {
+  auto *TPIDR2SaveTy =
+      FunctionType::get(Builder.getVoidTy(), {}, /*IsVarArgs=*/false);
+
+  auto Attrs =
+      AttributeList::get(M->getContext(), 0, {"aarch64_pstate_sm_compatible"});
+  FunctionCallee Callee =
+      M->getOrInsertFunction("__arm_tpidr2_save", TPIDR2SaveTy, Attrs);
+  Builder.CreateCall(Callee);
+
+  // A save to TPIDR2 should be followed by clearing TPIDR2_EL0.
+  Function *WriteIntr =
+      Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_set_tpidr2);
+  Builder.CreateCall(WriteIntr->getFunctionType(), WriteIntr,
+                     Builder.getInt64(0));
+}
+
+/// This function generates code to commit a lazy save at the beginning of a
+/// function marked with `aarch64_pstate_za_new`. If the value read from
+/// TPIDR2_EL0 is not null on entry to the function then the lazy-saving scheme
+/// is active and we should call __arm_tpidr2_save to commit the lazy save.
+/// Additionally, PSTATE.ZA should be enabled at the beginning of the function
+/// and disabled before returning.
+bool SMEABI::updateNewZAFunctions(Module *M, Function *F,
+                                  IRBuilder<> &Builder) {
+  LLVMContext &Context = F->getContext();
+  BasicBlock *OrigBB = &F->getEntryBlock();
+
+  // Create the new blocks for reading TPIDR2_EL0 & enabling ZA state.
+  auto *SaveBB = OrigBB->splitBasicBlock(OrigBB->begin(), "save.za", true);
+  auto *PreludeBB = BasicBlock::Create(Context, "prelude", F, SaveBB);
+
+  // Read TPIDR2_EL0 in PreludeBB & branch to SaveBB if not 0.
+  Builder.SetInsertPoint(PreludeBB);
+  Function *TPIDR2Intr =
+      Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_get_tpidr2);
+  auto *TPIDR2 = Builder.CreateCall(TPIDR2Intr->getFunctionType(), TPIDR2Intr,
+                                    {}, "tpidr2");
+  auto *Cmp =
+      Builder.CreateCmp(ICmpInst::ICMP_NE, TPIDR2, Builder.getInt64(0), "cmp");
+  Builder.CreateCondBr(Cmp, SaveBB, OrigBB);
+
+  // Create a call __arm_tpidr2_save, which commits the lazy save.
+  Builder.SetInsertPoint(&SaveBB->back());
+  emitTPIDR2Save(M, Builder);
+
+  // Enable pstate.za at the start of the function.
+  Builder.SetInsertPoint(&OrigBB->front());
+  Function *EnableZAIntr =
+      Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_start_pstateza);
+  Builder.CreateCall(EnableZAIntr->getFunctionType(), EnableZAIntr);
+
+  // Before returning, disable pstate.za
+  for (BasicBlock &BB : F->getBasicBlockList()) {
+    Instruction *T = BB.getTerminator();
+    if (!T || !isa<ReturnInst>(T))
+      continue;
+    Builder.SetInsertPoint(T);
+    Function *DisableZAIntr =
+        Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_stop_pstateza);
+    Builder.CreateCall(DisableZAIntr->getFunctionType(), DisableZAIntr);
+  }
+
+  F->addFnAttr("aarch64_expanded_pstate_za");
+  return true;
+}
+
+bool SMEABI::runOnModule(Module &M) {
+  bool Changed = false;
+
+  LLVMContext &Context = M.getContext();
+  IRBuilder<> Builder(Context);
+
+  for (auto &F : M.getFunctionList()) {
+    if (F.isDeclaration() || F.hasFnAttribute("aarch64_expanded_pstate_za"))
+      continue;
+
+    SMEAttrs FnAttrs(F);
+    if (FnAttrs.hasNewZAInterface())
+      Changed |= updateNewZAFunctions(&M, &F, Builder);
+  }
+
+  return Changed;
+}

llvm/test/CodeGen/AArch64/O0-pipeline.ll

 ; CHECK-NEXT: Create Garbage Collector Module Metadata
 ; CHECK-NEXT: Assumption Cache Tracker
 ; CHECK-NEXT: Profile summary info
 ; CHECK-NEXT: Machine Branch Probability Analysis
 ; CHECK-NEXT:   ModulePass Manager
 ; CHECK-NEXT:     Pre-ISel Intrinsic Lowering
 ; CHECK-NEXT:     FunctionPass Manager
 ; CHECK-NEXT:       Expand Atomic instructions
+; CHECK-NEXT:     SME ABI Pass
+; CHECK-NEXT:     FunctionPass Manager
 ; CHECK-NEXT:       Module Verifier
 ; CHECK-NEXT:       Lower Garbage Collection Instructions
 ; CHECK-NEXT:       Shadow Stack GC Lowering
 ; CHECK-NEXT:       Lower constant intrinsics
 ; CHECK-NEXT:       Remove unreachable blocks from the CFG
 ; CHECK-NEXT:       Expand vector predication intrinsics
 ; CHECK-NEXT:       Scalarize Masked Memory Intrinsics
 ; CHECK-NEXT:       Expand reduction intrinsics

llvm/test/CodeGen/AArch64/O3-pipeline.ll

 ; CHECK-NEXT: Default Regalloc Eviction Advisor
 ; CHECK-NEXT:   ModulePass Manager
 ; CHECK-NEXT:     Pre-ISel Intrinsic Lowering
 ; CHECK-NEXT:     FunctionPass Manager
 ; CHECK-NEXT:       Expand Atomic instructions
 ; CHECK-NEXT:     SVE intrinsics optimizations
 ; CHECK-NEXT:       FunctionPass Manager
 ; CHECK-NEXT:         Dominator Tree Construction
+; CHECK-NEXT:     SME ABI Pass
 ; CHECK-NEXT:     FunctionPass Manager
 ; CHECK-NEXT:       Simplify the CFG
 ; CHECK-NEXT:       Dominator Tree Construction
 ; CHECK-NEXT:       Natural Loop Information
 ; CHECK-NEXT:       Canonicalize natural loops
 ; CHECK-NEXT:       Lazy Branch Probability Analysis
 ; CHECK-NEXT:       Lazy Block Frequency Analysis
 ; CHECK-NEXT:       Optimization Remark Emitter

llvm/test/CodeGen/AArch64/sme-new-za-function.ll

+; RUN: opt -S -mtriple=aarch64-linux-gnu -aarch64-sme-abi %s | FileCheck %s
+; RUN: opt -S -mtriple=aarch64-linux-gnu -aarch64-sme-abi -aarch64-sme-abi %s | FileCheck %s
+
+declare void @shared_za_callee() "aarch64_pstate_za_shared"
+
+define void @private_za() "aarch64_pstate_za_new" {
+; CHECK-LABEL: @private_za(
+; CHECK-NEXT:  prelude:
+; CHECK-NEXT:    [[TPIDR2:%.*]] = call i64 @llvm.aarch64.sme.get.tpidr2()
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i64 [[TPIDR2]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[SAVE_ZA:%.*]], label [[TMP0:%.*]]
+; CHECK:       save.za:
+; CHECK-NEXT:    call void @__arm_tpidr2_save()
+; CHECK-NEXT:    call void @llvm.aarch64.sme.set.tpidr2(i64 0)
+; CHECK-NEXT:    br label [[TMP0]]
+; CHECK:       0:
+; CHECK-NEXT:    call void @llvm.aarch64.sme.start.pstateza()
+; CHECK-NEXT:    call void @shared_za_callee()
+; CHECK-NEXT:    call void @llvm.aarch64.sme.stop.pstateza()
+; CHECK-NEXT:    ret void
+;
+  call void @shared_za_callee()
+  ret void
+}
+
+define i32 @private_za_multiple_exit(i32 %a, i32 %b, i64 %cond) "aarch64_pstate_za_new" {
+; CHECK-LABEL: @private_za_multiple_exit(
+; CHECK-NEXT:  prelude:
+; CHECK-NEXT:    [[TPIDR2:%.*]] = call i64 @llvm.aarch64.sme.get.tpidr2()
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i64 [[TPIDR2]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[SAVE_ZA:%.*]], label [[ENTRY:%.*]]
+; CHECK:       save.za:
+; CHECK-NEXT:    call void @__arm_tpidr2_save()
+; CHECK-NEXT:    call void @llvm.aarch64.sme.set.tpidr2(i64 0)
+; CHECK-NEXT:    br label [[ENTRY]]
+; CHECK:       entry:
+; CHECK-NEXT:    call void @llvm.aarch64.sme.start.pstateza()
+; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i64 [[COND:%.*]], 1
+; CHECK-NEXT:    br i1 [[TOBOOL]], label [[IF_ELSE:%.*]], label [[IF_END:%.*]]
+; CHECK:       if.else:
+; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    call void @llvm.aarch64.sme.stop.pstateza()
+; CHECK-NEXT:    ret i32 [[ADD]]
+; CHECK:       if.end:
+; CHECK-NEXT:    [[SUB:%.*]] = sub i32 [[A]], [[B]]
+; CHECK-NEXT:    call void @llvm.aarch64.sme.stop.pstateza()
+; CHECK-NEXT:    ret i32 [[SUB]]
+;
+entry:
+  %tobool = icmp eq i64 %cond, 1
+  br i1 %tobool, label %if.else, label %if.end
+
+if.else:
+  %add = add i32 %a, %b
+  ret i32 %add
+
+if.end:
+  %sub = sub i32 %a, %b
+  ret i32 %sub
+}
+
+; CHECK: declare "aarch64_pstate_sm_compatible" void @__arm_tpidr2_save()