AMDGPU/SI: Assembler: improvements to support trap handlers.

[nhaustov]

Summary

Add ParseAMDGPURegister which can be invoked recursively for parsing lists.
Rename getRegForName to getSpecialRegForName.
Support legacy SP3 register list syntax: [s2,s3,s4,s5] or [flat_scratch_lo,flat_scratch_hi].
Add 64-bit registers TBA, TMA where missing.
Add some tests.

[artem.tamazov]

Looks good.

llvm/trunk/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp

 #include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
 
 namespace {
 
 struct OptionalOperand;
 
+enum RegisterKind { IS_UNKNOWN, IS_VGPR, IS_SGPR, IS_TTMP, IS_SPECIAL };
+
 class AMDGPUOperand : public MCParsedAsmOperand {
   enum KindTy {
     Token,
     Immediate,
     Register,
     Expression
   } Kind;
 
private:
   bool ParseSectionDirectiveHSAText();
   bool subtargetHasRegister(const MCRegisterInfo &MRI, unsigned RegNo) const;
   bool ParseDirectiveAMDGPUHsaKernel();
   bool ParseDirectiveAMDGPUHsaModuleGlobal();
   bool ParseDirectiveAMDGPUHsaProgramGlobal();
   bool ParseSectionDirectiveHSADataGlobalAgent();
   bool ParseSectionDirectiveHSADataGlobalProgram();
   bool ParseSectionDirectiveHSARodataReadonlyAgent();
+  bool AddNextRegisterToList(unsigned& Reg, unsigned& RegWidth, RegisterKind RegKind, unsigned Reg1, unsigned RegNum);
+  bool ParseAMDGPURegister(RegisterKind& RegKind, unsigned& Reg, unsigned& RegNum, unsigned& RegWidth);
 
 public:
   enum AMDGPUMatchResultTy {
     Match_PreferE32 = FIRST_TARGET_MATCH_RESULT_TY
   };
 
   AMDGPUAsmParser(const MCSubtargetInfo &STI, MCAsmParser &_Parser,
                const MCInstrInfo &MII,
struct OptionalOperand {
   AMDGPUOperand::ImmTy Type;
   bool IsBit;
   int64_t Default;
   bool (*ConvertResult)(int64_t&);
 };
 
 }
 
-enum  RegisterKind { IS_VGPR, IS_SGPR, IS_TTMP };
-
 static int getRegClass(RegisterKind Is, unsigned RegWidth) {
   if (Is == IS_VGPR) {
     switch (RegWidth) {
       default: return -1;
       case 1: return AMDGPU::VGPR_32RegClassID;
       case 2: return AMDGPU::VReg_64RegClassID;
       case 3: return AMDGPU::VReg_96RegClassID;
       case 4: return AMDGPU::VReg_128RegClassID;
switch (RegWidth) {
       case 4: return AMDGPU::SReg_128RegClassID;
       case 8: return AMDGPU::SReg_256RegClassID;
       case 16: return AMDGPU::SReg_512RegClassID;
     }
   }
   return -1;
 }
 
-static unsigned getRegForName(StringRef RegName) {
-
+static unsigned getSpecialRegForName(StringRef RegName) {
   return StringSwitch<unsigned>(RegName)
     .Case("exec", AMDGPU::EXEC)
     .Case("vcc", AMDGPU::VCC)
     .Case("flat_scratch", AMDGPU::FLAT_SCR)
     .Case("m0", AMDGPU::M0)
     .Case("scc", AMDGPU::SCC)
+    .Case("tba", AMDGPU::TBA)
+    .Case("tma", AMDGPU::TMA)
     .Case("flat_scratch_lo", AMDGPU::FLAT_SCR_LO)
     .Case("flat_scratch_hi", AMDGPU::FLAT_SCR_HI)
     .Case("vcc_lo", AMDGPU::VCC_LO)
     .Case("vcc_hi", AMDGPU::VCC_HI)
     .Case("exec_lo", AMDGPU::EXEC_LO)
     .Case("exec_hi", AMDGPU::EXEC_HI)
     .Case("tma_lo", AMDGPU::TMA_LO)
     .Case("tma_hi", AMDGPU::TMA_HI)
     .Case("tba_lo", AMDGPU::TBA_LO)
     .Case("tba_hi", AMDGPU::TBA_HI)
     .Default(0);
 }
 
 bool AMDGPUAsmParser::ParseRegister(unsigned &RegNo, SMLoc &StartLoc, SMLoc &EndLoc) {
   auto R = parseRegister();
   if (!R) return true;
   assert(R->isReg());
   RegNo = R->getReg();
   StartLoc = R->getStartLoc();
   EndLoc = R->getEndLoc();
   return false;
 }
 
-std::unique_ptr<AMDGPUOperand> AMDGPUAsmParser::parseRegister() {
-  const AsmToken &Tok = Parser.getTok();
-  SMLoc StartLoc = Tok.getLoc();
-  SMLoc EndLoc = Tok.getEndLoc();
-  const MCRegisterInfo *TRI = getContext().getRegisterInfo();
-
-  StringRef RegName = Tok.getString();
-  unsigned RegNo = getRegForName(RegName);
-
-  if (RegNo) {
-    Parser.Lex();
-    if (!subtargetHasRegister(*TRI, RegNo))
-      return nullptr;
-    return AMDGPUOperand::CreateReg(RegNo, StartLoc, EndLoc,
-                                    TRI, &getSTI(), false);
-  }
+bool AMDGPUAsmParser::AddNextRegisterToList(unsigned& Reg, unsigned& RegWidth, RegisterKind RegKind, unsigned Reg1, unsigned RegNum)
+{
+  switch (RegKind) {
+  case IS_SPECIAL:
+    if (Reg == AMDGPU::EXEC_LO && Reg1 == AMDGPU::EXEC_HI) { Reg = AMDGPU::EXEC; RegWidth = 2; return true; }
+    if (Reg == AMDGPU::FLAT_SCR_LO && Reg1 == AMDGPU::FLAT_SCR_HI) { Reg = AMDGPU::FLAT_SCR; RegWidth = 2; return true; }
+    if (Reg == AMDGPU::VCC_LO && Reg1 == AMDGPU::VCC_HI) { Reg = AMDGPU::VCC; RegWidth = 2; return true; }
+    if (Reg == AMDGPU::TBA_LO && Reg1 == AMDGPU::TBA_HI) { Reg = AMDGPU::TBA; RegWidth = 2; return true; }
+    if (Reg == AMDGPU::TMA_LO && Reg1 == AMDGPU::TMA_HI) { Reg = AMDGPU::TMA; RegWidth = 2; return true; }
+    return false;
+  case IS_VGPR:
+  case IS_SGPR:
+  case IS_TTMP:
+    if (Reg1 != Reg + RegWidth) { return false; }
+    RegWidth++;
+    return true;
+  default:
+    assert(false); return false;
+  }
 
-  // Match vgprs, sgprs and ttmps
-  if (RegName[0] != 's' && RegName[0] != 'v' && !RegName.startswith("ttmp"))
-    return nullptr;
-
-  const RegisterKind Is = RegName[0] == 'v' ? IS_VGPR : RegName[0] == 's' ? IS_SGPR : IS_TTMP;
-  unsigned RegWidth;
-  unsigned RegIndexInClass;
-  if (RegName.size() > (Is == IS_TTMP ? strlen("ttmp") : 1) ) {
-    // We have a single 32-bit register. Syntax: vXX
-    RegWidth = 1;
-    if (RegName.substr(Is == IS_TTMP ? strlen("ttmp") : 1).getAsInteger(10, RegIndexInClass))
-      return nullptr;
+bool AMDGPUAsmParser::ParseAMDGPURegister(RegisterKind& RegKind, unsigned& Reg, unsigned& RegNum, unsigned& RegWidth)
+{
+  const MCRegisterInfo *TRI = getContext().getRegisterInfo();
+  if (getLexer().is(AsmToken::Identifier)) {
+    StringRef RegName = Parser.getTok().getString();
+    if ((Reg = getSpecialRegForName(RegName))) {
       Parser.Lex();
+      RegKind = IS_SPECIAL;
     } else {
-    // We have a register greater than 32-bits (a range of single registers). Syntax: v[XX:YY]
-
-    int64_t RegLo, RegHi;
+      unsigned RegNumIndex = 0;
+      if (RegName[0] == 'v') { RegNumIndex = 1; RegKind = IS_VGPR; }
+      else if (RegName[0] == 's') { RegNumIndex = 1; RegKind = IS_SGPR; }
+      else if (RegName.startswith("ttmp")) { RegNumIndex = strlen("ttmp"); RegKind = IS_TTMP; }
+      else { return false; }
+      if (RegName.size() > RegNumIndex) {
+        // Single 32-bit register: vXX.
+        if (RegName.substr(RegNumIndex).getAsInteger(10, RegNum)) { return false; }
         Parser.Lex();
-    if (getLexer().isNot(AsmToken::LBrac))
-      return nullptr;
-
+        RegWidth = 1;
+      } else {
+        // Range of registers: v[XX:YY].
+        Parser.Lex();
+        int64_t RegLo, RegHi;
+        if (getLexer().isNot(AsmToken::LBrac)) { return false; }
         Parser.Lex();
-    if (getParser().parseAbsoluteExpression(RegLo))
-      return nullptr;
 
-    if (getLexer().isNot(AsmToken::Colon))
-      return nullptr;
+        if (getParser().parseAbsoluteExpression(RegLo)) { return false; }
 
+        if (getLexer().isNot(AsmToken::Colon)) { return false; }
         Parser.Lex();
-    if (getParser().parseAbsoluteExpression(RegHi))
-      return nullptr;
 
-    if (getLexer().isNot(AsmToken::RBrac))
-      return nullptr;
+        if (getParser().parseAbsoluteExpression(RegHi)) { return false; }
 
+        if (getLexer().isNot(AsmToken::RBrac)) { return false; }
         Parser.Lex();
+
+        RegNum = (unsigned) RegLo;
         RegWidth = (RegHi - RegLo) + 1;
-    if (Is == IS_VGPR) {
-      // VGPR registers aren't aligned.
-      RegIndexInClass = RegLo;
+      }
+    }
+  } else if (getLexer().is(AsmToken::LBrac)) {
+    // List of consecutive registers: [s0,s1,s2,s3]
+    Parser.Lex();
+    if (!ParseAMDGPURegister(RegKind, Reg, RegNum, RegWidth)) { return false; }
+    if (RegWidth != 1) { return false; }
+    RegisterKind RegKind1;
+    unsigned Reg1, RegNum1, RegWidth1;
+    do {
+      if (getLexer().is(AsmToken::Comma)) {
+        Parser.Lex();
+      } else if (getLexer().is(AsmToken::RBrac)) {
+        Parser.Lex();
+        break;
+      } else if (ParseAMDGPURegister(RegKind1, Reg1, RegNum1, RegWidth1)) {
+        if (RegWidth1 != 1) { return false; }
+        if (RegKind1 != RegKind) { return false; }
+        if (!AddNextRegisterToList(Reg, RegWidth, RegKind1, Reg1, RegNum1)) { return false; }
       } else {
+        return false;
+      }
+    } while (true);
+  } else {
+    return false;
+  }
+  switch (RegKind) {
+  case IS_SPECIAL:
+    RegNum = 0;
+    RegWidth = 1;
+    break;
+  case IS_VGPR:
+  case IS_SGPR:
+  case IS_TTMP:
+  {
+    unsigned Size = 1;
+    if (RegKind == IS_SGPR || RegKind == IS_TTMP) {
       // SGPR and TTMP registers must be are aligned. Max required alignment is 4 dwords.
-      unsigned Size = std::min(RegWidth, 4u);
-      if (RegLo % Size != 0)
-        return nullptr;
+      Size = std::min(RegWidth, 4u);
+    }
+    if (RegNum % Size != 0) { return false; }
+    RegNum = RegNum / Size;
+    int RCID = getRegClass(RegKind, RegWidth);
+    if (RCID == -1) { return false; }
+    const MCRegisterClass RC = TRI->getRegClass(RCID);
+    if (RegNum >= RC.getNumRegs()) { return false; }
+    Reg = RC.getRegister(RegNum);
+    break;
+  }
 
-      RegIndexInClass = RegLo / Size;
-    }
-  }
+  default:
+    assert(false); return false;
+
+  if (!subtargetHasRegister(*TRI, Reg)) { return false; }
+  return true;
 
-  int RCID = getRegClass(Is, RegWidth);
-  if (RCID == -1)
-    return nullptr;
+std::unique_ptr<AMDGPUOperand> AMDGPUAsmParser::parseRegister() {
+  const auto &Tok = Parser.getTok();
+  SMLoc StartLoc = Tok.getLoc();
+  SMLoc EndLoc = Tok.getEndLoc();
+  const MCRegisterInfo *TRI = getContext().getRegisterInfo();
 
-  const MCRegisterClass RC = TRI->getRegClass(RCID);
-  if (RegIndexInClass >= RC.getNumRegs())
-    return nullptr;
+  RegisterKind RegKind;
+  unsigned Reg, RegNum, RegWidth;
 
-  RegNo = RC.getRegister(RegIndexInClass);
-  if (!subtargetHasRegister(*TRI, RegNo))
+  if (!ParseAMDGPURegister(RegKind, Reg, RegNum, RegWidth)) {
     return nullptr;
-
-  return AMDGPUOperand::CreateReg(RegNo, StartLoc, EndLoc,
+  }
+  return AMDGPUOperand::CreateReg(Reg, StartLoc, EndLoc,
                                   TRI, &getSTI(), false);
 }
 
 unsigned AMDGPUAsmParser::checkTargetMatchPredicate(MCInst &Inst) {
 
   uint64_t TSFlags = MII.get(Inst.getOpcode()).TSFlags;
 
   if ((getForcedEncodingSize() == 32 && (TSFlags & SIInstrFlags::VOP3)) ||
case AsmToken::Real: {
 
       APFloat F((float)BitsToDouble(IntVal));
       if (Negate)
         F.changeSign();
       Operands.push_back(
           AMDGPUOperand::CreateImm(F.bitcastToAPInt().getZExtValue(), S));
       return MatchOperand_Success;
     }
+    case AsmToken::LBrac:
     case AsmToken::Identifier: {
       if (auto R = parseRegister()) {
         unsigned Modifiers = 0;
 
         if (Negate)
           Modifiers |= 0x1;
 
         if (Abs) {

llvm/trunk/lib/Target/AMDGPU/SIRegisterInfo.td

 
 def SGPR_64 : RegisterClass<"AMDGPU", [v2i32, i64, f64], 32, (add SGPR_64Regs)>;
 
 def TTMP_64 : RegisterClass<"AMDGPU", [v2i32, i64, f64], 32, (add TTMP_64Regs)> {
   let isAllocatable = 0;
 }
 
 def SReg_64 : RegisterClass<"AMDGPU", [v2i32, i64, f64, i1], 32,
-  (add SGPR_64, VCC, EXEC, FLAT_SCR, TTMP_64)
+  (add SGPR_64, VCC, EXEC, FLAT_SCR, TTMP_64, TBA, TMA)
 >;
 
 def SReg_128 : RegisterClass<"AMDGPU", [v4i32, v16i8, v2i64], 32, (add SGPR_128)> {
   // Requires 2 s_mov_b64 to copy
   let CopyCost = 2;
 }
 
 def SReg_256 : RegisterClass<"AMDGPU", [v8i32, v8f32], 32, (add SGPR_256)> {

llvm/trunk/test/MC/AMDGPU/trap.s

 
 //===----------------------------------------------------------------------===//
 // Trap Handler related - Pairs and quadruples of registers
 //===----------------------------------------------------------------------===//
 
 s_mov_b64     ttmp[4:5], exec
 // SICI: s_mov_b64 ttmp[4:5], exec       ; encoding: [0x7e,0x04,0xf4,0xbe]
 // VI:   s_mov_b64 ttmp[4:5], exec       ; encoding: [0x7e,0x01,0xf4,0xbe]
+
+s_mov_b64     tba, ttmp[4:5]
+// SICI: s_mov_b64 tba, ttmp[4:5]        ; encoding: [0x74,0x04,0xec,0xbe]
+// VI:   s_mov_b64 tba, ttmp[4:5]        ; encoding: [0x74,0x01,0xec,0xbe]
+
+s_mov_b64     ttmp[4:5], tba
+// SICI: s_mov_b64 ttmp[4:5], tba        ; encoding: [0x6c,0x04,0xf4,0xbe]
+// VI:   s_mov_b64 ttmp[4:5], tba        ; encoding: [0x6c,0x01,0xf4,0xbe]
+
+s_mov_b64     tma, ttmp[4:5]
+// SICI: s_mov_b64 tma, ttmp[4:5]        ; encoding: [0x74,0x04,0xee,0xbe]
+// VI:   s_mov_b64 tma, ttmp[4:5]        ; encoding: [0x74,0x01,0xee,0xbe]
+
+s_mov_b64     ttmp[4:5], tma
+// SICI: s_mov_b64 ttmp[4:5], tma        ; encoding: [0x6e,0x04,0xf4,0xbe]
+// VI:   s_mov_b64 ttmp[4:5], tma        ; encoding: [0x6e,0x01,0xf4,0xbe]