[AArch64][SME] Add ldr and str instructions

[c-rhodes]

Summary

This patch adds support for ldr/str instructions to load/store vector
to/from ZA array.

Depends on D105572.

The reference can be found here:
https://developer.arm.com/documentation/ddi0602/2021-06

[c-rhodes]

Changes:

• Updated sed line in tests to fix Windows precommit.
• Improved diagnostics tests.

[kmclaughlin]

LGTM!

llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td

 let Predicates = [HasSME] in {
 //===----------------------------------------------------------------------===//
 // Loads and stores
 //===----------------------------------------------------------------------===//
 
 defm LD1_MXIPXX : sme_mem_ld_ss<"ld1">;
 defm ST1_MXIPXX : sme_mem_st_ss<"st1">;
 
+//===----------------------------------------------------------------------===//
+// Spill + fill
+//===----------------------------------------------------------------------===//
+
+defm LDR_ZA : sme_fill<"ldr">;
+defm STR_ZA : sme_spill<"str">;
+
 } // End let Predicates = [HasSME]

llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp

AArch64AsmParser::tryParseMatrixRegister(OperandVector &Operands) {
 
   StringRef Name = Tok.getString();
 
   if (Name.equals_insensitive("za")) {
     Parser.Lex(); // eat "za"
     Operands.push_back(AArch64Operand::CreateMatrixRegister(
         AArch64::ZA, /*ElementWidth=*/0, MatrixKind::Array, S, getLoc(),
         getContext()));
+    if (getLexer().is(AsmToken::LBrac)) {
+      // There's no comma after matrix operand, so we can parse the next operand
+      // immediately.
+      if (parseOperand(Operands, false, false))
+        return MatchOperand_NoMatch;
+    }
     return MatchOperand_Success;
   }
 
   // Try to parse matrix register.
   unsigned Reg = matchRegisterNameAlias(Name, RegKind::Matrix);
   if (!Reg)
     return MatchOperand_NoMatch;
 

llvm/lib/Target/AArch64/Disassembler/AArch64Disassembler.cpp

DecodeStatus Result =
         decodeInstruction(Table, MI, Insn, Address, this, STI);
 
     switch (MI.getOpcode()) {
     default:
       break;
     // For Scalable Matrix Extension (SME) instructions that have an implicit
     // operand for the accumulator (ZA) which isn't encoded, manually insert
     // operand.
+    case AArch64::LDR_ZA:
+    case AArch64::STR_ZA: {
+      MI.insert(MI.begin(), MCOperand::createReg(AArch64::ZA));
+      // Spill and fill instructions have a single immediate used for both the
+      // vector select offset and optional memory offset. Replicate the decoded
+      // immediate.
+      const MCOperand &Imm4Op = MI.getOperand(2);
+      assert(Imm4Op.isImm() && "Unexpected operand type!");
+      MI.addOperand(Imm4Op);
+      break;
+    }
     case AArch64::LD1_MXIPXX_H_B:
     case AArch64::LD1_MXIPXX_V_B:
     case AArch64::ST1_MXIPXX_H_B:
     case AArch64::ST1_MXIPXX_V_B:
       MI.insert(MI.begin(), MCOperand::createReg(AArch64::ZAB0));
       break;
     }
 

llvm/lib/Target/AArch64/SMEInstrFormats.td

multiclass sme_mem_st_v_ss<string mnemonic, bit is_col> {
   defm : sme_mem_st_ss_aliases<NAME, is_col>;
 }
 
 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>
+    : 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, imm0_15:$imm4, GPR64sp:$Rn, 0), 1>;
+}
+
+multiclass sme_spill<string opcodestr> {
+  defm NAME : sme_spill_fill<0b1, (outs),
+                             (ins MatrixOp:$ZAt, MatrixIndexGPR32Op12_15:$Rv,
+                                  imm0_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,
+                                  imm0_15:$imm4, GPR64sp:$Rn,
+                                  imm0_15:$offset),
+                             opcodestr>;
+}
+

llvm/test/MC/AArch64/SME/ldr-diagnostics.s

