[AArch64][SME] Add SME addha/va intrinsics

[david-arm]

Summary

This patch adds new the following SME intrinsics:

@llvm.aarch64.sme.addva
@llvm.aarch64.sme.addha

[sagarkulkarni19]

llvm/lib/Target/AArch64/SMEInstrFormats.td

240–241

Need to enclose these instruction definitions in let Predicates = [HasSMEI64] {..}?

[david-arm]

• Added HasSMEI64 guards around 64-bit pseudos/patterns.

[david-arm]

llvm/lib/Target/AArch64/SMEInstrFormats.td

240–241

Thanks @sagarkulkarni19, that's a great spot! The real instruction definitions are guarded correctly, but the patterns and pseudos are not. I think that the mopa/mops may have a similar problem, so I'm going to look into those too.

[sdesmalen]

LGTM with nit addressed.

llvm/lib/Target/AArch64/SMEInstrFormats.td

255

nit: can you merge the definitions of the pseudos and the patterns together under a single HasSMEI64 condition?

llvm/include/llvm/IR/IntrinsicsAArch64.td

let TargetPrefix = "aarch64" in {
   def int_aarch64_sme_smops_wide  : SME_OuterProduct_Intrinsic;
   def int_aarch64_sme_umopa_wide  : SME_OuterProduct_Intrinsic;
   def int_aarch64_sme_umops_wide  : SME_OuterProduct_Intrinsic;
   def int_aarch64_sme_sumopa_wide : SME_OuterProduct_Intrinsic;
   def int_aarch64_sme_sumops_wide : SME_OuterProduct_Intrinsic;
   def int_aarch64_sme_usmopa_wide : SME_OuterProduct_Intrinsic;
   def int_aarch64_sme_usmops_wide : SME_OuterProduct_Intrinsic;
 
+  class SME_AddVectorToTile_Intrinsic
+      : DefaultAttrsIntrinsic<[],
+          [llvm_i64_ty,
+           LLVMScalarOrSameVectorWidth<0, llvm_i1_ty>,
+           LLVMScalarOrSameVectorWidth<0, llvm_i1_ty>,
+           llvm_anyvector_ty]>;
+
+  def int_aarch64_sme_addha : SME_AddVectorToTile_Intrinsic;
+  def int_aarch64_sme_addva : SME_AddVectorToTile_Intrinsic;
+
   //
   // Counting elements
   //
 
   class AdvSIMD_SME_CNTSB_Intrinsic
     : DefaultAttrsIntrinsic<[llvm_i64_ty], [], [IntrNoMem]>;
 
   def int_aarch64_sme_cntsb : AdvSIMD_SME_CNTSB_Intrinsic;

llvm/lib/Target/AArch64/AArch64ISelLowering.h

MachineBasicBlock *EmitTileLoad(unsigned Opc, unsigned BaseReg,
                                   MachineBasicBlock *BB) const;
   MachineBasicBlock *EmitFill(MachineInstr &MI, MachineBasicBlock *BB) const;
   MachineBasicBlock *EmitMopa(unsigned Opc, unsigned BaseReg, MachineInstr &MI,
                               MachineBasicBlock *BB) const;
   MachineBasicBlock *EmitInsertVectorToTile(unsigned Opc, unsigned BaseReg,
                                             MachineInstr &MI,
                                             MachineBasicBlock *BB) const;
   MachineBasicBlock *EmitZero(MachineInstr &MI, MachineBasicBlock *BB) const;
+  MachineBasicBlock *EmitAddVectorToTile(unsigned Opc, unsigned BaseReg,
+                                         MachineInstr &MI,
+                                         MachineBasicBlock *BB) const;
 
   MachineBasicBlock *
   EmitInstrWithCustomInserter(MachineInstr &MI,
                               MachineBasicBlock *MBB) const override;
 
   bool getTgtMemIntrinsic(IntrinsicInfo &Info, const CallInst &I,
                           MachineFunction &MF,
                           unsigned Intrinsic) const override;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

for (unsigned I = 0; I < 8; I++) {
     if (Mask & (1 << I))
       MIB.addDef(AArch64::ZAD0 + I, RegState::ImplicitDefine);
   }
 
   MI.eraseFromParent(); // The pseudo is gone now.
   return BB;
 }
 
