[AArch64] Provide a custom decoder for LDR_ZA/STR_ZA

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

@@ -1563,6 +1563,25 @@ static DecodeStatus DecodePRFMRegInstruction(MCInst &Inst, uint32_t insn,
   return Success;
 }
 
+static DecodeStatus
+DecodeSMESpillFillInstruction(MCInst &Inst, uint32_t Bits, uint64_t Addr,
+                              const MCDisassembler *Decoder) {
+  unsigned RvBits = fieldFromInstruction(Bits, 13, 2);
+  unsigned RnBits = fieldFromInstruction(Bits, 5, 5);
+  unsigned Imm4Bits = fieldFromInstruction(Bits, 0, 4);
+
+  DecodeSimpleRegisterClass<AArch64::MatrixIndexGPR32_12_15RegClassID, 0, 4>(
+      Inst, RvBits, Addr, Decoder);
+  Inst.addOperand(MCOperand::createImm(Imm4Bits));
+  DecodeSimpleRegisterClass<AArch64::GPR64spRegClassID, 0, 32>(Inst, RnBits,
+                                                               Addr, Decoder);
+  // Spill and fill instructions have a single immediate used for both
+  // the vector select offset and optional memory offset. Replicate
+  // the decoded immediate.
+  Inst.addOperand(MCOperand::createImm(Imm4Bits));
+  return Success;
+}
+
 #include "AArch64GenDisassemblerTables.inc"
 #include "AArch64GenInstrInfo.inc"
 
@@ -1621,16 +1640,6 @@ DecodeStatus AArch64Disassembler::getInstruction(MCInst &MI, uint64_t &Size,
       }
     }
 
-    if (MI.getOpcode() == AArch64::LDR_ZA ||
-        MI.getOpcode() == AArch64::STR_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);
-    }
-
     if (Result != MCDisassembler::Fail)
       return Result;
   }

llvm/lib/Target/AArch64/SMEInstrFormats.td

@@ -1108,6 +1108,10 @@ 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<[]> {
+  // 'offset' operand is encoded in the same bits as 'imm4'. There is currently
+  // no way to tell TableGen about this.
+  let DecoderMethod = "DecodeSMESpillFillInstruction";
+  bits<0> ZAt;
   bits<2> Rv;
   bits<5> Rn;
   bits<4> imm4;

llvm/test/MC/Disassembler/AArch64/armv9.2a-sme.txt

@@ -0,0 +1,4 @@
+# RUN: llvm-mc -triple=aarch64 -mattr=+sme -disassemble %s | FileCheck %s
+
+[0x45,0x41,0x00,0xe1] # CHECK: ldr za[w14, 5], [x10, #5, mul vl]
+[0x45,0x41,0x20,0xe1] # CHECK: str za[w14, 5], [x10, #5, mul vl]