[AArch64][AsmParser] Extend RegOp with integrated 'shift/extend'.

[sdesmalen]

Summary

In some cases the shift/extend needs to be explicitly parsed together
with the register, rather than as a separate operand. This is needed
for addressing modes where the instruction as a whole dictates the
scaling/extend, rather than specific bits in the instruction.
By parsing them as a single operand, we avoid the need to pass an
extra operand in all CodeGen patterns (because all operands need to
have an associated value), and we avoid the need to update TableGen to
accept operands that have no associated bits in the instruction.

An added benefit of parsing them together is that the assembler
can give a sensible diagnostic if the scaling is not correct.

This is patch [2/4] in a series to add assembler/disassembler support for
SVE's contiguous LD1 (scalar+scalar) instructions:

• Patch [1/4]: https://reviews.llvm.org/D45687
• Patch [2/4]: https://reviews.llvm.org/D45688
• Patch [3/4]: https://reviews.llvm.org/D45689
• Patch [4/4]: https://reviews.llvm.org/D45690

[SjoerdMeijer]

Looks reasonable to me.

[fhahn]

LGTM

lib/Target/AArch64/AArch64RegisterInfo.td

174

This change seems unrelated to handling "integrated 'shift/extend'", but I do not think it is worth updating the diff just for that.

lib/Target/AArch64/AArch64RegisterInfo.td

 }
 
 def GPR32sponly : RegisterClass<"AArch64", [i32], 32, (add WSP)>;
 def GPR64sponly : RegisterClass<"AArch64", [i64], 64, (add SP)>;
 
 def GPR64spPlus0Operand : AsmOperandClass {
   let Name = "GPR64sp0";
   let RenderMethod = "addRegOperands";
+  let PredicateMethod = "isGPR64<AArch64::GPR64spRegClassID>";
   let ParserMethod = "tryParseGPR64sp0Operand";
 }
 
 def GPR64sp0 : RegisterOperand<GPR64sp> {
   let ParserMatchClass = GPR64spPlus0Operand;
 }
 
 // GPR32/GPR64 but with zero-register substitution enabled.

lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp

#include "AArch64GenAsmMatcher.inc"
   OperandMatchResultTy tryParseAdrpLabel(OperandVector &Operands);
   OperandMatchResultTy tryParseAdrLabel(OperandVector &Operands);
   OperandMatchResultTy tryParseFPImm(OperandVector &Operands);
   OperandMatchResultTy tryParseAddSubImm(OperandVector &Operands);
   OperandMatchResultTy tryParseGPR64sp0Operand(OperandVector &Operands);
   bool tryParseNeonVectorRegister(OperandVector &Operands);
   OperandMatchResultTy tryParseVectorIndex(OperandVector &Operands);
   OperandMatchResultTy tryParseGPRSeqPair(OperandVector &Operands);
+  template <bool ParseShiftExtend>
+  OperandMatchResultTy tryParseGPROperand(OperandVector &Operands);
   template <bool ParseSuffix>
   OperandMatchResultTy tryParseSVEDataVector(OperandVector &Operands);
   OperandMatchResultTy tryParseSVEPredicateVector(OperandVector &Operands);
   template <RegKind VectorKind>
   OperandMatchResultTy tryParseVectorList(OperandVector &Operands,
                                           bool ExpectMatch = false);
   OperandMatchResultTy tryParseSVEPattern(OperandVector &Operands);
 
private:
   SMLoc StartLoc, EndLoc;
 
   struct TokOp {
     const char *Data;
     unsigned Length;
     bool IsSuffix; // Is the operand actually a suffix on the mnemonic.
   };
 
