[AArch64][SME] Add intrinsics for enabling/disabling ZA.

[sdesmalen]

Summary

This adds the intrinsics:

• void @llvm.aarch64.sme.start.pstateza() -> smstart za
• void @llvm.aarch64.sme.stop.pstateza() -> smstop za

[aemerson]

LGTM.

[paulwalker-arm]

llvm/test/CodeGen/AArch64/sme-toggle-pstateza.ll

11

Sorry for bikeshedding and feel free to ignore but what about llvm.aarch64.sme.za.enable/llvm.aarch64.sme.za.disable?

[aemerson]

llvm/test/CodeGen/AArch64/sme-toggle-pstateza.ll

11

Now that you bring it up, adding PSTATE to the intrinsic names seem a bit too low level, but I don't care enough to object.

llvm/include/llvm/IR/IntrinsicsAArch64.td

def int_aarch64_sme_get_pstatesm
                               [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]>;
+
+  def int_aarch64_sme_za_enable
+      : DefaultAttrsIntrinsic<[], [], [IntrNoMem, IntrHasSideEffects]>;
+  def int_aarch64_sme_za_disable
+      : DefaultAttrsIntrinsic<[], [], [IntrNoMem, IntrHasSideEffects]>;
+
   // Clamp
   //
 
   def int_aarch64_sve_sclamp : AdvSIMD_3VectorArg_Intrinsic;
   def int_aarch64_sve_uclamp : AdvSIMD_3VectorArg_Intrinsic;
 
   //
   // Reversal

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

#undef LCALLNAME5
   if (Subtarget->hasSVE() || Subtarget->hasSME()) {
     for (auto VT :
          {MVT::nxv16i1, MVT::nxv8i1, MVT::nxv4i1, MVT::nxv2i1, MVT::nxv1i1}) {
       setOperationAction(ISD::SPLAT_VECTOR, VT, Custom);
       setOperationAction(ISD::EXTRACT_VECTOR_ELT, VT, Custom);
     }
   }
 
+  if (Subtarget->hasSME())
+    setOperationAction(ISD::INTRINSIC_VOID, 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);
case Intrinsic::aarch64_mops_memset_tag: {
     return DAG.getMergeValues({MS.getValue(0), MS.getValue(2)}, DL);
   }
   case Intrinsic::aarch64_sme_get_pstatesm: {
     SDValue Chain = Op->getOperand(0);
     SMEAttrs Attrs(DAG.getMachineFunction().getFunction());
     SDValue PStateSM = getPStateSM(DAG, Chain, Attrs, DL, Op.getValueType());
     return DAG.getMergeValues({PStateSM, Chain}, DL);
   }
+  case Intrinsic::aarch64_sme_za_enable:
+    return DAG.getNode(
+        AArch64ISD::SMSTART, DL, MVT::Other,
+        Op->getOperand(0), // Chain
+        DAG.getTargetConstant((int32_t)(AArch64SVCR::SVCRZA), DL, MVT::i32),
+        DAG.getConstant(0, DL, MVT::i64), DAG.getConstant(1, DL, MVT::i64));
+  case Intrinsic::aarch64_sme_za_disable:
+    return DAG.getNode(
+        AArch64ISD::SMSTOP, DL, MVT::Other,
+        Op->getOperand(0), // Chain
+        DAG.getTargetConstant((int32_t)(AArch64SVCR::SVCRZA), DL, MVT::i32),
+        DAG.getConstant(0, DL, MVT::i64), DAG.getConstant(1, DL, MVT::i64));
   }
 }
 
 SDValue AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
                                                      SelectionDAG &DAG) const {
   unsigned IntNo = cast<ConstantSDNode>(Op.getOperand(0))->getZExtValue();
   SDLoc dl(Op);
   switch (IntNo) {
SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   case ISD::SET_ROUNDING:
     return LowerSET_ROUNDING(Op, DAG);
   case ISD::MUL:
     return LowerMUL(Op, DAG);
   case ISD::MULHS:
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::MULHS_PRED);
   case ISD::MULHU:
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::MULHU_PRED);
+  case ISD::INTRINSIC_VOID:
   case ISD::INTRINSIC_W_CHAIN:
     return LowerINTRINSIC_W_CHAIN(Op, DAG);
   case ISD::INTRINSIC_WO_CHAIN:
     return LowerINTRINSIC_WO_CHAIN(Op, DAG);
   case ISD::ATOMIC_STORE:
     if (cast<MemSDNode>(Op)->getMemoryVT() == MVT::i128) {
       assert(Subtarget->hasLSE2());
       return LowerStore128(Op, DAG);

llvm/test/CodeGen/AArch64/sme-toggle-pstateza.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64 -mattr=+sme -verify-machineinstrs < %s | FileCheck %s
+
+define void @toggle_pstate_za() {
+; CHECK-LABEL: toggle_pstate_za:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    smstart za
+; CHECK-NEXT:    smstop za
+; CHECK-NEXT:    ret
+  call void @llvm.aarch64.sme.za.enable()
+  call void @llvm.aarch64.sme.za.disable()
+  ret void
+}
+
+declare void @llvm.aarch64.sme.za.enable()
+declare void @llvm.aarch64.sme.za.disable()