[AArch64][SME] Fix lowering of llvm.aarch64.get.pstatesm()

[sdesmalen]

Summary

A thread may not have access to SME or TPIDR2_EL0, so in order to
safely query PSTATE.SM in a streaming-compatible function, the
code should call __arm_sme_state(), as described in the ABI:

https://github.com/ARM-software/abi-aa/pull/123/commits/c2bb09c4d4ee60a5787baf1ccc7e92e67e4240b7

This means that the value of pstate.sm is:

• 0 if the function is non-streaming.
• 1 if the function has arm_streaming or arm_locally_streaming.
• evaluated at runtime by a call to __arm_sme_state() otherwise.

[david-arm]

llvm/include/llvm/IR/CallingConv.h

255

Is this right? The ABI document in github suggests:

X2-X15, X19-X29 and SP are call-preserved.

for __arm_sme_state, which is different to the comment, specifically the X0-X13 bit.

llvm/test/CodeGen/AArch64/sme-get-pstatesm.ll

1

I assume you removed this line because you chopped out some unwanted CHECK lines after running the script?

[aemerson]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4486

Reminder to use the suggested merged method from the other patch here.

[sdesmalen]

Split out calling convention for SME support routines in two and added LLVM IR
parsing such that they can be used in a generic way by adding them to function
calls (rather than relying on an intermediary intrinsic to be expanded).

[sdesmalen]

llvm/include/llvm/IR/CallingConv.h

255

Good spot, I've fixed this to add two calling conventions instead.

llvm/test/CodeGen/AArch64/sme-get-pstatesm.ll

1

I removed it, because otherwise it generates loads of CHECK lines for the RUN line with the CHECK-CSRMASK check-prefix.

[sdesmalen]

Fixed sme-support-routines-calling-convention.ll

[sdesmalen]

Fixed up test (again)

[aemerson]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4541

Nit: this could be just SMEAttrs Attrs(DAG.getMachineFunction().getFunction());?

[sdesmalen]

Thanks @aemerson!

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4541

Good spot, I thought I fixed all of those up already.

llvm/include/llvm/IR/CallingConv.h

enum {
     WASM_EmscriptenInvoke = 99,
 
     /// Calling convention used for AMD graphics targets.
     AMDGPU_Gfx = 100,
 
     /// M68k_INTR - Calling convention used for M68k interrupt routines.
     M68k_INTR = 101,
 
+    /// Preserve X0-X13, X19-X29, SP, Z0-Z31, P0-P15.
+    AArch64_SME_ABI_Support_Routines_PreserveMost = 102,
+
     /// The highest possible calling convention ID. Must be some 2^k - 1.
     MaxID = 1023
   };
 
 } // end namespace CallingConv
 
 } // end namespace llvm
 
 #endif // LLVM_IR_CALLINGCONV_H

llvm/lib/Target/AArch64/AArch64CallingConvention.td

 
 // Functions taking SVE arguments or returning an SVE type
 // must (additionally) preserve full Z8-Z23 and predicate registers P4-P15
 def CSR_AArch64_SVE_AAPCS : CalleeSavedRegs<(add (sequence "Z%u", 8, 23),
                                                  (sequence "P%u", 4, 15),
                                                  X19, X20, X21, X22, X23, X24,
                                                  X25, X26, X27, X28, LR, FP)>;
 
+// SME ABI support routines such as __arm_sme_state and __arm_tpidr2_save/restore
+// preserve most registers.
+def CSR_AArch64_SME_ABI_Support_Routines_PreserveMost
+                          : CalleeSavedRegs<(add (sequence "Z%u", 0, 31),
+                                                 (sequence "P%u", 0, 15),
+                                                 (sequence "X%u", 0, 13),
+                                                 (sequence "X%u",19, 28),
+                                                 LR, FP)>;
+
 def CSR_AArch64_AAPCS_SwiftTail
     : CalleeSavedRegs<(sub CSR_AArch64_AAPCS, X20, X22)>;
 
 // Constructors and destructors return 'this' in the iOS 64-bit C++ ABI; since
 // 'this' and the pointer return value are both passed in X0 in these cases,
 // this can be partially modelled by treating X0 as a callee-saved register;
 // only the resulting RegMask is used; the SaveList is ignored
 //

