[AMDGPU][MC] Added arg checks for vmcnt() expcnt() lgkmcnt() helpers

[dp]

Summary

Summary of changes:

• corrected vmcnt(), expcnt(), lgkmcnt() helpers to checks their argument for truncation;
• added saturated versions of these helpers.

See bug 32711 for details: https://bugs.llvm.org//show_bug.cgi?id=32711

[vpykhtin]

LGTM, Thanks!

lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp

void AMDGPUAsmParser::cvtExp(MCInst &Inst, const OperandVector &Operands) {
 
   Inst.addOperand(MCOperand::createImm(EnMask));
 }
 
 //===----------------------------------------------------------------------===//
 // s_waitcnt
 //===----------------------------------------------------------------------===//
 
+static bool
+encodeCnt(
+  const AMDGPU::IsaInfo::IsaVersion ISA,
+  int64_t &IntVal,
+  int64_t CntVal,
+  bool Saturate,
+  unsigned (*encode)(const IsaInfo::IsaVersion &Version, unsigned, unsigned),
+  unsigned (*decode)(const IsaInfo::IsaVersion &Version, unsigned))
+{
+  bool Failed = false;
+
+  IntVal = encode(ISA, IntVal, CntVal);
+  if (CntVal != decode(ISA, IntVal)) {
+    if (Saturate) {
+      IntVal = encode(ISA, IntVal, -1);
+    } else {
+      Failed = true;
+    }
+  }
+  return Failed;
+}
+
 bool AMDGPUAsmParser::parseCnt(int64_t &IntVal) {
   StringRef CntName = Parser.getTok().getString();
   int64_t CntVal;
 
   Parser.Lex();
   if (getLexer().isNot(AsmToken::LParen))
     return true;
 
   Parser.Lex();
   if (getLexer().isNot(AsmToken::Integer))
     return true;
 
   if (getParser().parseAbsoluteExpression(CntVal))
     return true;
 
-  if (getLexer().isNot(AsmToken::RParen))
-    return true;
-
-  Parser.Lex();
-  if (getLexer().is(AsmToken::Amp) || getLexer().is(AsmToken::Comma))
-    Parser.Lex();
-
   AMDGPU::IsaInfo::IsaVersion ISA =
       AMDGPU::IsaInfo::getIsaVersion(getFeatureBits());
-  if (CntName == "vmcnt")
-    IntVal = encodeVmcnt(ISA, IntVal, CntVal);
-  else if (CntName == "expcnt")
-    IntVal = encodeExpcnt(ISA, IntVal, CntVal);
-  else if (CntName == "lgkmcnt")
-    IntVal = encodeLgkmcnt(ISA, IntVal, CntVal);
-  else
+
+  bool Failed = true;
+  bool Sat = CntName.endswith("_sat");
+
+  if (CntName == "vmcnt" || CntName == "vmcnt_sat") {
+    Failed = encodeCnt(ISA, IntVal, CntVal, Sat, encodeVmcnt, decodeVmcnt);
+  } else if (CntName == "expcnt" || CntName == "expcnt_sat") {
+    Failed = encodeCnt(ISA, IntVal, CntVal, Sat, encodeExpcnt, decodeExpcnt);
+  } else if (CntName == "lgkmcnt" || CntName == "lgkmcnt_sat") {
+    Failed = encodeCnt(ISA, IntVal, CntVal, Sat, encodeLgkmcnt, decodeLgkmcnt);
+  }
+
+  // To improve diagnostics, do not skip delimiters on errors
+  if (!Failed) {
+    if (getLexer().isNot(AsmToken::RParen)) {
       return true;
+    }
+    Parser.Lex();
+    if (getLexer().is(AsmToken::Amp) || getLexer().is(AsmToken::Comma)) {
+      const AsmToken NextToken = getLexer().peekTok();
+      if (NextToken.is(AsmToken::Identifier)) {
+        Parser.Lex();
+      }
+    }
+  }
 
-  return false;
+  return Failed;
 }
 
 OperandMatchResultTy
 AMDGPUAsmParser::parseSWaitCntOps(OperandVector &Operands) {
   AMDGPU::IsaInfo::IsaVersion ISA =
       AMDGPU::IsaInfo::getIsaVersion(getFeatureBits());
   int64_t Waitcnt = getWaitcntBitMask(ISA);
   SMLoc S = Parser.getTok().getLoc();

test/MC/AMDGPU/sopp-err.s

 s_sendmsg sendmsg(MSG_SYSMSG, SYSMSG_OP_ECC_ERR_INTERRUPT, 0)
 // GCN: error: failed parsing operand
 
 s_sendmsg sendmsg(MSG_SYSMSG, 0)
 // GCN: error: invalid/unsupported code of SYSMSG_OP
 
 s_sendmsg sendmsg(MSG_SYSMSG, 5)
 // GCN: error: invalid/unsupported code of SYSMSG_OP
+
+s_waitcnt lgkmcnt(16)
+// GCN: error: failed parsing operand
+
+s_waitcnt expcnt(8)
+// GCN: error: failed parsing operand
+
+s_waitcnt vmcnt(16)
+// GCN: error: failed parsing operand
+
+s_waitcnt vmcnt(0xFFFFFFFFFFFF0000)
+// GCN: error: failed parsing operand
+
+s_waitcnt vmcnt(0), expcnt(0), lgkmcnt(0),
+// GCN: error: failed parsing operand
+
+s_waitcnt vmcnt(0) & expcnt(0) & lgkmcnt(0)&
+// GCN: error: failed parsing operand

test/MC/AMDGPU/sopp.s

 // GCN: s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0) ; encoding: [0x00,0x00,0x8c,0xbf]
 
 s_waitcnt vmcnt(1)
 // GCN: s_waitcnt vmcnt(1) ; encoding: [0x71,0x0f,0x8c,0xbf]
 
 s_waitcnt vmcnt(9)
 // GCN: s_waitcnt vmcnt(9) ; encoding: [0x79,0x0f,0x8c,0xbf]
 
+s_waitcnt vmcnt(15)
+// GCN: s_waitcnt ; encoding: [0x7f,0x0f,0x8c,0xbf]
+
+s_waitcnt vmcnt_sat(9)
+// GCN: s_waitcnt vmcnt(9) ; encoding: [0x79,0x0f,0x8c,0xbf]
+
+s_waitcnt vmcnt_sat(15)
+// GCN: s_waitcnt ; encoding: [0x7f,0x0f,0x8c,0xbf]
+
+s_waitcnt vmcnt_sat(16)
+// GCN: s_waitcnt ; encoding: [0x7f,0x0f,0x8c,0xbf]
+
 s_waitcnt expcnt(2)
 // GCN: s_waitcnt expcnt(2) ; encoding: [0x2f,0x0f,0x8c,0xbf]
 
+s_waitcnt expcnt(7)
+// GCN: s_waitcnt ; encoding: [0x7f,0x0f,0x8c,0xbf]
+
+s_waitcnt expcnt_sat(2)
+// GCN: s_waitcnt expcnt(2) ; encoding: [0x2f,0x0f,0x8c,0xbf]
+
+s_waitcnt expcnt_sat(7)
+// GCN: s_waitcnt ; encoding: [0x7f,0x0f,0x8c,0xbf]
+
+s_waitcnt expcnt_sat(0xFFFF0000)
+// GCN: s_waitcnt ; encoding: [0x7f,0x0f,0x8c,0xbf]
+
 s_waitcnt lgkmcnt(3)
 // GCN: s_waitcnt lgkmcnt(3) ; encoding: [0x7f,0x03,0x8c,0xbf]
 
 s_waitcnt lgkmcnt(9)
 // GCN: s_waitcnt lgkmcnt(9) ; encoding: [0x7f,0x09,0x8c,0xbf]
 
+s_waitcnt lgkmcnt(15)
+// GCN: s_waitcnt ; encoding: [0x7f,0x0f,0x8c,0xbf]
+
 s_waitcnt vmcnt(0), expcnt(0)
 // GCN: s_waitcnt vmcnt(0) expcnt(0) ; encoding: [0x00,0x0f,0x8c,0xbf]
 
+s_waitcnt lgkmcnt_sat(3)
+// GCN: s_waitcnt lgkmcnt(3) ; encoding: [0x7f,0x03,0x8c,0xbf]
+
+s_waitcnt lgkmcnt_sat(9)
+// GCN: s_waitcnt lgkmcnt(9) ; encoding: [0x7f,0x09,0x8c,0xbf]
+
+s_waitcnt lgkmcnt_sat(15)
+// GCN: s_waitcnt ; encoding: [0x7f,0x0f,0x8c,0xbf]
+
+s_waitcnt lgkmcnt_sat(16)
+// GCN: s_waitcnt ; encoding: [0x7f,0x0f,0x8c,0xbf]
 
 s_sethalt 9
 // GCN: s_sethalt 9 ; encoding: [0x09,0x00,0x8d,0xbf]
 
 s_setkill 7
 // GCN: s_setkill 7 ; encoding: [0x07,0x00,0x8b,0xbf]
 
 s_sleep 10