[AArch64][SME] Fix broken intrinsics for ZA STR (vector)

[sdesmalen]

Summary

Avoid generating instructions like this:

str za[w12, 0], [x0, #15, mul vl]

Because the index into ZA and the immediate of the addressing mode
should match up, e.g.

str za[w12, 15], [x0, #15, mul vl]

[david-arm]

LGTM! Thanks for this fix @sdesmalen.

[david-arm]

Hi @sdesmalen, really sorry to do this, but following @rsandifo-arm's comment on D127317 I realised this patch is still incorrect because in @str_with_off_15mulvl it should really look like:

mov w12, #-15
str za[w12, 15], [x0, #15, mul vl]

sine the vector select value passed in to the intrinsic was 0. We therefore need to ensure the vector select value in str is also 0.

[david-arm]

@sdesmalen for now I've decided to disable the AddedComplexity = 2 patterns (see D147433).

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 *EmitZASpillFill(MachineInstr &MI, MachineBasicBlock *BB,
+                                     bool IsSpill) const;
   MachineBasicBlock *EmitZAInstr(unsigned Opc, unsigned BaseReg,
                                  MachineInstr &MI, MachineBasicBlock *BB,
                                  bool HasTile) const;
   MachineBasicBlock *EmitZero(MachineInstr &MI, MachineBasicBlock *BB) const;
 
   MachineBasicBlock *
   EmitInstrWithCustomInserter(MachineInstr &MI,
                               MachineBasicBlock *MBB) 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 {
+MachineBasicBlock *AArch64TargetLowering::EmitZASpillFill(MachineInstr &MI,
+                                                          MachineBasicBlock *BB,
+                                                          bool IsSpill) const {
   const TargetInstrInfo *TII = Subtarget->getInstrInfo();
   MachineInstrBuilder MIB =
-      BuildMI(*BB, MI, MI.getDebugLoc(), TII->get(AArch64::LDR_ZA));
+      BuildMI(*BB, MI, MI.getDebugLoc(),
+              TII->get(IsSpill ? AArch64::STR_ZA : AArch64::LDR_ZA));
 
+  if (IsSpill)
+    MIB.addReg(AArch64::ZA);
+  else
     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;
 }
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);
+    return EmitZASpillFill(MI, BB, /*IsSpill=*/false);
+  case AArch64::STR_ZA_PSEUDO:
+    return EmitZASpillFill(MI, BB, /*IsSpill=*/true);
   case AArch64::ZERO_M_PSEUDO:
     return EmitZero(MI, BB);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // AArch64 Lowering private implementation.
 //===----------------------------------------------------------------------===//

llvm/lib/Target/AArch64/SMEInstrFormats.td

: sme_spill_fill_base<0b0, (outs MatrixOp:$ZAt),
                                sme_elm_idx0_15:$imm4, GPR64sp:$Rn,
                                imm0_15:$offset),
                           opcodestr>;
 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>;
+  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 mayStore = 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)>;
+              (!cast<Instruction>(NAME # _PSEUDO) $idx, $imm4, $base)>;
   }
 }
 
 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>;

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

; CHECK-NEXT:    ret
   call void @llvm.aarch64.sme.str(i32 0, ptr %base)
   ret void;
 }
 
 define void @str_with_off_15mulvl(ptr %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:    str za[w12, 15], [x0, #15, mul vl]
 ; CHECK-NEXT:    ret
   %vscale = call i64 @llvm.vscale.i64()
   %mulvl = mul i64 %vscale, 240
   %base = getelementptr i8, ptr %ptr, i64 %mulvl
   call void @llvm.aarch64.sme.str(i32 0, ptr %base)
   ret void;
 }