+  // Separate shift/extend operand.
+  struct ShiftExtendOp {
+    AArch64_AM::ShiftExtendType Type;
+    unsigned Amount;
+    bool HasExplicitAmount;
+  };
+
   struct RegOp {
     unsigned RegNum;
     RegKind Kind;
-
     int ElementWidth;
+
+    // In some cases the shift/extend needs to be explicitly parsed together
+    // with the register, rather than as a separate operand. This is needed
+    // for addressing modes where the instruction as a whole dictates the
+    // scaling/extend, rather than specific bits in the instruction.
+    // By parsing them as a single operand, we avoid the need to pass an
+    // extra operand in all CodeGen patterns (because all operands need to
+    // have an associated value), and we avoid the need to update TableGen to
+    // accept operands that have no associated bits in the instruction.
+    //
+    // An added benefit of parsing them together is that the assembler
+    // can give a sensible diagnostic if the scaling is not correct.
+    //
+    // The default is 'lsl #0' (HasExplicitAmount = false) if no
+    // ShiftExtend is specified.
+    ShiftExtendOp ShiftExtend;
   };
 
   struct VectorListOp {
     unsigned RegNum;
     unsigned Count;
     unsigned NumElements;
     unsigned ElementWidth;
     RegKind  RegisterKind;
private:
   };
 
   struct PSBHintOp {
     const char *Data;
     unsigned Length;
     unsigned Val;
   };
 
-  struct ShiftExtendOp {
-    AArch64_AM::ShiftExtendType Type;
-    unsigned Amount;
-    bool HasExplicitAmount;
-  };
-
   struct ExtendOp {
     unsigned Val;
   };
 
   union {
     struct TokOp Tok;
     struct RegOp Reg;
     struct VectorListOp VectorList;
public:
   }
 
   StringRef getPrefetchName() const {
     assert(Kind == k_Prefetch && "Invalid access!");
     return StringRef(Prefetch.Data, Prefetch.Length);
   }
 
   AArch64_AM::ShiftExtendType getShiftExtendType() const {
-    assert(Kind == k_ShiftExtend && "Invalid access!");
+    if (Kind == k_ShiftExtend)
       return ShiftExtend.Type;
+    if (Kind == k_Register)
+      return Reg.ShiftExtend.Type;
+    llvm_unreachable("Invalid access!");
   }
 
   unsigned getShiftExtendAmount() const {
-    assert(Kind == k_ShiftExtend && "Invalid access!");
+    if (Kind == k_ShiftExtend)
       return ShiftExtend.Amount;
+    if (Kind == k_Register)
+      return Reg.ShiftExtend.Amount;
+    llvm_unreachable("Invalid access!");
   }
 
   bool hasShiftExtendAmount() const {
-    assert(Kind == k_ShiftExtend && "Invalid access!");
+    if (Kind == k_ShiftExtend)
       return ShiftExtend.HasExplicitAmount;
+    if (Kind == k_Register)
+      return Reg.ShiftExtend.HasExplicitAmount;
+    llvm_unreachable("Invalid access!");
   }
 
   bool isImm() const override { return Kind == k_Immediate; }
   bool isMem() const override { return false; }
 
   template <int Width> bool isSImm() const { return isSImmScaled<Width, 1>(); }
 
   template <int Bits, int Scale> bool isSImmScaled() const {
public:
   }
 
   bool isXSeqPair() const {
     return Kind == k_Register && Reg.Kind == RegKind::Scalar &&
            AArch64MCRegisterClasses[AArch64::XSeqPairsClassRegClassID].contains(
                Reg.RegNum);
   }
 
-  bool isGPR64sp0() const {
-    return Kind == k_Register && Reg.Kind == RegKind::Scalar &&
-      AArch64MCRegisterClasses[AArch64::GPR64spRegClassID].contains(Reg.RegNum);
-  }
-
   template<int64_t Angle, int64_t Remainder>
   bool isComplexRotation() const {
     if (!isImm()) return false;
 
     const MCConstantExpr *CE = dyn_cast<MCConstantExpr>(getImm());
     if (!CE) return false;
     uint64_t Value = CE->getValue();
 
     return (Value % Angle == Remainder && Value <= 270);
   }
 
