[AArch64][SME] Add some SME PSTATE setting/query intrinsics

[david-arm]

Summary

This patch adds support for:

• Querying the PSTATE.SM state with @llvm.aarch64.sme.get.pstatesm
• Reading/writing the TPIDR2 register with new @llvm.aarch64.sme.get.tpidr2 and @llvm.aarch64.sme.set.tpidr2 intrinsics.

Tests added here:

CodeGen/AArch64/sme-get-pstatesm.ll
CodeGen/AArch64/sme-read-write-tpidr2.ll

[aemerson]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19346

This is a lowering but it’s being implemented in a combine method, which although it may work, doesn’t seem the right place to do it.

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19346

So you're absolutely right @aemerson, it doesn't feel like the right place to do it. We wanted the intrinsic to return a boolean i1 value, but the problem is that when you do that there is a promotion error:

PromoteIntegerResult #0: t2: i1,ch = llvm.aarch64.sme.get.pstatesm t0, TargetConstant:i64<644>

and it's not easy to fix.

So there are at least two simple solutions:

1. Lower this in a DAG combine, which is a fairly maverick thing to do I know. :)
2. Change the intrinsic to return a i32 so we don't have to worry about promotion.

We chose one because we thought a i1 return made sense, but I'm happy to consider 2 if you prefer it!

[aemerson]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19346

Can't this be done in AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN()? I would have thought if you do a lowering into custom nodes then you bypass any legalization issues with intrinsics. Even if you do still run into that issue, changing it to return i32 seems reasonable to me.

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19346

So there is still a problem here and it looks like Intrinsic::aarch64_rndr is a similar example where someone also had to lower this using performDAGCombine. Even if I make aarch64_sme_get_pstatesm return a i64 value we still can't lower this normally in AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN. The problem is that the AArch64ISD::MRS node needs a chain, which we have during DAG combine. However, in AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN we don't have the chain - I could potentially try to create one, but at least it seems like for Intrinsic::aarch64_rndr someone decided not to do that

[c-rhodes]

llvm/include/llvm/IR/IntrinsicsAArch64.td

2673–2678

can we lower to the generic llvm.read/write_register intrinsics for TPIDR2_EL0?

[david-arm]

• Used LowerINTRINSIC_W_CHAIN to lower get.pstatesm intrinsic instead of performDAGCombine.

[david-arm]

llvm/include/llvm/IR/IntrinsicsAArch64.td

2673–2678

I had a look and this doesn't seem to be something used very often to be honest. It has quite a cumbersome interface because you have to pass in metadata for the first argument specifying the register and I worry it makes it more difficult if we ever want to look for read/writes of TPIDR2_EL0. I feel the approach we have in this patch is at least consistent with all the other AArch64 intrinsics that read/write system registers, e.g. int_aarch64_sve_rdffr, int_aarch64_sve_wrffr, etc.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19346

So it needs a chain because the MRS node needs a chain. It was a bit fiddly trying to get this to work, but eventually discovered you have to mark INTRINSIC_W_CHAIN for MVT::Other as needing custom lowering.

llvm/include/llvm/IR/IntrinsicsAArch64.td

let TargetPrefix = "aarch64" in {
 
   class AdvSIMD_SME_CNTSB_Intrinsic
     : DefaultAttrsIntrinsic<[llvm_i64_ty], [], [IntrNoMem]>;
 
   def int_aarch64_sme_cntsb : AdvSIMD_SME_CNTSB_Intrinsic;
   def int_aarch64_sme_cntsh : AdvSIMD_SME_CNTSB_Intrinsic;
   def int_aarch64_sme_cntsw : AdvSIMD_SME_CNTSB_Intrinsic;
   def int_aarch64_sme_cntsd : AdvSIMD_SME_CNTSB_Intrinsic;
+
+  //
+  // PSTATE Functions
+  //
+
+  def int_aarch64_sme_get_pstatesm
+      : DefaultAttrsIntrinsic<[llvm_i64_ty], [],
+                              [IntrReadMem, IntrInaccessibleMemOnly]>;
+
+  def int_aarch64_sme_get_tpidr2
+      : DefaultAttrsIntrinsic<[llvm_i64_ty], [],
+                              [IntrNoMem, IntrHasSideEffects]>;
+  def int_aarch64_sme_set_tpidr2
+      : DefaultAttrsIntrinsic<[], [llvm_i64_ty],
+                              [IntrNoMem, IntrHasSideEffects]>;
 }

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Subtarget->hasNEON()) {
     // ADDP custom lowering
     for (MVT VT : { MVT::v32i8, MVT::v16i16, MVT::v8i32, MVT::v4i64 })
       setOperationAction(ISD::ADD, VT, Custom);
     // FADDP custom lowering
     for (MVT VT : { MVT::v16f16, MVT::v8f32, MVT::v4f64 })
       setOperationAction(ISD::FADD, VT, Custom);
   }
 
