[AArch64][SME] Add ldr/str (fill/spill) intrinsics

[david-arm]

Summary

This patch adds implementations for the fill/spill SME ACLE intrinsics:

@llvm.aarch64.sme.ldr
@llvm.aarch64.sme.str

[aemerson]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

2354

Nit: Comments usually start with capitals.

llvm/lib/Target/AArch64/AArch64InstrFormats.td

3210 ↗

(On Diff #435164)

s/offset/Offset

llvm/test/CodeGen/AArch64/SME/sme-intrinsics-loads.ll

265

I'm not following why a byte offset from a pointer would lower into an offset scaled by VL?

[david-arm]

llvm/test/CodeGen/AArch64/SME/sme-intrinsics-loads.ll

265

That's an excellent point @aemerson! Something has clearly gone wrong here - I'll investigate!

[david-arm]

• Fixed bug in isel pattern for imm offset forms of ldr/str

[aemerson]

LGTM.

[c-rhodes]

llvm/lib/Target/AArch64/SMEInstrFormats.td

25

I think s4 implies 4-bit signed?

[bryanpkc]

llvm/test/CodeGen/AArch64/SME/sme-intrinsics-stores.ll

289

@david-arm How does the VL multiplier in the address operand affect the vector offset? Could you explain why za[w12, 0] through za[w12, 15] are selected for VL multipliers 0-15, but for all multipliers larger than 16, only za[w12, 0] is selected (w12 remains zero in all cases)? I understand that ACLE intrinsics won't actually generate slice offsets larger than 15, but the behaviour of the LLVM intrinsics is confusing (as they do not accept an explicit slice offset argument).

[sdesmalen]

llvm/test/CodeGen/AArch64/SME/sme-intrinsics-stores.ll

289

I'm not sure if this answers your question, but it seemed we had a fix for the stores that we hadn't pushed upstream yet. It avoids creating e.g. str za[w12, 0], [x0, #15, mul vl], since the immediates must match up. The fix can be found here: D147136.

[bryanpkc]

llvm/test/CodeGen/AArch64/SME/sme-intrinsics-stores.ll

289

Thanks! I noticed that discrepancy between LDR and STR later yesterday too. But how the LLVM intrinsics are designed to work is still not clear to me. Determining the slice offset based on the addend in the address operand, and resetting the slice offset to 0 if the addend is larger than 256 bytes, seems rather arbitrary. Am I missing something?

[rsandifo-arm]

llvm/test/CodeGen/AArch64/SME/sme-intrinsics-loads.ll

280

I think this is the same as @bryanpkc's comment, but I'm not sure how this can legitimately become:

mov w12, wzr
ldr za[w12, 15], [x0, #15, mul vl]

If the intrinsic uses a single argument for the ZA index, shouldn't the register part of the index be adjusted to match the immediate offset? I.e. shouldn't w12 + 15 == 0, to match the 0 argument to the intrinsic?

[david-arm]

llvm/test/CodeGen/AArch64/SME/sme-intrinsics-loads.ll

280

Hmm, yes I see your point. The instruction encoding is a little odd in that there is a single offset immediate used for both the vector select and for the address offset. I thought we had a fix for one of the issues with ldr/str in D147136, but I think there might still be a problem with the complex pattern. @rsandifo-arm @bryanpkc I'll take a look!

llvm/include/llvm/IR/IntrinsicsAArch64.td

let TargetPrefix = "aarch64" in {
   def int_aarch64_sme_st1w_horiz : SME_Load_Store_S_Intrinsic;
   def int_aarch64_sme_st1d_horiz : SME_Load_Store_D_Intrinsic;
   def int_aarch64_sme_st1q_horiz : SME_Load_Store_Q_Intrinsic;
   def int_aarch64_sme_st1b_vert  : SME_Load_Store_B_Intrinsic;
   def int_aarch64_sme_st1h_vert  : SME_Load_Store_H_Intrinsic;
   def int_aarch64_sme_st1w_vert  : SME_Load_Store_S_Intrinsic;
   def int_aarch64_sme_st1d_vert  : SME_Load_Store_D_Intrinsic;
   def int_aarch64_sme_st1q_vert  : SME_Load_Store_Q_Intrinsic;
+
+  // Spill + fill
+  def int_aarch64_sme_ldr : DefaultAttrsIntrinsic<
+    [], [llvm_i32_ty, llvm_ptr_ty]>;
+  def int_aarch64_sme_str : DefaultAttrsIntrinsic<
+    [], [llvm_i32_ty, llvm_ptr_ty]>;
 }

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

default:
     break;
   }
 
   if (Opcode != ISD::INTRINSIC_VOID)
     return EVT();
 
   const unsigned IntNo =
       cast<ConstantSDNode>(Root->getOperand(1))->getZExtValue();
+  if (IntNo == Intrinsic::aarch64_sme_ldr ||
+      IntNo == Intrinsic::aarch64_sme_str)
+    return MVT::nxv16i8;
+
   if (IntNo != Intrinsic::aarch64_sve_prf)
     return EVT();
 
   // We are using an SVE prefetch intrinsic. Type must be inferred
   // from the width of the predicate.
   return getPackedVectorTypeFromPredicateType(
       Ctx, Root->getOperand(2)->getValueType(0), /*NumVec=*/1);
 }

llvm/lib/Target/AArch64/AArch64ISelLowering.h

MachineBasicBlock *EmitF128CSEL(MachineInstr &MI,
                                   MachineBasicBlock *BB) const;
 
   MachineBasicBlock *EmitLoweredCatchRet(MachineInstr &MI,
                                            MachineBasicBlock *BB) const;
 
   MachineBasicBlock *EmitTileLoad(unsigned Opc, unsigned BaseReg,
                                   MachineInstr &MI,
                                   MachineBasicBlock *BB) const;
+  MachineBasicBlock *EmitFill(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

AArch64TargetLowering::EmitTileLoad(unsigned Opc, unsigned BaseReg,
   MIB.add(MI.getOperand(3)); // pg
   MIB.add(MI.getOperand(4)); // base
   MIB.add(MI.getOperand(5)); // offset
 
   MI.eraseFromParent(); // The pseudo is gone now.
   return BB;
 }
 
+MachineBasicBlock *
+AArch64TargetLowering::EmitFill(MachineInstr &MI, MachineBasicBlock *BB) const {
+  const TargetInstrInfo *TII = Subtarget->getInstrInfo();
+  MachineInstrBuilder MIB =
+      BuildMI(*BB, MI, MI.getDebugLoc(), TII->get(AArch64::LDR_ZA));
+
+  MIB.addReg(AArch64::ZA, RegState::Define);
+  MIB.add(MI.getOperand(0)); // Vector select register
+  MIB.add(MI.getOperand(1)); // Vector select offset
+  MIB.add(MI.getOperand(2)); // Base
+  MIB.add(MI.getOperand(1)); // Offset, same as vector select offset
+
+  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::LD1_MXIPXX_V_PSEUDO_H:
     return EmitTileLoad(AArch64::LD1_MXIPXX_V_H, AArch64::ZAH0, MI, BB);
   case AArch64::LD1_MXIPXX_V_PSEUDO_S:
     return EmitTileLoad(AArch64::LD1_MXIPXX_V_S, AArch64::ZAS0, MI, BB);
   case AArch64::LD1_MXIPXX_V_PSEUDO_D:
     return EmitTileLoad(AArch64::LD1_MXIPXX_V_D, AArch64::ZAD0, MI, BB);
   case AArch64::LD1_MXIPXX_V_PSEUDO_Q:
     return EmitTileLoad(AArch64::LD1_MXIPXX_V_Q, AArch64::ZAQ0, MI, BB);
+  case AArch64::LDR_ZA_PSEUDO:
+    return EmitFill(MI, BB);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // AArch64 Lowering private implementation.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//

llvm/lib/Target/AArch64/SMEInstrFormats.td

 def imm_to_tile128 : ComplexPattern<i64, 1, "ImmToTile<AArch64::ZAQ0>", []>;
 
 def tileslice8   : ComplexPattern<i32 , 2, "SelectSMETileSlice<4>", []>;
 def tileslice16  : ComplexPattern<i32 , 2, "SelectSMETileSlice<3>", []>;
 def tileslice32  : ComplexPattern<i32 , 2, "SelectSMETileSlice<2>", []>;
 def tileslice64  : ComplexPattern<i32 , 2, "SelectSMETileSlice<1>", []>;
 def tileslice128 : ComplexPattern<i32 , 2, "SelectSMETileSlice<0>", []>; // nop
 
+def am_sme_indexed_b4 :ComplexPattern<iPTR, 2, "SelectAddrModeIndexedSVE<0,15>", [], [SDNPWantRoot]>;
+
 //===----------------------------------------------------------------------===//
 // SME Outer Products
 //===----------------------------------------------------------------------===//
 
 class sme_fp_outer_product_inst<bit S, bit sz, MatrixTileOperand za_ty,
                                 ZPRRegOp zpr_ty, string mnemonic>
     : I<(outs za_ty:$ZAda),
       (ins za_ty:$_ZAda, PPR3bAny:$Pn, PPR3bAny:$Pm, zpr_ty:$Zn, zpr_ty:$Zm),
multiclass sme_mem_st_ss<string mnemonic> {
   defm _H : sme_mem_st_v_ss<mnemonic, /*is_col=*/0b0>;
   defm _V : sme_mem_st_v_ss<mnemonic, /*is_col=*/0b1>;
 }
 
 //===----------------------------------------------------------------------===//
 // SME Save and Restore Array
 //===----------------------------------------------------------------------===//
 
-class sme_spill_fill_inst<bit isStore, dag outs, dag ins, string opcodestr>
+class sme_spill_fill_base<bit isStore, dag outs, dag ins, string opcodestr>
     : I<outs, ins, opcodestr, "\t$ZAt[$Rv, $imm4], [$Rn, $offset, mul vl]", "",
         []>,
       Sched<[]> {
   bits<2> Rv;
   bits<5> Rn;
   bits<4> imm4;
   let Inst{31-22} = 0b1110000100;
   let Inst{21}    = isStore;
   let Inst{20-15} = 0b000000;
   let Inst{14-13} = Rv;
   let Inst{12-10} = 0b000;
   let Inst{9-5}   = Rn;
   let Inst{4}     = 0b0;
   let Inst{3-0}   = imm4;
-
-  let mayLoad = !not(isStore);
-  let mayStore = isStore;
-}
-
-multiclass sme_spill_fill<bit isStore, dag outs, dag ins, string opcodestr> {
-  def NAME : sme_spill_fill_inst<isStore, outs, ins, opcodestr>;
-
-  def : InstAlias<opcodestr # "\t$ZAt[$Rv, $imm4], [$Rn]",
-                  (!cast<Instruction>(NAME) MatrixOp:$ZAt,
-                   MatrixIndexGPR32Op12_15:$Rv, sme_elm_idx0_15:$imm4, GPR64sp:$Rn, 0), 1>;
 }
 
-multiclass sme_spill<string opcodestr> {
-  defm NAME : sme_spill_fill<0b1, (outs),
-                             (ins MatrixOp:$ZAt, MatrixIndexGPR32Op12_15:$Rv,
-                                  sme_elm_idx0_15:$imm4, GPR64sp:$Rn,
-                                  imm0_15:$offset),
-                             opcodestr>;
-}
-
-multiclass sme_fill<string opcodestr> {
-  defm NAME : sme_spill_fill<0b0, (outs MatrixOp:$ZAt),
-                             (ins MatrixIndexGPR32Op12_15:$Rv,
-                                  sme_elm_idx0_15:$imm4, GPR64sp:$Rn,
-                                  imm0_15:$offset),
-                             opcodestr>;
+let mayStore = 1 in
+class sme_spill_inst<string opcodestr>
+    : sme_spill_fill_base<0b1, (outs),
+let mayLoad = 1 in
+class sme_fill_inst<string opcodestr>
+    : sme_spill_fill_base<0b0, (outs MatrixOp:$ZAt),
+multiclass sme_spill<string opcodestr> {
+  def NAME : sme_spill_inst<opcodestr>;
+  def : InstAlias<opcodestr # "\t$ZAt[$Rv, $imm4], [$Rn]",
+                  (!cast<Instruction>(NAME) MatrixOp:$ZAt,
+                   MatrixIndexGPR32Op12_15:$Rv, sme_elm_idx0_15:$imm4, GPR64sp:$Rn, 0), 1>;
+  // base
+  def : Pat<(int_aarch64_sme_str MatrixIndexGPR32Op12_15:$idx, GPR64sp:$base),
+            (!cast<Instruction>(NAME) ZA, $idx, 0, $base, 0)>;
+  // scalar + immediate (mul vl)
+  let AddedComplexity = 2 in {
+    def : Pat<(int_aarch64_sme_str MatrixIndexGPR32Op12_15:$idx,
+                                   (am_sme_indexed_b4 GPR64sp:$base, imm0_15:$imm4)),
+              (!cast<Instruction>(NAME) ZA, $idx, 0, $base, $imm4)>;
+  }
+}
+
+multiclass sme_fill<string opcodestr> {
+  def NAME : sme_fill_inst<opcodestr>;
+  def : InstAlias<opcodestr # "\t$ZAt[$Rv, $imm4], [$Rn]",
+                  (!cast<Instruction>(NAME) MatrixOp:$ZAt,
+                   MatrixIndexGPR32Op12_15:$Rv, sme_elm_idx0_15:$imm4, GPR64sp:$Rn, 0), 1>;
+  def NAME # _PSEUDO
+      : Pseudo<(outs),
+               (ins MatrixIndexGPR32Op12_15:$idx, imm0_15:$imm4,
+                    GPR64sp:$base), []>,
+        Sched<[]> {
+    // Translated to actual instruction in AArch64ISelLowering.cpp
+    let usesCustomInserter = 1;
+    let mayLoad = 1;
+  }
+  // base
+  def : Pat<(int_aarch64_sme_ldr MatrixIndexGPR32Op12_15:$idx, GPR64sp:$base),
+            (!cast<Instruction>(NAME # _PSEUDO) $idx, 0, $base)>;
+  // scalar + immediate (mul vl)
+  let AddedComplexity = 2 in {
+    def : Pat<(int_aarch64_sme_ldr MatrixIndexGPR32Op12_15:$idx,
+                                   (am_sme_indexed_b4 GPR64sp:$base, imm0_15:$imm4)),
+              (!cast<Instruction>(NAME # _PSEUDO) $idx, $imm4, $base)>;
+  }
 }
 
 //===----------------------------------------------------------------------===//
 // Move instructions
 //===----------------------------------------------------------------------===//
 
 class sme_vector_to_tile_base<bit Q, bit V, bits<2> sz, dag outs, dag ins,
                               string mnemonic, string argstr>

llvm/test/CodeGen/AArch64/SME/sme-intrinsics-loads.ll

 ; CHECK-NEXT:    ld1q {za15v.q[w12, 0]}, p0/z, [x0, x1, lsl #4]
 ; CHECK-NEXT:    ret
   %base = getelementptr i128, i128* %ptr, i64 %index
   call void @llvm.aarch64.sme.ld1q.horiz(<vscale x 16 x i1> %pg, i128* %base, i64 0, i32 0)
   call void @llvm.aarch64.sme.ld1q.vert(<vscale x 16 x i1> %pg, i128* %base, i64 15, i32 0)
   ret void;
 }
 
+define void @ldr(i8* %ptr) {
+; CHECK-LABEL: ldr:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w12, wzr
+; CHECK-NEXT:    ldr za[w12, 0], [x0]
+; CHECK-NEXT:    ret
+  call void @llvm.aarch64.sme.ldr(i32 0, i8* %ptr)
+  ret void;
+}
+
+define void @ldr_with_off_15(i8* %ptr) {
+; CHECK-LABEL: ldr_with_off_15:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w12, wzr
+; CHECK-NEXT:    add x8, x0, #15
+; CHECK-NEXT:    ldr za[w12, 0], [x8]
+; CHECK-NEXT:    ret
+  %base = getelementptr i8, i8* %ptr, i64 15
+  call void @llvm.aarch64.sme.ldr(i32 0, i8* %base)
+  ret void;
+}
+
+define void @ldr_with_off_15mulvl(i8* %ptr) {
+; CHECK-LABEL: ldr_with_off_15mulvl:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w12, wzr
+; CHECK-NEXT:    ldr za[w12, 15], [x0, #15, mul vl]
+; CHECK-NEXT:    ret
+  %vscale = call i64 @llvm.vscale.i64()
+  %mulvl = mul i64 %vscale, 240
+  %base = getelementptr i8, i8* %ptr, i64 %mulvl
+  call void @llvm.aarch64.sme.ldr(i32 0, i8* %base)
+  ret void;
+}
+
+define void @ldr_with_off_16mulvl(i8* %ptr) {
+; CHECK-LABEL: ldr_with_off_16mulvl:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w12, wzr
+; CHECK-NEXT:    addvl x8, x0, #16
+; CHECK-NEXT:    ldr za[w12, 0], [x8]
+; CHECK-NEXT:    ret
+  %vscale = call i64 @llvm.vscale.i64()
+  %mulvl = mul i64 %vscale, 256
+  %base = getelementptr i8, i8* %ptr, i64 %mulvl
+  call void @llvm.aarch64.sme.ldr(i32 0, i8* %base)
+  ret void;
+}
+
 declare void @llvm.aarch64.sme.ld1b.horiz(<vscale x 16 x i1>, i8*, i64, i32)
 declare void @llvm.aarch64.sme.ld1h.horiz(<vscale x 16 x i1>, i16*, i64, i32)
 declare void @llvm.aarch64.sme.ld1w.horiz(<vscale x 16 x i1>, i32*, i64, i32)
 declare void @llvm.aarch64.sme.ld1d.horiz(<vscale x 16 x i1>, i64*, i64, i32)
 declare void @llvm.aarch64.sme.ld1q.horiz(<vscale x 16 x i1>, i128*, i64, i32)
 declare void @llvm.aarch64.sme.ld1b.vert(<vscale x 16 x i1>, i8*, i64, i32)
 declare void @llvm.aarch64.sme.ld1h.vert(<vscale x 16 x i1>, i16*, i64, i32)
 declare void @llvm.aarch64.sme.ld1w.vert(<vscale x 16 x i1>, i32*, i64, i32)
 declare void @llvm.aarch64.sme.ld1d.vert(<vscale x 16 x i1>, i64*, i64, i32)
 declare void @llvm.aarch64.sme.ld1q.vert(<vscale x 16 x i1>, i128*, i64, i32)
+
+declare void @llvm.aarch64.sme.ldr(i32, i8*)
+declare i64 @llvm.vscale.i64()

llvm/test/CodeGen/AArch64/SME/sme-intrinsics-stores.ll

 ; CHECK-NEXT:    st1q {za15v.q[w12, 0]}, p0, [x0, x1, lsl #4]
 ; CHECK-NEXT:    ret
   %base = getelementptr i128, i128* %ptr, i64 %index
   call void @llvm.aarch64.sme.st1q.horiz(<vscale x 16 x i1> %pg, i128* %base, i64 0, i32 0)
   call void @llvm.aarch64.sme.st1q.vert(<vscale x 16 x i1> %pg, i128* %base, i64 15, i32 0)
   ret void;
 }
 
+define void @str(i8* %ptr) {
+; CHECK-LABEL: str:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w12, wzr
+; CHECK-NEXT:    str za[w12, 0], [x0]
+; CHECK-NEXT:    ret
+  call void @llvm.aarch64.sme.str(i32 0, i8* %ptr)
+  ret void;
+}
+
+define void @str_with_off_15(i8* %ptr) {
+; CHECK-LABEL: str_with_off_15:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w12, wzr
+; CHECK-NEXT:    add x8, x0, #15
+; CHECK-NEXT:    str za[w12, 0], [x8]
+; CHECK-NEXT:    ret
+  %base = getelementptr i8, i8* %ptr, i64 15
+  call void @llvm.aarch64.sme.str(i32 0, i8* %base)
+  ret void;
+}
+
+define void @str_with_off_15mulvl(i8* %ptr) {
+; CHECK-LABEL: str_with_off_15mulvl:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w12, wzr
+; CHECK-NEXT:    str za[w12, 0], [x0, #15, mul vl]
+; CHECK-NEXT:    ret
+  %vscale = call i64 @llvm.vscale.i64()
+  %mulvl = mul i64 %vscale, 240
+  %base = getelementptr i8, i8* %ptr, i64 %mulvl
+  call void @llvm.aarch64.sme.str(i32 0, i8* %base)
+  ret void;
+}
+
+define void @str_with_off_16mulvl(i8* %ptr) {
+; CHECK-LABEL: str_with_off_16mulvl:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w12, wzr
+; CHECK-NEXT:    addvl x8, x0, #16
+; CHECK-NEXT:    str za[w12, 0], [x8]
+; CHECK-NEXT:    ret
+  %vscale = call i64 @llvm.vscale.i64()
+  %mulvl = mul i64 %vscale, 256
+  %base = getelementptr i8, i8* %ptr, i64 %mulvl
+  call void @llvm.aarch64.sme.str(i32 0, i8* %base)
+  ret void;
+}
+
 declare void @llvm.aarch64.sme.st1b.horiz(<vscale x 16 x i1>, i8*, i64, i32)
 declare void @llvm.aarch64.sme.st1h.horiz(<vscale x 16 x i1>, i16*, i64, i32)
 declare void @llvm.aarch64.sme.st1w.horiz(<vscale x 16 x i1>, i32*, i64, i32)
 declare void @llvm.aarch64.sme.st1d.horiz(<vscale x 16 x i1>, i64*, i64, i32)
 declare void @llvm.aarch64.sme.st1q.horiz(<vscale x 16 x i1>, i128*, i64, i32)
 declare void @llvm.aarch64.sme.st1b.vert(<vscale x 16 x i1>, i8*, i64, i32)
 declare void @llvm.aarch64.sme.st1h.vert(<vscale x 16 x i1>, i16*, i64, i32)
 declare void @llvm.aarch64.sme.st1w.vert(<vscale x 16 x i1>, i32*, i64, i32)
 declare void @llvm.aarch64.sme.st1d.vert(<vscale x 16 x i1>, i64*, i64, i32)
 declare void @llvm.aarch64.sme.st1q.vert(<vscale x 16 x i1>, i128*, i64, i32)
+
+declare void @llvm.aarch64.sme.str(i32, i8*)
+declare i64 @llvm.vscale.i64()