llvm/lib/Target/AArch64/AArch64ISelLowering.h

 // selection DAG.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_AARCH64_AARCH64ISELLOWERING_H
 #define LLVM_LIB_TARGET_AARCH64_AARCH64ISELLOWERING_H
 
 #include "AArch64.h"
+#include "Utils/AArch64SMEAttributes.h"
 #include "llvm/CodeGen/CallingConvLower.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Instruction.h"
 
 namespace llvm {
private:
   //  2. Most unpacked integer types are not legal and thus integer extends
   //     cannot be used to convert between unpacked and packed types.
   // These can make "bitcasting" a multiphase process. REINTERPRET_CAST is used
   // to transition between unpacked and packed types of the same element type,
   // with BITCAST used otherwise.
   // This function does not handle predicate bitcasts.
   SDValue getSVESafeBitCast(EVT VT, SDValue Op, SelectionDAG &DAG) const;
 
+  // Returns the runtime value for PSTATE.SM. When the function is streaming-
+  // compatible, this generates a call to __arm_sme_state.
+  SDValue getPStateSM(SelectionDAG &DAG, SDValue Chain, SMEAttrs Attrs,
+                      SDLoc DL, EVT VT) const;
+
   bool isConstantUnsignedBitfieldExtractLegal(unsigned Opc, LLT Ty1,
                                               LLT Ty2) const override;
 };
 
 namespace AArch64 {
 FastISel *createFastISel(FunctionLoweringInfo &funcInfo,
                          const TargetLibraryInfo *libInfo);
 } // end namespace AArch64
 
 } // end namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (isZeroingInactiveLanes(Op))
     return Reinterpret;
 
   // Zero the newly introduced lanes.
   SDValue Mask = DAG.getConstant(1, DL, InVT);
   Mask = DAG.getNode(AArch64ISD::REINTERPRET_CAST, DL, VT, Mask);
   return DAG.getNode(ISD::AND, DL, VT, Reinterpret, Mask);
 }
 
