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

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

Differential Revision: https://reviews.llvm.org/D127957

Author: david-arm

llvm/include/llvm/IR/IntrinsicsAArch64.td

@@ -2663,4 +2663,19 @@ let TargetPrefix = "aarch64" in {
   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

@@ -1076,6 +1076,10 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
       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);
@@ -4309,12 +4313,12 @@ static SDValue lowerConvertToSVBool(SDValue Op, SelectionDAG &DAG) {
 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);
@@ -4336,6 +4340,15 @@ SDValue AArch64TargetLowering::LowerINTRINSIC_W_CHAIN(SDValue Op,
     // 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);
+  }
   }
 }
 

llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td

@@ -138,6 +138,12 @@ 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
 //===----------------------------------------------------------------------===//

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

@@ -0,0 +1,14 @@
+; 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

@@ -0,0 +1,23 @@
+; 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)