+  template <unsigned RegClassID> bool isGPR64() const {
+    return Kind == k_Register && Reg.Kind == RegKind::Scalar &&
+           AArch64MCRegisterClasses[RegClassID].contains(getReg());
+  }
+
+  template <unsigned RegClassID, int ExtWidth>
+  bool isGPR64WithShiftExtend() const {
+    if (!isGPR64<RegClassID>())
+      return false;
+
+    return getShiftExtendType() == AArch64_AM::LSL &&
+           getShiftExtendAmount() == Log2_32(ExtWidth / 8);
+  }
+
   /// Is this a vector list with the type implicit (presumably attached to the
   /// instruction itself)?
   template <RegKind VectorKind, unsigned NumRegs>
   bool isImplicitlyTypedVectorList() const {
     return Kind == k_VectorList && VectorList.Count == NumRegs &&
            VectorList.NumElements == 0 &&
            VectorList.RegisterKind == VectorKind;
   }
CreateToken(StringRef Str, bool IsSuffix, SMLoc S, MCContext &Ctx) {
     Op->Tok.Length = Str.size();
     Op->Tok.IsSuffix = IsSuffix;
     Op->StartLoc = S;
     Op->EndLoc = S;
     return Op;
   }
 
   static std::unique_ptr<AArch64Operand>
-  CreateReg(unsigned RegNum, RegKind Kind, SMLoc S, SMLoc E, MCContext &Ctx) {
+  CreateReg(unsigned RegNum, RegKind Kind, SMLoc S, SMLoc E, MCContext &Ctx,
+            AArch64_AM::ShiftExtendType ExtTy = AArch64_AM::LSL,
+            unsigned ShiftAmount = 0,
+            unsigned HasExplicitAmount = false) {
     auto Op = make_unique<AArch64Operand>(k_Register, Ctx);
     Op->Reg.RegNum = RegNum;
     Op->Reg.Kind = Kind;
+    Op->Reg.ElementWidth = 0;
+    Op->Reg.ShiftExtend.Type = ExtTy;
+    Op->Reg.ShiftExtend.Amount = ShiftAmount;
+    Op->Reg.ShiftExtend.HasExplicitAmount = HasExplicitAmount;
     Op->StartLoc = S;
     Op->EndLoc = E;
     return Op;
   }
 
   static std::unique_ptr<AArch64Operand>
   CreateVectorReg(unsigned RegNum, RegKind Kind, unsigned ElementWidth,
-                  SMLoc S, SMLoc E, MCContext &Ctx) {
+                  SMLoc S, SMLoc E, MCContext &Ctx,
+                  AArch64_AM::ShiftExtendType ExtTy = AArch64_AM::LSL,
+                  unsigned ShiftAmount = 0,
+                  unsigned HasExplicitAmount = false) {
     assert((Kind == RegKind::NeonVector || Kind == RegKind::SVEDataVector ||
             Kind == RegKind::SVEPredicateVector) &&
            "Invalid vector kind");
-    auto Op = make_unique<AArch64Operand>(k_Register, Ctx);
-    Op->Reg.RegNum = RegNum;
+    auto Op = CreateReg(RegNum, Kind, S, E, Ctx, ExtTy, ShiftAmount,
+                        HasExplicitAmount);
     Op->Reg.ElementWidth = ElementWidth;
-    Op->Reg.Kind = Kind;
-    Op->StartLoc = S;
-    Op->EndLoc = E;
     return Op;
   }
 
   static std::unique_ptr<AArch64Operand>
   CreateVectorList(unsigned RegNum, unsigned Count, unsigned NumElements,
                    unsigned ElementWidth, RegKind RegisterKind, SMLoc S, SMLoc E,
                    MCContext &Ctx) {
     auto Op = make_unique<AArch64Operand>(k_VectorList, Ctx);
case k_ShiftedImm: {
     OS << "<shiftedimm ";
     OS << *getShiftedImmVal();
     OS << ", lsl #" << AArch64_AM::getShiftValue(Shift) << ">";
     break;
   }
   case k_CondCode:
     OS << "<condcode " << getCondCode() << ">";
     break;
-  case k_Register:
-    OS << "<register " << getReg() << ">";
-    break;
   case k_VectorList: {
     OS << "<vectorlist ";
     unsigned Reg = getVectorListStart();
     for (unsigned i = 0, e = getVectorListCount(); i != e; ++i)
       OS << Reg + i << " ";
     OS << ">";
     break;
   }
if (!Name.empty())
       OS << "<prfop " << Name << ">";
     else
       OS << "<prfop invalid #" << getPrefetch() << ">";
     break;
   }
   case k_PSBHint:
     OS << getPSBHintName();
     break;
+  case k_Register:
+    OS << "<register " << getReg() << ">";
+    if (!getShiftExtendAmount() && !hasShiftExtendAmount())
+      break;
+    LLVM_FALLTHROUGH;
   case k_ShiftExtend:
     OS << "<" << AArch64_AM::getShiftExtendName(getShiftExtendType()) << " #"
        << getShiftExtendAmount();
     if (!hasShiftExtendAmount())
       OS << "<imp>";
     OS << '>';
     break;
   }
