[AArch64][SME][NFC] Add implicit operands for SME instructions in the disassembly.

This patch simplifies the switch statement in getInstruction to add
implicit operands (register ZA and Immediate  equal to zero)
in the SME operands when disassembly.

The register ZA and the zero immediate  can be added by checking the operand
in MCInstDesc.

Differential Revision: https://reviews.llvm.org/D125534

llvm/lib/Target/AArch64/AArch64InstrFormats.td

@@ -1278,8 +1278,12 @@ def VectorIndexHOperand : AsmVectorIndex<0, 7>;
 def VectorIndexSOperand : AsmVectorIndex<0, 3>;
 def VectorIndexDOperand : AsmVectorIndex<0, 1>;
 
-defm VectorIndex0 : VectorIndex<i64, VectorIndex0Operand,
+let OperandNamespace = "AArch64" in {
+  let OperandType = "OPERAND_IMPLICIT_IMM_0" in {
+    defm VectorIndex0 : VectorIndex<i64, VectorIndex0Operand,
                                 [{ return ((uint64_t)Imm) == 0; }]>;
+  }
+}
 defm VectorIndex1 : VectorIndex<i64, VectorIndex1Operand,
                                 [{ return ((uint64_t)Imm) == 1; }]>;
 defm VectorIndexB : VectorIndex<i64, VectorIndexBOperand,
