[LoongArch] Add custom parser for atomic instructions' memory operand

[wangleiat]

Summary

In order to be compatible with the form of the atomic instruction in
GAS that accepts the fourth operand as 0 (i.e. am* $rd, $rk, $rj, 0),
we need to treat $rj, 0 as one operand, but only print $rj.

For this, the number of result operands of inline assembly memory
operand ZB constraint is modified to 2 (reg + 0).

Restrictions on register usage in am* instructions have also been
adjusted. When $rd is equal to $r0, the instruction must be
considered legal, because of some special usage like PseudoUNIMP.

[SixWeining]

LGTM but let's wait for others.

llvm/lib/Target/LoongArch/MCTargetDesc/LoongArchInstPrinter.cpp

64

Seems this line is useless.

[xen0n]

LGTM, thanks!

I've fixed the formatting of the summary for you (the ordered lists can't be used to write Chinese-style paragraphs) while at it.

[wangleiat]

In D139303#3980494, @xen0n wrote:

LGTM, thanks!

I've fixed the formatting of the summary for you (the ordered lists can't be used to write Chinese-style paragraphs) while at it.

Thank you very much!

llvm/lib/Target/LoongArch/MCTargetDesc/LoongArchInstPrinter.cpp

64

Thanks, i'll remove it.

[wangleiat]

Address @SixWeining's comment.

llvm/lib/Target/LoongArch/AsmParser/LoongArchAsmParser.cpp

 // Auto-generated instruction matching functions.
 #define GET_ASSEMBLER_HEADER
 #include "LoongArchGenAsmMatcher.inc"
 
   OperandMatchResultTy parseRegister(OperandVector &Operands);
   OperandMatchResultTy parseImmediate(OperandVector &Operands);
   OperandMatchResultTy parseOperandWithModifier(OperandVector &Operands);
   OperandMatchResultTy parseSImm26Operand(OperandVector &Operands);
+  OperandMatchResultTy parseAtomicMemOp(OperandVector &Operands);
 
   bool parseOperand(OperandVector &Operands, StringRef Mnemonic);
 
   // Helper to emit the sequence of instructions generated by the
   // "emitLoadAddress*" functions.
   void emitLAInstSeq(MCRegister DestReg, MCRegister TmpReg,
                      const MCExpr *Symbol, SmallVectorImpl<Inst> &Insts,
                      SMLoc IDLoc, MCStreamer &Out);

 public:
   LoongArchOperand(KindTy K) : MCParsedAsmOperand(), Kind(K) {}
 
   bool isToken() const override { return Kind == KindTy::Token; }
   bool isReg() const override { return Kind == KindTy::Register; }
   bool isImm() const override { return Kind == KindTy::Immediate; }
   bool isMem() const override { return false; }
   void setReg(MCRegister PhysReg) { Reg.RegNum = PhysReg; }
