[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

138

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

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

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

2

Is there a reason for specifying the pass twice?

[sdesmalen]

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

2

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

2

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/docs/AArch64SME.rst

``aarch64_pstate_za_new``
   is used for functions with ``__attribute__((arm_new_za))``
 
 ``aarch64_pstate_za_shared``
   is used for functions with ``__attribute__((arm_shared_za))``
 
 ``aarch64_pstate_za_preserved``
   is used for functions with ``__attribute__((arm_preserves_za))``
 
+``aarch64_expanded_pstate_za``
+  is used for functions with ``__attribute__((arm_new_za))``
+
 Clang must ensure that the above attributes are added both to the
 function's declaration/definition as well as to their call-sites. This is
 important for calls to attributed function pointers, where there is no
 definition or declaration available.
 
 
 2. Handling PSTATE.SM
 =====================

 =====================
 
 In contrast to PSTATE.SM, enabling PSTATE.ZA does not affect the SVE vector
 length and also doesn't clobber FP/AdvSIMD/SVE registers. This means it is safe
 to toggle PSTATE.ZA using intrinsics. This also makes it simpler to setup a
 lazy-save mechanism for calls to private-ZA functions (i.e. functions that may
 either directly or indirectly clobber ZA state).
 
-For this purpose, we'll introduce a new LLVM IR pass that is run just before
-SelectionDAG.
+For the purpose of handling functions marked with ``aarch64_pstate_za_new``,
+we have introduced a new LLVM IR pass (SMEABIPass) that is run just before
+SelectionDAG. Any such functions dealt with by this pass are marked with
+``aarch64_expanded_pstate_za``.
 
 Setting up a lazy-save
 ----------------------
 
 Committing a lazy-save
 ----------------------
 
 Exception handling and ZA

llvm/lib/Target/AArch64/AArch64.h

 FunctionPass *createFalkorHWPFFixPass();
 FunctionPass *createFalkorMarkStridedAccessesPass();
 FunctionPass *createAArch64BranchTargetsPass();
 FunctionPass *createAArch64MIPeepholeOptPass();
 
 FunctionPass *createAArch64CleanupLocalDynamicTLSPass();
 
 FunctionPass *createAArch64CollectLOHPass();
+FunctionPass *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);
 }
 
addPass(createAArch64StackTaggingPass(
       /*IsOptNone=*/TM->getOptLevel() == CodeGenOpt::None));
 
   // Match interleaved memory accesses to ldN/stN intrinsics.
   if (TM->getOptLevel() != CodeGenOpt::None) {
     addPass(createInterleavedLoadCombinePass());
     addPass(createInterleavedAccessPass());
   }
 
+  // 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());
+
   // Add Control Flow Guard checks.
   if (TM->getTargetTriple().isOSWindows())
     addPass(createCFGuardCheckPass());
 
   if (TM->Options.JMCInstrument)
     addPass(createJMCInstrumenterPass());
 }
 

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 FunctionPass {
+  static char ID; // Pass identification, replacement for typeid
+  SMEABI() : FunctionPass(ID) {
+    initializeSMEABIPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnFunction(Function &F) 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)
+
+FunctionPass *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_za_enable);
+  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_za_disable);
+    Builder.CreateCall(DisableZAIntr->getFunctionType(), DisableZAIntr);
+  }
+
+  F->addFnAttr("aarch64_expanded_pstate_za");
+  return true;
+}
+
+bool SMEABI::runOnFunction(Function &F) {
+  Module *M = F.getParent();
+  LLVMContext &Context = F.getContext();
+  IRBuilder<> Builder(Context);
+
+  if (F.isDeclaration() || F.hasFnAttribute("aarch64_expanded_pstate_za"))
+    return false;
+
+  bool Changed = false;
+  SMEAttrs FnAttrs(F);
+  if (FnAttrs.hasNewZAInterface())
+    Changed |= updateNewZAFunctions(M, &F, Builder);
+
+  return Changed;
+}

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

 ; 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
 ; CHECK-NEXT:       AArch64 Stack Tagging
+; CHECK-NEXT:       SME ABI Pass
 ; CHECK-NEXT:       Exception handling preparation
 ; CHECK-NEXT:       Safe Stack instrumentation pass
 ; CHECK-NEXT:       Insert stack protectors
 ; CHECK-NEXT:       Module Verifier
 ; CHECK-NEXT:       Analysis containing CSE Info
 ; CHECK-NEXT:       IRTranslator
 ; CHECK-NEXT:       Analysis for ComputingKnownBits
 ; CHECK-NEXT:       AArch64O0PreLegalizerCombiner

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

 ; CHECK-NEXT:       Basic Alias Analysis (stateless AA impl)
 ; CHECK-NEXT:       Function Alias Analysis Results
 ; CHECK-NEXT:       AArch64 Stack Tagging
 ; CHECK-NEXT:       Function Alias Analysis Results
 ; CHECK-NEXT:       Memory SSA
 ; CHECK-NEXT:       Interleaved Load Combine Pass
 ; CHECK-NEXT:       Dominator Tree Construction
 ; CHECK-NEXT:       Interleaved Access Pass
+; CHECK-NEXT:       SME ABI Pass
 ; CHECK-NEXT:       Natural Loop Information
 ; CHECK-NEXT:       Type Promotion
 ; CHECK-NEXT:       CodeGen Prepare
 ; CHECK-NEXT:       Dominator Tree Construction
 ; CHECK-NEXT:       Exception handling preparation
 ; CHECK-NEXT:     AArch64 Promote Constant
 ; CHECK-NEXT:       FunctionPass Manager
 ; CHECK-NEXT:         Dominator Tree Construction

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.za.enable()
+; CHECK-NEXT:    call void @shared_za_callee()
+; CHECK-NEXT:    call void @llvm.aarch64.sme.za.disable()
+; 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.za.enable()
+; 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.za.disable()
+; CHECK-NEXT:    ret i32 [[ADD]]
+; CHECK:       if.end:
+; CHECK-NEXT:    [[SUB:%.*]] = sub i32 [[A]], [[B]]
+; CHECK-NEXT:    call void @llvm.aarch64.sme.za.disable()
+; 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()