@@ -1329,6 +1333,8 @@ def sme_elm_idx0_0 : Operand<i64>, ImmLeaf<i64, [{
 }]> {
   let ParserMatchClass = Imm0_0Operand;
   let PrintMethod = "printMatrixIndex";
+  let OperandNamespace = "AArch64";
+  let OperandType = "OPERAND_IMPLICIT_IMM_0";
 }
 def sme_elm_idx0_1 : Operand<i64>, ImmLeaf<i64, [{
   return ((uint64_t)Imm) <= 1;

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

@@ -19,6 +19,7 @@
 #include "llvm/MC/MCDecoderOps.h"
 #include "llvm/MC/MCDisassembler/MCRelocationInfo.h"
 #include "llvm/MC/MCInst.h"
+#include "llvm/MC/MCInstrDesc.h"
 #include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/MC/MCSubtargetInfo.h"
 #include "llvm/MC/TargetRegistry.h"
@@ -283,7 +284,8 @@ static bool Check(DecodeStatus &Out, DecodeStatus In) {
 static MCDisassembler *createAArch64Disassembler(const Target &T,
                                                const MCSubtargetInfo &STI,
                                                MCContext &Ctx) {
-  return new AArch64Disassembler(STI, Ctx);
+
+  return new AArch64Disassembler(STI, Ctx, T.createMCInstrInfo());
 }
 
 DecodeStatus AArch64Disassembler::getInstruction(MCInst &MI, uint64_t &Size,
@@ -308,67 +310,37 @@ DecodeStatus AArch64Disassembler::getInstruction(MCInst &MI, uint64_t &Size,
     DecodeStatus Result =
         decodeInstruction(Table, MI, Insn, Address, this, STI);
 
-    switch (MI.getOpcode()) {
-    default:
-      break;
+    const MCInstrDesc &Desc = MCII->get(MI.getOpcode());
+
     // 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.
+    // operand for the accumulator (ZA) or implicit immediate zero which isn't
+    // encoded, manually insert operand.
+    for (unsigned i = 0; i < Desc.getNumOperands(); i++) {
+      if (Desc.OpInfo[i].OperandType == MCOI::OPERAND_REGISTER) {
+        switch (Desc.OpInfo[i].RegClass) {
+        default:
+          break;
+        case AArch64::MPRRegClassID:
+          MI.insert(MI.begin() + i, MCOperand::createReg(AArch64::ZA));
+          break;
+        case AArch64::MPR8RegClassID:
+          MI.insert(MI.begin() + i, MCOperand::createReg(AArch64::ZAB0));
+          break;
+        }
+      } else if (Desc.OpInfo[i].OperandType ==
+                 AArch64::OPERAND_IMPLICIT_IMM_0) {
+        MI.insert(MI.begin() + i, MCOperand::createImm(0));
+      }
+    }
+
+    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);
-      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:
-    case AArch64::INSERT_MXIPZ_H_B:
-    case AArch64::INSERT_MXIPZ_V_B:
-      // e.g.
-      // MOVA ZA0<HV>.B[<Ws>, <imm>], <Pg>/M, <Zn>.B
-      //      ^ insert implicit 8-bit element tile
-      MI.insert(MI.begin(), MCOperand::createReg(AArch64::ZAB0));
-      break;
-    case AArch64::EXTRACT_ZPMXI_H_B:
-    case AArch64::EXTRACT_ZPMXI_V_B:
-      // MOVA <Zd>.B, <Pg>/M, ZA0<HV>.B[<Ws>, <imm>]
-      //                      ^ insert implicit 8-bit element tile
-      MI.insert(MI.begin()+2, MCOperand::createReg(AArch64::ZAB0));
-      break;
-    case AArch64::LD1_MXIPXX_H_Q:
-    case AArch64::LD1_MXIPXX_V_Q:
-    case AArch64::ST1_MXIPXX_H_Q:
-    case AArch64::ST1_MXIPXX_V_Q:
-      // 128-bit load/store have implicit zero vector index.
-      MI.insert(MI.begin()+2, MCOperand::createImm(0));
-      break;
-    // 128-bit mova have implicit zero vector index.
-    case AArch64::INSERT_MXIPZ_H_Q:
-    case AArch64::INSERT_MXIPZ_V_Q:
-      MI.insert(MI.begin()+2, MCOperand::createImm(0));
-      break;
-    case AArch64::EXTRACT_ZPMXI_H_Q:
-    case AArch64::EXTRACT_ZPMXI_V_Q:
-      MI.addOperand(MCOperand::createImm(0));
-      break;
-    case AArch64::SMOVvi8to32_idx0:
-    case AArch64::SMOVvi8to64_idx0:
-    case AArch64::SMOVvi16to32_idx0:
-    case AArch64::SMOVvi16to64_idx0:
-    case AArch64::SMOVvi32to64_idx0:
-    case AArch64::UMOVvi8_idx0:
-    case AArch64::UMOVvi16_idx0:
-    case AArch64::UMOVvi32_idx0:
-    case AArch64::UMOVvi64_idx0:
-      MI.addOperand(MCOperand::createImm(0));
-      break;
     }
 
     if (Result != MCDisassembler::Fail)

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

@@ -13,13 +13,17 @@
 #define LLVM_LIB_TARGET_AARCH64_DISASSEMBLER_AARCH64DISASSEMBLER_H
 
 #include "llvm/MC/MCDisassembler/MCDisassembler.h"
+#include "llvm/MC/MCInstrInfo.h"
 
 namespace llvm {
 
 class AArch64Disassembler : public MCDisassembler {
+  std::unique_ptr<const MCInstrInfo> const MCII;
+
 public:
-  AArch64Disassembler(const MCSubtargetInfo &STI, MCContext &Ctx)
-    : MCDisassembler(STI, Ctx) {}
+  AArch64Disassembler(const MCSubtargetInfo &STI, MCContext &Ctx,
+                      MCInstrInfo const *MCII)
+      : MCDisassembler(STI, Ctx), MCII(MCII) {}
 
   ~AArch64Disassembler() override = default;
 

llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCTargetDesc.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_LIB_TARGET_AARCH64_MCTARGETDESC_AARCH64MCTARGETDESC_H
 #define LLVM_LIB_TARGET_AARCH64_MCTARGETDESC_AARCH64MCTARGETDESC_H
 
+#include "llvm/MC/MCInstrDesc.h"
 #include "llvm/Support/DataTypes.h"
 
 #include <memory>
@@ -64,6 +65,12 @@ bool isQForm(const MCInst &MI, const MCInstrInfo *MCII);
 bool isFpOrNEON(const MCInst &MI, const MCInstrInfo *MCII);
 }
 
+namespace AArch64 {
+enum OperandType {
+  OPERAND_IMPLICIT_IMM_0 = MCOI::OPERAND_FIRST_TARGET,
+};
+} // namespace AArch64
+
 } // End llvm namespace
 
 // Defines symbolic names for AArch64 registers.  This defines a mapping from