+MachineBasicBlock *
+AArch64TargetLowering::EmitAddVectorToTile(unsigned Opc, unsigned BaseReg,
+                                           MachineInstr &MI,
+                                           MachineBasicBlock *BB) const {
+  const TargetInstrInfo *TII = Subtarget->getInstrInfo();
+  MachineInstrBuilder MIB = BuildMI(*BB, MI, MI.getDebugLoc(), TII->get(Opc));
+
+  MIB.addReg(BaseReg + MI.getOperand(0).getImm(), RegState::Define);
+  MIB.addReg(BaseReg + MI.getOperand(0).getImm());
+  MIB.add(MI.getOperand(1)); // pn
+  MIB.add(MI.getOperand(2)); // pm
+  MIB.add(MI.getOperand(3)); // zn
+
+  MI.eraseFromParent(); // The pseudo is gone now.
+  return BB;
+}
+
 MachineBasicBlock *AArch64TargetLowering::EmitInstrWithCustomInserter(
     MachineInstr &MI, MachineBasicBlock *BB) const {
   switch (MI.getOpcode()) {
   default:
 #ifndef NDEBUG
     MI.dump();
 #endif
     llvm_unreachable("Unexpected instruction for custom inserter!");
#endif
   case AArch64::INSERT_MXIPZ_V_PSEUDO_D:
     return EmitInsertVectorToTile(AArch64::INSERT_MXIPZ_V_D, AArch64::ZAD0, MI,
                                   BB);
   case AArch64::INSERT_MXIPZ_V_PSEUDO_Q:
     return EmitInsertVectorToTile(AArch64::INSERT_MXIPZ_V_Q, AArch64::ZAQ0, MI,
                                   BB);
   case AArch64::ZERO_M_PSEUDO:
     return EmitZero(MI, BB);
+  case AArch64::ADDHA_MPPZ_PSEUDO_S:
+    return EmitAddVectorToTile(AArch64::ADDHA_MPPZ_S, AArch64::ZAS0, MI, BB);
+  case AArch64::ADDVA_MPPZ_PSEUDO_S:
+    return EmitAddVectorToTile(AArch64::ADDVA_MPPZ_S, AArch64::ZAS0, MI, BB);
+  case AArch64::ADDHA_MPPZ_PSEUDO_D:
+    return EmitAddVectorToTile(AArch64::ADDHA_MPPZ_D, AArch64::ZAD0, MI, BB);
+  case AArch64::ADDVA_MPPZ_PSEUDO_D:
+    return EmitAddVectorToTile(AArch64::ADDVA_MPPZ_D, AArch64::ZAD0, MI, BB);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // AArch64 Lowering private implementation.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//

llvm/lib/Target/AArch64/SMEInstrFormats.td

 }
 
 class sme_add_vector_to_tile_u64<bit V, string mnemonic>
     : sme_add_vector_to_tile_inst<0b1, V, TileOp64, ZPR64, mnemonic> {
   bits<3> ZAda;
   let Inst{2-0} = ZAda;
 }
 
+class sme_add_vector_to_tile_pseudo<ZPRRegOp zpr_ty>
+    : Pseudo<(outs),
+             (ins i64imm:$tile, PPR3bAny:$Pn, PPR3bAny:$Pm, zpr_ty:$Zn), []>,
+      Sched<[]> {
+  // Translated to the actual instructions in AArch64ISelLowering.cpp
+  let usesCustomInserter = 1;
+}
+
+def ADDHA_MPPZ_PSEUDO_S : sme_add_vector_to_tile_pseudo<ZPR32>;
+def ADDVA_MPPZ_PSEUDO_S : sme_add_vector_to_tile_pseudo<ZPR32>;
+
+def : Pat<(int_aarch64_sme_addha
+            imm0_3:$tile, (nxv4i1 PPR3bAny:$pn), (nxv4i1 PPR3bAny:$pm),
+            (nxv4i32 ZPR32:$zn)),
+          (ADDHA_MPPZ_PSEUDO_S imm0_3:$tile, $pn, $pm, $zn)>;
+def : Pat<(int_aarch64_sme_addva
+            imm0_3:$tile, (nxv4i1 PPR3bAny:$pn), (nxv4i1 PPR3bAny:$pm),
+            (nxv4i32 ZPR32:$zn)),
+          (ADDVA_MPPZ_PSEUDO_S imm0_3:$tile, $pn, $pm, $zn)>;
+
+let Predicates = [HasSMEI64] in {
+def ADDHA_MPPZ_PSEUDO_D : sme_add_vector_to_tile_pseudo<ZPR64>;
+def ADDVA_MPPZ_PSEUDO_D : sme_add_vector_to_tile_pseudo<ZPR64>;
+
+def : Pat<(int_aarch64_sme_addha
+            imm0_7:$tile, (nxv2i1 PPR3bAny:$pn), (nxv2i1 PPR3bAny:$pm),
+            (nxv2i64 ZPR64:$zn)),
+          (ADDHA_MPPZ_PSEUDO_D imm0_7:$tile, $pn, $pm, $zn)>;
+def : Pat<(int_aarch64_sme_addva
+            imm0_7:$tile, (nxv2i1 PPR3bAny:$pn), (nxv2i1 PPR3bAny:$pm),
+            (nxv2i64 ZPR64:$zn)),
+          (ADDVA_MPPZ_PSEUDO_D imm0_7:$tile, $pn, $pm, $zn)>;
+}
+
 //===----------------------------------------------------------------------===//
 // SME Contiguous Loads
 //===----------------------------------------------------------------------===//
 
 class sme_mem_ld_ss_base<bit Q, bit V, bits<2> msz, dag outs, dag ins,
                          string mnemonic, string argstr>
     : I<outs, ins, mnemonic, argstr, "", []>, Sched<[]> {
   bits<5> Rm;

llvm/test/CodeGen/AArch64/sme-intrinsics-add.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme-i64 -verify-machineinstrs < %s | FileCheck %s
+
+define void @addha_s(<vscale x 4 x i1> %pn, <vscale x 4 x i1> %pm, <vscale x 4 x i32> %zn) {
+; CHECK-LABEL: addha_s:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    addha za0.s, p0/m, p1/m, z0.s
+; CHECK-NEXT:    ret
+  call void @llvm.aarch64.sme.addha.nxv4i32(i64 0, <vscale x 4 x i1> %pn, <vscale x 4 x i1> %pm, <vscale x 4 x i32> %zn)
+  ret void
+}
+
+define void @addva_s(<vscale x 4 x i1> %pn, <vscale x 4 x i1> %pm, <vscale x 4 x i32> %zn) {
+; CHECK-LABEL: addva_s:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    addva za3.s, p0/m, p1/m, z0.s
+; CHECK-NEXT:    ret
+  call void @llvm.aarch64.sme.addva.nxv4i32(i64 3, <vscale x 4 x i1> %pn, <vscale x 4 x i1> %pm, <vscale x 4 x i32> %zn)
+  ret void
+}
+
+define void @addha_d(<vscale x 2 x i1> %pn, <vscale x 2 x i1> %pm, <vscale x 2 x i64> %zn) {
+; CHECK-LABEL: addha_d:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    addha za0.d, p0/m, p1/m, z0.d
+; CHECK-NEXT:    ret
+  call void @llvm.aarch64.sme.addha.nxv2i64(i64 0, <vscale x 2 x i1> %pn, <vscale x 2 x i1> %pm, <vscale x 2 x i64> %zn)
+  ret void
+}
+
+define void @addva_d(<vscale x 2 x i1> %pn, <vscale x 2 x i1> %pm, <vscale x 2 x i64> %zn) {
+; CHECK-LABEL: addva_d:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    addva za7.d, p0/m, p1/m, z0.d
+; CHECK-NEXT:    ret
+  call void @llvm.aarch64.sme.addva.nxv2i64(i64 7, <vscale x 2 x i1> %pn, <vscale x 2 x i1> %pm, <vscale x 2 x i64> %zn)
+  ret void
+}
+
+declare void @llvm.aarch64.sme.addha.nxv4i32(i64, <vscale x 4 x i1>, <vscale x 4 x i1>, <vscale x 4 x i32>)
+declare void @llvm.aarch64.sme.addha.nxv2i64(i64, <vscale x 2 x i1>, <vscale x 2 x i1>, <vscale x 2 x i64>)
+declare void @llvm.aarch64.sme.addva.nxv4i32(i64, <vscale x 4 x i1>, <vscale x 4 x i1>, <vscale x 4 x i32>)
+declare void @llvm.aarch64.sme.addva.nxv2i64(i64, <vscale x 2 x i1>, <vscale x 2 x i1>, <vscale x 2 x i64>)