+  bool isGPR() const {
+    return Kind == KindTy::Register &&
+           LoongArchMCRegisterClasses[LoongArch::GPRRegClassID].contains(
+               Reg.RegNum);
+  }
 
   static bool evaluateConstantImm(const MCExpr *Expr, int64_t &Imm,
                                   LoongArchMCExpr::VariantKind &VK) {
     if (auto *LE = dyn_cast<LoongArchMCExpr>(Expr)) {
       VK = LE->getKind();
       return false;
     }
 
LoongArchAsmParser::parseSImm26Operand(OperandVector &Operands) {
   MCSymbol *Sym = getContext().getOrCreateSymbol(Identifier);
   Res = MCSymbolRefExpr::create(Sym, MCSymbolRefExpr::VK_None, getContext());
   Res = LoongArchMCExpr::create(Res, LoongArchMCExpr::VK_LoongArch_CALL,
                                 getContext());
   Operands.push_back(LoongArchOperand::createImm(Res, S, E));
   return MatchOperand_Success;
 }
 
+OperandMatchResultTy
+LoongArchAsmParser::parseAtomicMemOp(OperandVector &Operands) {
+  // Parse "$r*".
+  if (parseRegister(Operands) != MatchOperand_Success)
+    return MatchOperand_NoMatch;
+
+  // If there is a next operand and it is 0, ignore it. Otherwise print a
+  // diagnostic message.
+  if (getLexer().is(AsmToken::Comma)) {
+    getLexer().Lex(); // Consume comma token.
+    int64_t ImmVal;
+    SMLoc ImmStart = getLoc();
+    if (getParser().parseIntToken(ImmVal, "expected optional integer offset"))
+      return MatchOperand_ParseFail;
+    if (ImmVal) {
+      Error(ImmStart, "optional integer offset must be 0");
+      return MatchOperand_ParseFail;
+    }
+  }
+
+  return MatchOperand_Success;
+}
 /// Looks at a token type and creates the relevant operand from this
 /// information, adding to Operands. Return true upon an error.
 bool LoongArchAsmParser::parseOperand(OperandVector &Operands,
                                       StringRef Mnemonic) {
   // Check if the current operand has a custom associated parser, if so, try to
   // custom parse the operand, or fallback to the general approach.
   OperandMatchResultTy Result =
       MatchOperandParserImpl(Operands, Mnemonic, /*ParseForAllFeatures=*/true);
unsigned LoongArchAsmParser::checkTargetMatchPredicate(MCInst &Inst) {
   unsigned Opc = Inst.getOpcode();
   switch (Opc) {
   default:
     if (Opc >= LoongArch::AMADD_D && Opc <= LoongArch::AMXOR_W) {
       unsigned Rd = Inst.getOperand(0).getReg();
       unsigned Rk = Inst.getOperand(1).getReg();
       unsigned Rj = Inst.getOperand(2).getReg();
       if (Rd == Rk || Rd == Rj)
-        return Match_RequiresAMORdDifferRkRj;
+        return Rd == LoongArch::R0 ? Match_Success
+                                   : Match_RequiresAMORdDifferRkRj;
     }
     break;
   case LoongArch::PseudoLA_PCREL_LARGE:
   case LoongArch::PseudoLA_GOT_LARGE:
   case LoongArch::PseudoLA_TLS_IE_LARGE:
   case LoongArch::PseudoLA_TLS_LD_LARGE:
   case LoongArch::PseudoLA_TLS_GD_LARGE: {
     unsigned Rd = Inst.getOperand(0).getReg();

llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp

 bool LoongArchAsmPrinter::PrintAsmMemoryOperand(const MachineInstr *MI,
                                                 unsigned OpNo,
                                                 const char *ExtraCode,
                                                 raw_ostream &OS) {
   // TODO: handle extra code.
   if (ExtraCode)
     return true;
 
+  // We only support memory operands like "Base + Offset", where base must be a
+  // register, and offset can be a register or an immediate value.
   const MachineOperand &BaseMO = MI->getOperand(OpNo);
   // Base address must be a register.
   if (!BaseMO.isReg())
     return true;
   // Print the base address register.
   OS << "$" << LoongArchInstPrinter::getRegisterName(BaseMO.getReg());
-  // Print the offset register or immediate if has.
-  if (OpNo + 1 < MI->getNumOperands()) {
-    const MachineOperand &OffsetMO = MI->getOperand(OpNo + 1);
-    if (OffsetMO.isReg())
-      OS << ", $" << LoongArchInstPrinter::getRegisterName(OffsetMO.getReg());
-    else if (OffsetMO.isImm())
-      OS << ", " << OffsetMO.getImm();
-    else
-      return true;
-  }
+  // Print the offset operand.
+
   return false;
 }
 
 bool LoongArchAsmPrinter::runOnMachineFunction(MachineFunction &MF) {
   AsmPrinter::runOnMachineFunction(MF);
   return true;
 }
 
 // Force static initialization.
 extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeLoongArchAsmPrinter() {
   RegisterAsmPrinter<LoongArchAsmPrinter> X(getTheLoongArch32Target());
   RegisterAsmPrinter<LoongArchAsmPrinter> Y(getTheLoongArch64Target());
 }

llvm/lib/Target/LoongArch/LoongArchISelDAGToDAG.cpp

void LoongArchDAGToDAGISel::Select(SDNode *Node) {
   }
 
   // Select the default instruction.
   SelectCode(Node);
 }
 
 bool LoongArchDAGToDAGISel::SelectInlineAsmMemoryOperand(
     const SDValue &Op, unsigned ConstraintID, std::vector<SDValue> &OutOps) {
+  SDValue Base = Op;
+  SDValue Offset =
+      CurDAG->getTargetConstant(0, SDLoc(Op), Subtarget->getGRLenVT());
   switch (ConstraintID) {
   default:
     llvm_unreachable("unexpected asm memory constraint");
   // Reg+Reg addressing.
   case InlineAsm::Constraint_k:
-    OutOps.push_back(Op.getOperand(0));
-    OutOps.push_back(Op.getOperand(1));
-    return false;
+    Base = Op.getOperand(0);
+    Offset = Op.getOperand(1);
+    break;
   // Reg+simm12 addressing.
-  case InlineAsm::Constraint_m: {
-    SDValue Base = Op;
-    SDValue Offset =
-        CurDAG->getTargetConstant(0, SDLoc(Op), Subtarget->getGRLenVT());
+  case InlineAsm::Constraint_m:
     if (CurDAG->isBaseWithConstantOffset(Op)) {
       ConstantSDNode *CN = dyn_cast<ConstantSDNode>(Op.getOperand(1));
       if (isIntN(12, CN->getSExtValue())) {
         Base = Op.getOperand(0);
         Offset = CurDAG->getTargetConstant(CN->getZExtValue(), SDLoc(Op),
                                            Op.getValueType());
       }
     }
-    OutOps.push_back(Base);
-    OutOps.push_back(Offset);
-    return false;
-  }
+    break;
+  // Reg+0 addressing.
   case InlineAsm::Constraint_ZB:
-    OutOps.push_back(Op);
-    // No offset.
-    return false;
+    break;
   // Reg+(simm14<<2) addressing.
-  case InlineAsm::Constraint_ZC: {
-    SDValue Base = Op;
-    SDValue Offset =
-        CurDAG->getTargetConstant(0, SDLoc(Op), Subtarget->getGRLenVT());
+  case InlineAsm::Constraint_ZC:
     if (CurDAG->isBaseWithConstantOffset(Op)) {
       ConstantSDNode *CN = dyn_cast<ConstantSDNode>(Op.getOperand(1));
       if (isIntN(16, CN->getSExtValue()) &&
           isAligned(Align(4ULL), CN->getZExtValue())) {
         Base = Op.getOperand(0);
         Offset = CurDAG->getTargetConstant(CN->getZExtValue(), SDLoc(Op),
                                            Op.getValueType());
       }
     }
-    OutOps.push_back(Base);
-    OutOps.push_back(Offset);
-    return false;
-  }
-  }
-  return true;
-}
+    break;
+  }
 
 bool LoongArchDAGToDAGISel::SelectBaseAddr(SDValue Addr, SDValue &Base) {
   // If this is FrameIndex, select it directly. Otherwise just let it get
   // selected to a register independently.
   if (auto *FIN = dyn_cast<FrameIndexSDNode>(Addr))
     Base =
         CurDAG->getTargetFrameIndex(FIN->getIndex(), Subtarget->getGRLenVT());
   else

llvm/lib/Target/LoongArch/LoongArchInstrInfo.td

 class SImmAsmOperand<int width, string suffix = "">
     : ImmAsmOperand<"S", width, suffix> {
 }
 
 class UImmAsmOperand<int width, string suffix = "">
     : ImmAsmOperand<"U", width, suffix> {
 }
 
+// A parse method for "$r*" or "$r*, 0", where the 0 is be silently ignored.
+// Only used for "AM*" instructions, in order to be compatible with GAS.
+def AtomicMemAsmOperand : AsmOperandClass {
+  let Name = "AtomicMemAsmOperand";
+  let RenderMethod = "addRegOperands";
+  let PredicateMethod = "isGPR";
+  let ParserMethod = "parseAtomicMemOp";
+}
+
+def GPRMemAtomic : RegisterOperand<GPR> {
+  let ParserMatchClass = AtomicMemAsmOperand;
+  let PrintMethod = "printAtomicMemOp";
+}
+
 // A parameterized register class alternative to i32imm/i64imm from Target.td.
 def grlenimm : Operand<GRLenVT>;
 def imm32 : Operand<GRLenVT> {
   let ParserMatchClass = ImmAsmOperand<"", 32, "">;
 }
 
 def uimm2 : Operand<GRLenVT> {
   let ParserMatchClass = UImmAsmOperand<2>;
: Fmt2RI12<op, (outs), (ins GPR:$rd, GPR:$rj, simm12_addlike:$imm12), opstr,
                "$rd, $rj, $imm12">;
 class STORE_2RI14<bits<8> op, string opstr>
     : Fmt2RI14<op, (outs), (ins GPR:$rd, GPR:$rj, simm14_lsl2:$imm14), opstr,
                "$rd, $rj, $imm14">;
 } // mayStore = 1
 
 let mayLoad = 1, mayStore = 1, Constraints = "@earlyclobber $rd" in
 class AM_3R<bits<17> op, string opstr>
-    : Fmt3R<op, (outs GPR:$rd), (ins GPR:$rk, GPR:$rj), opstr, "$rd, $rk, $rj">;
+    : Fmt3R<op, (outs GPR:$rd), (ins GPR:$rk, GPRMemAtomic:$rj), opstr,
+            "$rd, $rk, $rj">;
 
 let mayLoad = 1 in
 class LLBase<bits<8> op, string opstr>
     : Fmt2RI14<op, (outs GPR:$rd), (ins GPR:$rj, simm14_lsl2:$imm14), opstr,
                "$rd, $rj, $imm14">;
 
 let mayStore = 1, Constraints = "$rd = $dst" in
 class SCBase<bits<8> op, string opstr>
def : Pat<(StoreOp (AddLike BaseAddr:$ptr, simm12:$imm12), (vt StTy:$val)),
             (Inst StTy:$val, BaseAddr:$ptr, simm12:$imm12)>;
 }
 
 defm : AtomicStPat<atomic_store_8, ST_B, GPR, GRLenVT>;
 defm : AtomicStPat<atomic_store_16, ST_H, GPR, GRLenVT>;
 defm : AtomicStPat<atomic_store_unordered_monotonic_32, ST_W, GPR, i32>,
                    Requires<[IsLA32]>;
 
-def PseudoAtomicStoreW : Pseudo<(outs GPR:$dst), (ins GPR:$rj, GPR:$rk)>,
-                                 PseudoInstExpansion<(AMSWAP_DB_W R0,
-                                                      GPR:$rk, GPR:$rj)>;
+def PseudoAtomicStoreW
+  : Pseudo<(outs GPR:$dst), (ins GPR:$rj, GPR:$rk)>,
+           PseudoInstExpansion<(AMSWAP_DB_W R0, GPR:$rk, GPRMemAtomic:$rj)>;
 
 def : Pat<(atomic_store_release_seqcst_32 GPR:$rj, GPR:$rk),
           (PseudoAtomicStoreW GPR:$rj, GPR:$rk)>;
 
 let Predicates = [IsLA64] in {
-def PseudoAtomicStoreD : Pseudo<(outs GPR:$dst), (ins GPR:$rj, GPR:$rk)>,
-                                 PseudoInstExpansion<(AMSWAP_DB_D R0,
-                                                      GPR:$rk, GPR:$rj)>;
+def PseudoAtomicStoreD
+  : Pseudo<(outs GPR:$dst), (ins GPR:$rj, GPR:$rk)>,
+           PseudoInstExpansion<(AMSWAP_DB_D R0, GPR:$rk, GPRMemAtomic:$rj)>;
 
 def : Pat<(atomic_store_release_seqcst_64 GPR:$rj, GPR:$rk),
           (PseudoAtomicStoreD GPR:$rj, GPR:$rk)>;
 
 defm : LdPat<atomic_load_64, LD_D>;
 defm : AtomicStPat<atomic_store_unordered_monotonic_32, ST_W, GPR, i64>;
 defm : AtomicStPat<atomic_store_unordered_monotonic_64, ST_D, GPR, i64>;
 } // Predicates = [IsLA64]

llvm/lib/Target/LoongArch/MCTargetDesc/LoongArchInstPrinter.h

 public:
   LoongArchInstPrinter(const MCAsmInfo &MAI, const MCInstrInfo &MII,
                        const MCRegisterInfo &MRI)
       : MCInstPrinter(MAI, MII, MRI) {}
 
   void printInst(const MCInst *MI, uint64_t Address, StringRef Annot,
                  const MCSubtargetInfo &STI, raw_ostream &O) override;
   void printRegName(raw_ostream &O, unsigned RegNo) const override;
+  void printAtomicMemOp(const MCInst *MI, unsigned OpNo,
+                        const MCSubtargetInfo &STI, raw_ostream &O);
 
   // Autogenerated by tblgen.
   std::pair<const char *, uint64_t> getMnemonic(const MCInst *MI) override;
   void printInstruction(const MCInst *MI, uint64_t Address,
                         const MCSubtargetInfo &STI, raw_ostream &O);
   bool printAliasInstr(const MCInst *MI, uint64_t Address,
                        const MCSubtargetInfo &STI, raw_ostream &O);
   void printCustomAliasOperand(const MCInst *MI, uint64_t Address,

llvm/lib/Target/LoongArch/MCTargetDesc/LoongArchInstPrinter.cpp

if (MO.isImm()) {
     O << MO.getImm();
     return;
   }
 
   assert(MO.isExpr() && "Unknown operand kind in printOperand");
   MO.getExpr()->print(O, &MAI);
 }
 
+void LoongArchInstPrinter::printAtomicMemOp(const MCInst *MI, unsigned OpNo,
+                                            const MCSubtargetInfo &STI,
+                                            raw_ostream &O) {
+  const MCOperand &MO = MI->getOperand(OpNo);
+  assert(MO.isReg() && "printAtomicMemOp can only print register operands");
+  printRegName(O, MO.getReg());
+}
+
 const char *LoongArchInstPrinter::getRegisterName(unsigned RegNo) {
   // Default print reg alias name
   return getRegisterName(RegNo, LoongArch::RegAliasName);
 }

llvm/test/CodeGen/LoongArch/inline-asm-constraint-ZB.ll

 ; ASM-NEXT:    #NO_APP
 ; ASM-NEXT:    ret
   %1 = getelementptr inbounds i8, ptr %p, i32 %idx
 ;; Make sure machine instr with this "ZB" constraint is printed correctly.
 ; MACHINE-INSTR: INLINEASM{{.*}}[mem:ZB]
   call void asm "amswap.w $$r12, $$r13, $0", "*^ZB"(ptr elementtype(i32) %1)
   ret void
 }
+
+define void @ZB_Input_Output(ptr %p) nounwind {
+; ASM-LABEL: ZB_Input_Output:
+; ASM:       # %bb.0:
+; ASM-NEXT:    #APP
+; ASM-NEXT:    amadd_db.d $zero, $t1, $a0
+; ASM-NEXT:    #NO_APP
+; ASM-NEXT:    ret
+;; Make sure machine instr with this "ZB" constraint is printed correctly.
+; MACHINE-INSTR: INLINEASM{{.*}}[mem:ZB], %0:gpr, 0
+  call void asm "amadd_db.d $$zero, $$r13, $0", "=*^ZB,*^ZB,~{memory}"(ptr elementtype(i64) %p, ptr elementtype(i64) %p)
+  ret void
+}

llvm/test/MC/LoongArch/Basic/Integer/atomic.s

 #############################################################
 ## Instructions only for loongarch64
 #############################################################
 
 .ifdef LA64
 
 # CHECK64-ASM-AND-OBJ: amswap.w $a2, $t0, $s1
 # CHECK64-ASM: encoding: [0x06,0x33,0x60,0x38]
+amswap.w $a2, $t0, $s1, 0
+
+# CHECK64-ASM-AND-OBJ: amswap.w $zero, $t0, $zero
+# CHECK64-ASM: encoding: [0x00,0x30,0x60,0x38]
+amswap.w $zero, $t0, $zero
+
+# CHECK64-ASM-AND-OBJ: amadd_db.w $zero, $zero, $a1
+# CHECK64-ASM: encoding: [0xa0,0x00,0x6a,0x38]
+amadd_db.w $zero, $zero, $a1
+
+# CHECK64-ASM-AND-OBJ: amswap.w $a2, $t0, $s1
+# CHECK64-ASM: encoding: [0x06,0x33,0x60,0x38]
 amswap.w $a2, $t0, $s1
 
 # CHECK64-ASM-AND-OBJ: amswap.d $tp, $t2, $fp
 # CHECK64-ASM: encoding: [0xc2,0xba,0x60,0x38]
 amswap.d $tp, $t2, $fp
 
 # CHECK64-ASM-AND-OBJ: amadd.w $a4, $t0, $r21
 # CHECK64-ASM: encoding: [0xa8,0x32,0x61,0x38]

llvm/test/MC/LoongArch/Basic/Integer/invalid64.s

 bstrins.d $a0, $a0, 1, 2
 # CHECK:            ^~~~
 
 # CHECK: :[[#@LINE+1]]:22: error: msb is less than lsb
 bstrpick.d $a0, $a0, 32, 63
 # CHECK:             ^~~~~~
 
 # CHECK: :[[#@LINE+1]]:10: error: $rd must be different from both $rk and $rj
-amadd.d $zero, $zero, $zero
-# CHECK: :[[#@LINE+1]]:10: error: $rd must be different from both $rk and $rj
-ammin.w $zero, $zero, $a0
-# CHECK: :[[#@LINE+1]]:10: error: $rd must be different from both $rk and $rj
-amxor.w $zero, $a0, $zero
-
-# CHECK: :[[#@LINE+1]]:10: error: $rd must be different from both $rk and $rj
 amadd.d $a0, $a0, $a0
 # CHECK: :[[#@LINE+1]]:10: error: $rd must be different from both $rk and $rj
 ammin.w $a0, $a0, $a1
 # CHECK: :[[#@LINE+1]]:10: error: $rd must be different from both $rk and $rj
 amxor.w $a0, $a1, $a0
+
+# CHECK: :[[#@LINE+1]]:24: error: expected optional integer offset
+amadd.d $a0, $a1, $a2, $a3
+# CHECK: :[[#@LINE+1]]:24: error: optional integer offset must be 0
+amadd.d $a0, $a1, $a2, 1