+SDValue AArch64TargetLowering::getPStateSM(SelectionDAG &DAG, SDValue Chain,
+                                           SMEAttrs Attrs, SDLoc DL,
+                                           EVT VT) const {
+  if (Attrs.hasStreamingInterface() || Attrs.hasStreamingBody())
+    return DAG.getConstant(1, DL, VT);
+
+  if (Attrs.hasNonStreamingInterfaceAndBody())
+    return DAG.getConstant(0, DL, VT);
+
+  assert(Attrs.hasStreamingCompatibleInterface() && "Unexpected interface");
+
+  SDValue Callee = DAG.getExternalSymbol("__arm_sme_state",
+                                         getPointerTy(DAG.getDataLayout()));
+  Type *Int64Ty = Type::getInt64Ty(*DAG.getContext());
+  Type *RetTy = StructType::get(Int64Ty, Int64Ty);
+  TargetLowering::CallLoweringInfo CLI(DAG);
+  ArgListTy Args;
+  CLI.setDebugLoc(DL).setChain(Chain).setLibCallee(
+      CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost, RetTy, Callee,
+      std::move(Args));
+  std::pair<SDValue, SDValue> CallResult = LowerCallTo(CLI);
+  SDValue Mask = DAG.getConstant(/*PSTATE.SM*/ 1, DL, MVT::i64);
+  return DAG.getNode(ISD::AND, DL, MVT::i64, CallResult.first.getOperand(0),
+                     Mask);
+}
+
 SDValue AArch64TargetLowering::LowerINTRINSIC_W_CHAIN(SDValue Op,
                                                       SelectionDAG &DAG) const {
   unsigned IntNo = Op.getConstantOperandVal(1);
   SDLoc DL(Op);
   switch (IntNo) {
   default:
     return SDValue(); // Don't custom lower most intrinsics.
   case Intrinsic::aarch64_mops_memset_tag: {
case Intrinsic::aarch64_mops_memset_tag: {
 
     // MOPS_MEMSET_TAGGING has 3 results (DstWb, SizeWb, Chain) whereas the
     // intrinsic has 2. So hide SizeWb using MERGE_VALUES. Otherwise
     // LowerOperationWrapper will complain that the number of results has
     // changed.
     return DAG.getMergeValues({MS.getValue(0), MS.getValue(2)}, DL);
   }
   case Intrinsic::aarch64_sme_get_pstatesm: {
-    SDValue Chain = Op.getOperand(0);
-    SDValue MRS = DAG.getNode(
-        AArch64ISD::MRS, DL, DAG.getVTList(MVT::i64, MVT::Glue, MVT::Other),
-        Chain, DAG.getConstant(AArch64SysReg::SVCR, DL, MVT::i64));
-    SDValue Mask = DAG.getConstant(/* PSTATE.SM */ 1, DL, MVT::i64);
-    SDValue And = DAG.getNode(ISD::AND, DL, MVT::i64, MRS, Mask);
-    return DAG.getMergeValues({And, Chain}, DL);
+    SDValue Chain = Op->getOperand(0);
+    SMEAttrs Attrs =
+        SMEAttrs::getFromFunction(DAG.getMachineFunction().getFunction());
+    SDValue PStateSM = getPStateSM(DAG, Chain, Attrs, DL, Op.getValueType());
+    return DAG.getMergeValues({PStateSM, Chain}, DL);
   }
   }
 }
 
 SDValue AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
                                                      SelectionDAG &DAG) const {
   unsigned IntNo = cast<ConstantSDNode>(Op.getOperand(0))->getZExtValue();
   SDLoc dl(Op);
case CallingConv::Tail:
     return Subtarget->isTargetILP32() ? CC_AArch64_DarwinPCS_ILP32_VarArg
                                       : CC_AArch64_DarwinPCS_VarArg;
    case CallingConv::Win64:
     return IsVarArg ? CC_AArch64_Win64_VarArg : CC_AArch64_AAPCS;
    case CallingConv::CFGuard_Check:
      return CC_AArch64_Win64_CFGuard_Check;
    case CallingConv::AArch64_VectorCall:
    case CallingConv::AArch64_SVE_VectorCall:
+   case CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost:
      return CC_AArch64_AAPCS;
   }
 }
 
 CCAssignFn *
 AArch64TargetLowering::CCAssignFnForReturn(CallingConv::ID CC) const {
   return CC == CallingConv::WebKit_JS ? RetCC_AArch64_WebKit_JS
                                       : RetCC_AArch64_AAPCS;

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

AArch64RegisterInfo::getCalleeSavedRegs(const MachineFunction *MF) const {
   if (MF->getFunction().getCallingConv() == CallingConv::CFGuard_Check)
     return CSR_Win_AArch64_CFGuard_Check_SaveList;
   if (MF->getSubtarget<AArch64Subtarget>().isTargetWindows())
     return CSR_Win_AArch64_AAPCS_SaveList;
   if (MF->getFunction().getCallingConv() == CallingConv::AArch64_VectorCall)
     return CSR_AArch64_AAVPCS_SaveList;
   if (MF->getFunction().getCallingConv() == CallingConv::AArch64_SVE_VectorCall)
     return CSR_AArch64_SVE_AAPCS_SaveList;
+  if (MF->getFunction().getCallingConv() ==
+      CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost)
+    report_fatal_error(
+        "Calling convention AArch64_SME_ABI_Support_Routines_PreserveMost is "
+        "only supported to improve calls to SME ACLE save/restore/state "
+        "functions, and is not intended to be used beyond that scope.");
   if (MF->getSubtarget<AArch64Subtarget>().getTargetLowering()
           ->supportSwiftError() &&
       MF->getFunction().getAttributes().hasAttrSomewhere(
           Attribute::SwiftError))
     return CSR_AArch64_AAPCS_SwiftError_SaveList;
   if (MF->getFunction().getCallingConv() == CallingConv::SwiftTail)
     return CSR_AArch64_AAPCS_SwiftTail_SaveList;
   if (MF->getFunction().getCallingConv() == CallingConv::PreserveMost)
AArch64RegisterInfo::getDarwinCalleeSavedRegs(const MachineFunction *MF) const {
   if (MF->getFunction().getCallingConv() == CallingConv::CFGuard_Check)
     report_fatal_error(
         "Calling convention CFGuard_Check is unsupported on Darwin.");
   if (MF->getFunction().getCallingConv() == CallingConv::AArch64_VectorCall)
     return CSR_Darwin_AArch64_AAVPCS_SaveList;
   if (MF->getFunction().getCallingConv() == CallingConv::AArch64_SVE_VectorCall)
     report_fatal_error(
         "Calling convention SVE_VectorCall is unsupported on Darwin.");
+  if (MF->getFunction().getCallingConv() ==
+      CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost)
+    report_fatal_error(
+        "Calling convention AArch64_SME_ABI_Support_Routines_PreserveMost is "
+        "only supported to improve calls to SME ACLE save/restore/state "
+        "functions, and is not intended to be used beyond that scope.");
   if (MF->getFunction().getCallingConv() == CallingConv::CXX_FAST_TLS)
     return MF->getInfo<AArch64FunctionInfo>()->isSplitCSR()
                ? CSR_Darwin_AArch64_CXX_TLS_PE_SaveList
                : CSR_Darwin_AArch64_CXX_TLS_SaveList;
   if (MF->getSubtarget<AArch64Subtarget>().getTargetLowering()
           ->supportSwiftError() &&
       MF->getFunction().getAttributes().hasAttrSomewhere(
           Attribute::SwiftError))
AArch64RegisterInfo::getDarwinCallPreservedMask(const MachineFunction &MF,
 
   if (CC == CallingConv::CXX_FAST_TLS)
     return CSR_Darwin_AArch64_CXX_TLS_RegMask;
   if (CC == CallingConv::AArch64_VectorCall)
     return CSR_Darwin_AArch64_AAVPCS_RegMask;
   if (CC == CallingConv::AArch64_SVE_VectorCall)
     report_fatal_error(
         "Calling convention SVE_VectorCall is unsupported on Darwin.");
+  if (CC == CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost)
+    report_fatal_error(
+        "Calling convention AArch64_SME_ABI_Support_Routines_PreserveMost is "
+        "unsupported on Darwin.");
   if (CC == CallingConv::CFGuard_Check)
     report_fatal_error(
         "Calling convention CFGuard_Check is unsupported on Darwin.");
   if (MF.getSubtarget<AArch64Subtarget>()
           .getTargetLowering()
           ->supportSwiftError() &&
       MF.getFunction().getAttributes().hasAttrSomewhere(Attribute::SwiftError))
     return CSR_Darwin_AArch64_AAPCS_SwiftError_RegMask;
if (MF.getSubtarget<AArch64Subtarget>().isTargetDarwin()) {
     return getDarwinCallPreservedMask(MF, CC);
   }
 
   if (CC == CallingConv::AArch64_VectorCall)
     return SCS ? CSR_AArch64_AAVPCS_SCS_RegMask : CSR_AArch64_AAVPCS_RegMask;
   if (CC == CallingConv::AArch64_SVE_VectorCall)
     return SCS ? CSR_AArch64_SVE_AAPCS_SCS_RegMask
                : CSR_AArch64_SVE_AAPCS_RegMask;
+  if (CC == CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost)
+    return CSR_AArch64_SME_ABI_Support_Routines_PreserveMost_RegMask;
   if (CC == CallingConv::CFGuard_Check)
     return CSR_Win_AArch64_CFGuard_Check_RegMask;
   if (MF.getSubtarget<AArch64Subtarget>().getTargetLowering()
           ->supportSwiftError() &&
       MF.getFunction().getAttributes().hasAttrSomewhere(Attribute::SwiftError))
     return SCS ? CSR_AArch64_AAPCS_SwiftError_SCS_RegMask
                : CSR_AArch64_AAPCS_SwiftError_RegMask;
   if (CC == CallingConv::SwiftTail) {
bool AArch64RegisterInfo::isArgumentRegister(const MachineFunction &MF,
   case CallingConv::Win64:
     if (IsVarArg)
       HasReg(CC_AArch64_Win64_VarArg_ArgRegs, Reg);
     return HasReg(CC_AArch64_AAPCS_ArgRegs, Reg);
   case CallingConv::CFGuard_Check:
     return HasReg(CC_AArch64_Win64_CFGuard_Check_ArgRegs, Reg);
   case CallingConv::AArch64_VectorCall:
   case CallingConv::AArch64_SVE_VectorCall:
+  case CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost:
     return HasReg(CC_AArch64_AAPCS_ArgRegs, Reg);
   }
 }
 
 Register
 AArch64RegisterInfo::getFrameRegister(const MachineFunction &MF) const {
   const AArch64FrameLowering *TFI = getFrameLowering(MF);
   return TFI->hasFP(MF) ? AArch64::FP : AArch64::SP;

llvm/test/CodeGen/AArch64/sme-get-pstatesm.ll

-; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme -verify-machineinstrs < %s | FileCheck %s
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme -verify-machineinstrs -stop-after=finalize-isel < %s | FileCheck %s --check-prefix=CHECK-CSRMASK
 
-define i64 @is_streaming() {
-; CHECK-LABEL: is_streaming:
+define i64 @get_pstatesm_normal() nounwind {
+; CHECK-LABEL: get_pstatesm_normal:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mrs x8, SVCR
-; CHECK-NEXT:    and x0, x8, #0x1
+; CHECK-NEXT:    mov x0, xzr
 ; CHECK-NEXT:    ret
   %pstate = call i64 @llvm.aarch64.sme.get.pstatesm()
   ret i64 %pstate
 }
 
+define i64 @get_pstatesm_streaming() nounwind "aarch64_pstate_sm_enabled" {
+; CHECK-LABEL: get_pstatesm_streaming:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w0, #1
+; CHECK-NEXT:    ret
+  %pstate = call i64 @llvm.aarch64.sme.get.pstatesm()
+  ret i64 %pstate
+}
+
+define i64 @get_pstatesm_locally_streaming() nounwind "aarch64_pstate_sm_body" {
+; CHECK-LABEL: get_pstatesm_locally_streaming:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w0, #1
+; CHECK-NEXT:    ret
+  %pstate = call i64 @llvm.aarch64.sme.get.pstatesm()
+  ret i64 %pstate
+}
+
+define i64 @get_pstatesm_streaming_compatible() nounwind "aarch64_pstate_sm_compatible" {
+; CHECK-LABEL: get_pstatesm_streaming_compatible:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    bl __arm_sme_state
+; CHECK-NEXT:    and x0, x0, #0x1
+; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+;
+; CHECK-CSRMASK-LABEL: name: get_pstatesm_streaming_compatible
+; CHECK-CSRMASK: BL &__arm_sme_state, csr_aarch64_sme_abi_support_routines_preservemost
+  %pstate = call i64 @llvm.aarch64.sme.get.pstatesm()
+  ret i64 %pstate
+}
+
 declare i64 @llvm.aarch64.sme.get.pstatesm()