Operands.push_back(
     AArch64Operand::CreateToken(ZM, false, getLoc(), getContext()));
 
   Parser.Lex(); // Eat zero/merge token.
   return MatchOperand_Success;
 }
 
 /// parseRegister - Parse a register operand.
 bool AArch64AsmParser::parseRegister(OperandVector &Operands) {
-  SMLoc S = getLoc();
   // Try for a Neon vector register.
   if (!tryParseNeonVectorRegister(Operands))
     return false;
 
-  // Try for a scalar register.
-  unsigned Reg;
-  if (tryParseScalarRegister(Reg) != MatchOperand_Success)
-    return true;
-  Operands.push_back(AArch64Operand::CreateReg(Reg, RegKind::Scalar, S,
-      getLoc(), getContext()));
-
+  // Otherwise try for a scalar register.
+  if (tryParseGPROperand<false>(Operands) == MatchOperand_Success)
     return false;
+
+  return true;
 }
 
 bool AArch64AsmParser::parseSymbolicImmVal(const MCExpr *&ImmVal) {
   MCAsmParser &Parser = getParser();
   bool HasELFModifier = false;
   AArch64MCExpr::VariantKind RefKind;
 
   if (parseOptionalToken(AsmToken::Colon)) {
if (getParser().parseExpression(ImmVal) || !isa<MCConstantExpr>(ImmVal) ||
     return MatchOperand_ParseFail;
   }
 
   Operands.push_back(AArch64Operand::CreateReg(
       RegNum, RegKind::Scalar, StartLoc, getLoc(), getContext()));
   return MatchOperand_Success;
 }
 
+template <bool ParseShiftExtend>
+OperandMatchResultTy
+AArch64AsmParser::tryParseGPROperand(OperandVector &Operands) {
+  SMLoc StartLoc = getLoc();
+
+  unsigned RegNum;
+  OperandMatchResultTy Res = tryParseScalarRegister(RegNum);
+  if (Res != MatchOperand_Success)
+    return Res;
+
+  // No shift/extend is the default.
+  if (!ParseShiftExtend || getParser().getTok().isNot(AsmToken::Comma)) {
+    Operands.push_back(AArch64Operand::CreateReg(
+        RegNum, RegKind::Scalar, StartLoc, getLoc(), getContext()));
+    return MatchOperand_Success;
+  }
+
+  // Eat the comma
+  getParser().Lex();
+
+  // Match the shift
+  SmallVector<std::unique_ptr<MCParsedAsmOperand>, 1> ExtOpnd;
+  Res = tryParseOptionalShiftExtend(ExtOpnd);
+  if (Res != MatchOperand_Success)
+    return Res;
+
+  auto Ext = static_cast<AArch64Operand*>(ExtOpnd.back().get());
+  Operands.push_back(AArch64Operand::CreateReg(RegNum, RegKind::Scalar,
+                     StartLoc, Ext->getEndLoc(), getContext(),
+                     Ext->getShiftExtendType(), Ext->getShiftExtendAmount(),
+                     Ext->hasShiftExtendAmount()));
+
+  return MatchOperand_Success;
+}
+
 bool AArch64AsmParser::parseOptionalMulVl(OperandVector &Operands) {
   MCAsmParser &Parser = getParser();
 
   // Some SVE instructions have a decoration after the immediate, i.e.
   // "mul vl". We parse them here and add tokens, which must be present in the
   // asm string in the tablegen instruction.
   if (!Parser.getTok().getString().equals_lower("mul") ||
       !Parser.getLexer().peekTok().getString().equals_lower("vl"))