+// RUN: not llvm-mc -triple=aarch64 -show-encoding -mattr=+sme 2>&1 < %s| FileCheck %s
+
+// ------------------------------------------------------------------------- //
+// Invalid matrix operand (expected: za)
+
+ldr za0h.b[w12, #0], [x0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid matrix operand, expected za
+// CHECK-NEXT: ldr za0h.b[w12, #0], [x0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+ldr za3.s[w12, #0], [x0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid matrix operand, expected za
+// CHECK-NEXT: ldr za3.s[w12, #0], [x0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+// ------------------------------------------------------------------------- //
+// Invalid vector select register (expected: w12-w15)
+
+ldr za[w11, #0], [x0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid operand for instruction
+// CHECK-NEXT: ldr za[w11, #0], [x0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+ldr za[w16, #0], [x0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid operand for instruction
+// CHECK-NEXT: ldr za[w16, #0], [x0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+// ------------------------------------------------------------------------- //
+// Invalid vector select offset (expected: 0-15)
+
+ldr za[w12, #16], [x0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: immediate must be an integer in range [0, 15].
+// CHECK-NEXT: ldr za[w12, #16], [x0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+// ------------------------------------------------------------------------- //
+// Invalid memory operands
+
+ldr za[w12, #0], [w0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid operand for instruction
+// CHECK-NEXT: ldr za[w12, #0], [w0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+ldr za[w12, #0], [x0, #16, mul vl]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: immediate must be an integer in range [0, 15].
+// CHECK-NEXT: ldr za[w12, #0], [x0, #16, mul vl]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+ldr za[w12, #0], [x0, #0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid operand for instruction
+// CHECK-NEXT: ldr za[w12, #0], [x0, #0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:

llvm/test/MC/AArch64/SME/ldr.s

+// RUN: llvm-mc -triple=aarch64 -show-encoding -mattr=+sme < %s \
+// RUN:        | FileCheck %s --check-prefixes=CHECK-ENCODING,CHECK-INST
+// RUN: not llvm-mc -triple=aarch64 -show-encoding < %s 2>&1 \
+// RUN:        | FileCheck %s --check-prefix=CHECK-ERROR
+// RUN: llvm-mc -triple=aarch64 -filetype=obj -mattr=+sme < %s \
+// RUN:        | llvm-objdump -d --mattr=+sme - | FileCheck %s --check-prefix=CHECK-INST
+// RUN: llvm-mc -triple=aarch64 -filetype=obj -mattr=+sme < %s \
+// RUN:        | llvm-objdump -d - | FileCheck %s --check-prefix=CHECK-UNKNOWN
+// Disassemble encoding and check the re-encoding (-show-encoding) matches.
+// RUN: llvm-mc -triple=aarch64 -show-encoding -mattr=+sme < %s \
+// RUN:        | sed '/.text/d' | sed 's/.*encoding: //g' \
+// RUN:        | llvm-mc -triple=aarch64 -mattr=+sme -disassemble -show-encoding \
+// RUN:        | FileCheck %s --check-prefixes=CHECK-ENCODING,CHECK-INST
+
+ldr     za[w12, #0], [x0]
+// CHECK-INST: ldr     za[w12, #0], [x0]
+// CHECK-ENCODING: [0x00,0x00,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 00 00 00 e1 <unknown>
+
+ldr     za[w14, #5], [x10, #5, mul vl]
+// CHECK-INST: ldr     za[w14, #5], [x10, #5, mul vl]
+// CHECK-ENCODING: [0x45,0x41,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 45 41 00 e1 <unknown>
+
+ldr     za[w15, #7], [x13, #7, mul vl]
+// CHECK-INST: ldr     za[w15, #7], [x13, #7, mul vl]
+// CHECK-ENCODING: [0xa7,0x61,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: a7 61 00 e1 <unknown>
+
+ldr     za[w15, #15], [sp, #15, mul vl]
+// CHECK-INST: ldr     za[w15, #15], [sp, #15, mul vl]
+// CHECK-ENCODING: [0xef,0x63,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: ef 63 00 e1 <unknown>
+
+ldr     za[w12, #5], [x17, #5, mul vl]
+// CHECK-INST: ldr     za[w12, #5], [x17, #5, mul vl]
+// CHECK-ENCODING: [0x25,0x02,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 25 02 00 e1 <unknown>
+
+ldr     za[w12, #1], [x1, #1, mul vl]
+// CHECK-INST: ldr     za[w12, #1], [x1, #1, mul vl]
+// CHECK-ENCODING: [0x21,0x00,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 21 00 00 e1 <unknown>
+
+ldr     za[w14, #8], [x19, #8, mul vl]
+// CHECK-INST: ldr     za[w14, #8], [x19, #8, mul vl]
+// CHECK-ENCODING: [0x68,0x42,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 68 42 00 e1 <unknown>
+
+ldr     za[w12, #0], [x12]
+// CHECK-INST: ldr     za[w12, #0], [x12]
+// CHECK-ENCODING: [0x80,0x01,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 80 01 00 e1 <unknown>
+
+ldr     za[w14, #1], [x1, #1, mul vl]
+// CHECK-INST: ldr     za[w14, #1], [x1, #1, mul vl]
+// CHECK-ENCODING: [0x21,0x40,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 21 40 00 e1 <unknown>
+
+ldr     za[w12, #13], [x22, #13, mul vl]
+// CHECK-INST: ldr     za[w12, #13], [x22, #13, mul vl]
+// CHECK-ENCODING: [0xcd,0x02,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: cd 02 00 e1 <unknown>
+
+ldr     za[w15, #2], [x9, #2, mul vl]
+// CHECK-INST: ldr     za[w15, #2], [x9, #2, mul vl]
+// CHECK-ENCODING: [0x22,0x61,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 22 61 00 e1 <unknown>
+
+ldr     za[w13, #7], [x12, #7, mul vl]
+// CHECK-INST: ldr     za[w13, #7], [x12, #7, mul vl]
+// CHECK-ENCODING: [0x87,0x21,0x00,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 87 21 00 e1 <unknown>

llvm/test/MC/AArch64/SME/str-diagnostics.s

+// RUN: not llvm-mc -triple=aarch64 -show-encoding -mattr=+sme 2>&1 < %s| FileCheck %s
+
+// ------------------------------------------------------------------------- //
+// Invalid matrix operand (expected: za)
+
+str za0h.b[w12, #0], [x0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid matrix operand, expected za
+// CHECK-NEXT: str za0h.b[w12, #0], [x0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+str za3.s[w12, #0], [x0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid matrix operand, expected za
+// CHECK-NEXT: str za3.s[w12, #0], [x0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+// ------------------------------------------------------------------------- //
+// Invalid vector select register (expected: w12-w15)
+
+str za[w11, #0], [x0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid operand for instruction
+// CHECK-NEXT: str za[w11, #0], [x0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+str za[w16, #0], [x0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid operand for instruction
+// CHECK-NEXT: str za[w16, #0], [x0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+// ------------------------------------------------------------------------- //
+// Invalid vector select offset (expected: 0-15)
+
+str za[w12, #16], [x0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: immediate must be an integer in range [0, 15].
+// CHECK-NEXT: str za[w12, #16], [x0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+// ------------------------------------------------------------------------- //
+// Invalid memory operands
+
+str za[w12, #0], [w0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid operand for instruction
+// CHECK-NEXT: str za[w12, #0], [w0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+str za[w12, #0], [x0, #16, mul vl]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: immediate must be an integer in range [0, 15].
+// CHECK-NEXT: str za[w12, #0], [x0, #16, mul vl]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:
+
+str za[w12, #0], [x0, #0]
+// CHECK: [[@LINE-1]]:{{[0-9]+}}: error: invalid operand for instruction
+// CHECK-NEXT: str za[w12, #0], [x0, #0]
+// CHECK-NOT: [[@LINE-1]]:{{[0-9]+}}:

llvm/test/MC/AArch64/SME/str.s

+// RUN: llvm-mc -triple=aarch64 -show-encoding -mattr=+sme < %s \
+// RUN:        | FileCheck %s --check-prefixes=CHECK-ENCODING,CHECK-INST
+// RUN: not llvm-mc -triple=aarch64 -show-encoding < %s 2>&1 \
+// RUN:        | FileCheck %s --check-prefix=CHECK-ERROR
+// RUN: llvm-mc -triple=aarch64 -filetype=obj -mattr=+sme < %s \
+// RUN:        | llvm-objdump -d --mattr=+sme - | FileCheck %s --check-prefix=CHECK-INST
+// RUN: llvm-mc -triple=aarch64 -filetype=obj -mattr=+sme < %s \
+// RUN:        | llvm-objdump -d - | FileCheck %s --check-prefix=CHECK-UNKNOWN
+// Disassemble encoding and check the re-encoding (-show-encoding) matches.
+// RUN: llvm-mc -triple=aarch64 -show-encoding -mattr=+sme < %s \
+// RUN:        | sed '/.text/d' | sed 's/.*encoding: //g' \
+// RUN:        | llvm-mc -triple=aarch64 -mattr=+sme -disassemble -show-encoding \
+// RUN:        | FileCheck %s --check-prefixes=CHECK-ENCODING,CHECK-INST
+
+str     za[w12, #0], [x0]
+// CHECK-INST: str     za[w12, #0], [x0]
+// CHECK-ENCODING: [0x00,0x00,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 00 00 20 e1 <unknown>
+
+str     za[w14, #5], [x10, #5, mul vl]
+// CHECK-INST: str     za[w14, #5], [x10, #5, mul vl]
+// CHECK-ENCODING: [0x45,0x41,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 45 41 20 e1 <unknown>
+
+str     za[w15, #7], [x13, #7, mul vl]
+// CHECK-INST: str     za[w15, #7], [x13, #7, mul vl]
+// CHECK-ENCODING: [0xa7,0x61,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: a7 61 20 e1 <unknown>
+
+str     za[w15, #15], [sp, #15, mul vl]
+// CHECK-INST: str     za[w15, #15], [sp, #15, mul vl]
+// CHECK-ENCODING: [0xef,0x63,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: ef 63 20 e1 <unknown>
+
+str     za[w12, #5], [x17, #5, mul vl]
+// CHECK-INST: str     za[w12, #5], [x17, #5, mul vl]
+// CHECK-ENCODING: [0x25,0x02,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 25 02 20 e1 <unknown>
+
+str     za[w12, #1], [x1, #1, mul vl]
+// CHECK-INST: str     za[w12, #1], [x1, #1, mul vl]
+// CHECK-ENCODING: [0x21,0x00,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 21 00 20 e1 <unknown>
+
+str     za[w14, #8], [x19, #8, mul vl]
+// CHECK-INST: str     za[w14, #8], [x19, #8, mul vl]
+// CHECK-ENCODING: [0x68,0x42,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 68 42 20 e1 <unknown>
+
+str     za[w12, #0], [x12]
+// CHECK-INST: str     za[w12, #0], [x12]
+// CHECK-ENCODING: [0x80,0x01,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 80 01 20 e1 <unknown>
+
+str     za[w14, #1], [x1, #1, mul vl]
+// CHECK-INST: str     za[w14, #1], [x1, #1, mul vl]
+// CHECK-ENCODING: [0x21,0x40,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 21 40 20 e1 <unknown>
+
+str     za[w12, #13], [x22, #13, mul vl]
+// CHECK-INST: str     za[w12, #13], [x22, #13, mul vl]
+// CHECK-ENCODING: [0xcd,0x02,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: cd 02 20 e1 <unknown>
+
+str     za[w15, #2], [x9, #2, mul vl]
+// CHECK-INST: str     za[w15, #2], [x9, #2, mul vl]
+// CHECK-ENCODING: [0x22,0x61,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 22 61 20 e1 <unknown>
+
+str     za[w13, #7], [x12, #7, mul vl]
+// CHECK-INST: str     za[w13, #7], [x12, #7, mul vl]
+// CHECK-ENCODING: [0x87,0x21,0x20,0xe1]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 87 21 20 e1 <unknown>