+  if (Subtarget->hasSME()) {
+    setOperationAction(ISD::INTRINSIC_W_CHAIN, MVT::Other, Custom);
+  }
+
   if (Subtarget->hasSVE()) {
     for (auto VT : {MVT::nxv16i8, MVT::nxv8i16, MVT::nxv4i32, MVT::nxv2i64}) {
       setOperationAction(ISD::BITREVERSE, VT, Custom);
       setOperationAction(ISD::BSWAP, VT, Custom);
       setOperationAction(ISD::CTLZ, VT, Custom);
       setOperationAction(ISD::CTPOP, VT, Custom);
       setOperationAction(ISD::CTTZ, VT, Custom);
       setOperationAction(ISD::INSERT_SUBVECTOR, VT, Custom);
static SDValue lowerConvertToSVBool(SDValue Op, SelectionDAG &DAG) {
   SDValue MaskReinterpret =
       DAG.getNode(AArch64ISD::REINTERPRET_CAST, DL, OutVT, Mask);
   return DAG.getNode(ISD::AND, DL, OutVT, Reinterpret, MaskReinterpret);
 }
 
 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: {
     auto Node = cast<MemIntrinsicSDNode>(Op.getNode());
-    SDLoc DL(Op);
     SDValue Chain = Node->getChain();
     SDValue Dst = Op.getOperand(2);
     SDValue Val = Op.getOperand(3);
     Val = DAG.getAnyExtOrTrunc(Val, DL, MVT::i64);
     SDValue Size = Op.getOperand(4);
     auto Alignment = Node->getMemOperand()->getAlign();
     bool IsVol = Node->isVolatile();
     auto DstPtrInfo = Node->getPointerInfo();
 
     const auto &SDI =
         static_cast<const AArch64SelectionDAGInfo &>(DAG.getSelectionDAGInfo());
     SDValue MS =
         SDI.EmitMOPS(AArch64ISD::MOPS_MEMSET_TAGGING, DAG, DL, Chain, Dst, Val,
                      Size, Alignment, IsVol, DstPtrInfo, MachinePointerInfo{});
 
     // 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 AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
                                                      SelectionDAG &DAG) const {
   unsigned IntNo = cast<ConstantSDNode>(Op.getOperand(0))->getZExtValue();
   SDLoc dl(Op);
   switch (IntNo) {
case Intrinsic::aarch64_rndrrs: {
           N->getOperand(0), DAG.getConstant(Register, DL, MVT::i64));
       SDValue B = DAG.getNode(
           AArch64ISD::CSINC, DL, MVT::i32, DAG.getConstant(0, DL, MVT::i32),
           DAG.getConstant(0, DL, MVT::i32),
           DAG.getConstant(AArch64CC::NE, DL, MVT::i32), A.getValue(1));
       return DAG.getMergeValues(
           {A, DAG.getZExtOrTrunc(B, DL, MVT::i1), A.getValue(2)}, DL);
     }
     default:
       break;
     }
     break;
   case ISD::GlobalAddress:
     return performGlobalAddressCombine(N, DAG, Subtarget, getTargetMachine());
   }
   return SDValue();
 }

llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td

 def : InstAlias<"smstart",    (MSRpstatesvcrImm1 0b011, 0b1)>;
 def : InstAlias<"smstart sm", (MSRpstatesvcrImm1 0b001, 0b1)>;
 def : InstAlias<"smstart za", (MSRpstatesvcrImm1 0b010, 0b1)>;
 
 def : InstAlias<"smstop",     (MSRpstatesvcrImm1 0b011, 0b0)>;
 def : InstAlias<"smstop sm",  (MSRpstatesvcrImm1 0b001, 0b0)>;
 def : InstAlias<"smstop za",  (MSRpstatesvcrImm1 0b010, 0b0)>;
 
+// Read and write TPIDR2_EL0
+def : Pat<(int_aarch64_sme_set_tpidr2 i64:$val),
+          (MSR 0xde85, GPR64:$val)>;
+def : Pat<(i64 (int_aarch64_sme_get_tpidr2)),
+          (MRS 0xde85)>;
+
 //===----------------------------------------------------------------------===//
 // SVE2 instructions
 //===----------------------------------------------------------------------===//
 
 def REVD_ZPmZ : sve2_int_perm_revd<"revd">;
 
 defm SCLAMP_ZZZ : sve2_clamp<"sclamp", 0b0>;
 defm UCLAMP_ZZZ : sve2_clamp<"uclamp", 0b1>;
 
 defm PSEL_PPPRI : sve2_int_perm_sel_p<"psel">;
 
 } // End let Predicates = [HasSME]

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
+
+define i64 @is_streaming() {
+; CHECK-LABEL: is_streaming:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mrs x8, SVCR
+; CHECK-NEXT:    and x0, x8, #0x1
+; CHECK-NEXT:    ret
+  %pstate = call i64 @llvm.aarch64.sme.get.pstatesm()
+  ret i64 %pstate
+}
+
+declare i64 @llvm.aarch64.sme.get.pstatesm()

llvm/test/CodeGen/AArch64/sme-read-write-tpidr2.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64 -mattr=+sme < %s | FileCheck %s
+
+define i64 @get_tpidr2_el0() {
+; CHECK-LABEL: get_tpidr2_el0:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mrs x0, TPIDR2_EL0
+; CHECK-NEXT:    ret
+  %res = call i64 @llvm.aarch64.sme.get.tpidr2()
+  ret i64 %res
+}
+
+define void @set_tpidr2_el0(i64 %v) {
+; CHECK-LABEL: set_tpidr2_el0:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    msr TPIDR2_EL0, x0
+; CHECK-NEXT:    ret
+  call void @llvm.aarch64.sme.set.tpidr2(i64 %v)
+  ret void
+}
+
+declare i64 @llvm.aarch64.sme.get.tpidr2()
+declare void @llvm.aarch64.sme.set.tpidr2(i64 %v)