[llvm] [AArch64] Provide a custom decoder for LDR_ZA/STR_ZA (PR #156363)

[Sergei Barannikov via llvm-commits]

https://github.com/s-barannikov updated https://github.com/llvm/llvm-project/pull/156363

>From 92ae372fd41e9ae7cd38c754a9a5d19e71b4372e Mon Sep 17 00:00:00 2001
From: Sergei Barannikov <barannikov88 at gmail.com>
Date: Mon, 1 Sep 2025 20:27:48 +0300
Subject: [PATCH] [AArch64] Provide a custom decoder for LDR_ZA/STR_ZA

These are the only instructions that encode two operands in the same
field. Instead of fixing them after they have been incorrectly decoded,
provide a custom decoder.
---
 .../Disassembler/AArch64Disassembler.cpp      | 29 ++++++++++++-------
 llvm/lib/Target/AArch64/SMEInstrFormats.td    |  4 +++
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Target/AArch64/Disassembler/AArch64Disassembler.cpp b/llvm/lib/Target/AArch64/Disassembler/AArch64Disassembler.cpp
index 23e46b84f6278..8c1e9f61693fb 100644
--- a/llvm/lib/Target/AArch64/Disassembler/AArch64Disassembler.cpp
+++ b/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;
   }
diff --git a/llvm/lib/Target/AArch64/SMEInstrFormats.td b/llvm/lib/Target/AArch64/SMEInstrFormats.td
index b3005d5120229..40ec371fe79d3 100644
--- a/llvm/lib/Target/AArch64/SMEInstrFormats.td
+++ b/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;