diff --git a/llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp b/llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp
index 859761bd81ffed..7a1b75c0c885d1 100644
--- a/llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp
+++ b/llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp
@@ -1578,6 +1578,7 @@ class AMDGPUAsmParser : public MCTargetAsmParser {
   OperandMatchResultTy parseFORMAT(OperandVector &Operands);
   OperandMatchResultTy parseSymbolicOrNumericFormat(int64_t &Format);
   OperandMatchResultTy parseNumericFormat(int64_t &Format);
+  OperandMatchResultTy parseFlatOffset(OperandVector &Operands);
   bool tryParseFmt(const char *Pref, int64_t MaxVal, int64_t &Val);
   bool matchDfmtNfmt(int64_t &Dfmt, int64_t &Nfmt, StringRef FormatStr, SMLoc Loc);
 
@@ -6359,6 +6360,16 @@ AMDGPUAsmParser::parseFORMAT(OperandVector &Operands) {
   return MatchOperand_Success;
 }
 
+OperandMatchResultTy AMDGPUAsmParser::parseFlatOffset(OperandVector &Operands) {
+  OperandMatchResultTy Res =
+      parseIntWithPrefix("offset", Operands, AMDGPUOperand::ImmTyOffset);
+  if (Res == MatchOperand_NoMatch) {
+    Res = parseIntWithPrefix("inst_offset", Operands,
+                             AMDGPUOperand::ImmTyInstOffset);
+  }
+  return Res;
+}
+
 //===----------------------------------------------------------------------===//
 // ds
 //===----------------------------------------------------------------------===//
@@ -9107,8 +9118,6 @@ AMDGPUAsmParser::parseCustomOperand(OperandVector &Operands, unsigned MCK) {
     return parseNamedBit("d16", Operands, AMDGPUOperand::ImmTyD16);
   case MCK_ImmDA:
     return parseNamedBit("da", Operands, AMDGPUOperand::ImmTyDA);
-  case MCK_ImmDMask:
-    return parseIntWithPrefix("dmask", Operands, AMDGPUOperand::ImmTyDMask);
   case MCK_ImmExpCompr:
     return parseNamedBit("compr", Operands, AMDGPUOperand::ImmTyExpCompr);
   case MCK_ImmExpVM:
@@ -9130,22 +9139,8 @@ AMDGPUAsmParser::parseCustomOperand(OperandVector &Operands, unsigned MCK) {
   case MCK_ImmNegLo:
     return parseOperandArrayWithPrefix("neg_lo", Operands,
                                        AMDGPUOperand::ImmTyNegLo);
-  case MCK_ImmOffset:
   case MCK_ImmSMEMOffset:
     return parseIntWithPrefix("offset", Operands, AMDGPUOperand::ImmTyOffset);
-  case MCK_ImmFlatOffset: {
-    OperandMatchResultTy Res =
-        parseIntWithPrefix("offset", Operands, AMDGPUOperand::ImmTyOffset);
-    if (Res == MatchOperand_NoMatch) {
-      Res = parseIntWithPrefix("inst_offset", Operands,
-                               AMDGPUOperand::ImmTyInstOffset);
-    }
-    return Res;
-  }
-  case MCK_ImmOffset0:
-    return parseIntWithPrefix("offset0", Operands, AMDGPUOperand::ImmTyOffset0);
-  case MCK_ImmOffset1:
-    return parseIntWithPrefix("offset1", Operands, AMDGPUOperand::ImmTyOffset1);
   case MCK_ImmOModSI:
     return parseOModOperand(Operands);
   case MCK_ImmOpSel:
@@ -9179,12 +9174,6 @@ AMDGPUAsmParser::parseCustomOperand(OperandVector &Operands, unsigned MCK) {
     return parseNamedBit("tfe", Operands, AMDGPUOperand::ImmTyTFE);
   case MCK_ImmUNorm:
     return parseNamedBit("unorm", Operands, AMDGPUOperand::ImmTyUNorm);
-  case MCK_ImmWaitEXP:
-    return parseIntWithPrefix("wait_exp", Operands,
-                              AMDGPUOperand::ImmTyWaitEXP);
-  case MCK_ImmWaitVDST:
-    return parseIntWithPrefix("wait_vdst", Operands,
-                              AMDGPUOperand::ImmTyWaitVDST);
   }
   return tryCustomParseOperand(Operands, MCK);
 }
diff --git a/llvm/lib/Target/AMDGPU/SIInstrInfo.td b/llvm/lib/Target/AMDGPU/SIInstrInfo.td
index 8bc8f43990e713..cffedd49dad35c 100644
--- a/llvm/lib/Target/AMDGPU/SIInstrInfo.td
+++ b/llvm/lib/Target/AMDGPU/SIInstrInfo.td
@@ -1174,16 +1174,33 @@ class NamedOperandBit_1<string Name, AsmOperandClass MatchClass> :
   let ParserMatchClass = MatchClass;
 }
 
-class NamedOperandU8<string Name, AsmOperandClass MatchClass> : Operand<i8> {
-  let PrintMethod = "print"#Name;
-  let ParserMatchClass = MatchClass;
+class CustomOperandClass<string CName, bit Optional> : AsmOperandClass {
+  let Name = CName;
+  let PredicateMethod = "is"#CName;
+  let ParserMethod = "parse"#CName;
+  let RenderMethod = "addImmOperands";
+  let IsOptional = Optional;
+  let DefaultMethod = "default"#CName;
 }
 
-class NamedOperandU16<string Name, AsmOperandClass MatchClass> : Operand<i16> {
+class CustomOperand<ValueType Type, bit Optional = 0, string Name = NAME,
+    AsmOperandClass Class = CustomOperandClass<Name, Optional>>
+    : Operand<Type> {
   let PrintMethod = "print"#Name;
-  let ParserMatchClass = MatchClass;
+  let ParserMatchClass = Class;
+}
+
+class NamedIntOperandClass<string Prefix, string Name>
+    : CustomOperandClass<Name, 1> {
+  string ImmTy = "AMDGPUOperand::ImmTy"#Name;
+  let ParserMethod =
+    "[this](OperandVector &Operands) -> OperandMatchResultTy { "#
+    "return parseIntWithPrefix(\""#Prefix#"\", Operands, "#ImmTy#"); }";
 }
 
+class NamedIntOperand<ValueType Type, string Prefix, string Name = NAME>
+  : CustomOperand<Type, 1, Name, NamedIntOperandClass<Prefix, Name>>;
+
 class NamedOperandU32<string Name, AsmOperandClass MatchClass> : Operand<i32> {
   let PrintMethod = "print"#Name;
   let ParserMatchClass = MatchClass;
@@ -1209,10 +1226,10 @@ class NamedOperandU32Default1<string Name, AsmOperandClass MatchClass> :
 
 let OperandType = "OPERAND_IMMEDIATE" in {
 
-def flat_offset : NamedOperandU16<"FlatOffset", NamedMatchClass<"FlatOffset">>;
-def offset : NamedOperandU16<"Offset", NamedMatchClass<"Offset">>;
-def offset0 : NamedOperandU8<"Offset0", NamedMatchClass<"Offset0">>;
-def offset1 : NamedOperandU8<"Offset1", NamedMatchClass<"Offset1">>;
+def flat_offset : CustomOperand<i16, 1, "FlatOffset">;
+def offset : NamedIntOperand<i16, "offset", "Offset">;
+def offset0 : NamedIntOperand<i8, "offset0", "Offset0">;
+def offset1 : NamedIntOperand<i8, "offset1", "Offset1">;
 
 def gds : NamedOperandBit<"GDS", NamedMatchClass<"GDS">>;
 
@@ -1242,10 +1259,10 @@ def LWE : NamedOperandBit<"LWE", NamedMatchClass<"LWE">>;
 def exp_compr : NamedOperandBit<"ExpCompr", NamedMatchClass<"ExpCompr">>;
 def exp_vm : NamedOperandBit<"ExpVM", NamedMatchClass<"ExpVM">>;
 
-def FORMAT : NamedOperandU8<"FORMAT", NamedMatchClass<"FORMAT", 0>>;
+def FORMAT : CustomOperand<i8>;
 
-def DMask : NamedOperandU16<"DMask", NamedMatchClass<"DMask">>;
-def Dim : NamedOperandU8<"Dim", NamedMatchClass<"Dim", 0>>;
+def DMask : NamedIntOperand<i16, "dmask">;
+def Dim : CustomOperand<i8>;
 
 def dst_sel : NamedOperandU32<"SDWADstSel", NamedMatchClass<"SDWADstSel">>;
 def src0_sel : NamedOperandU32<"SDWASrc0Sel", NamedMatchClass<"SDWASrc0Sel">>;
@@ -1275,8 +1292,8 @@ def exp_tgt : NamedOperandU32<"ExpTgt", NamedMatchClass<"ExpTgt", 0>> {
 
 }
 
-def wait_vdst : NamedOperandU8<"WaitVDST", NamedMatchClass<"WaitVDST">>;
-def wait_exp : NamedOperandU8<"WaitEXP", NamedMatchClass<"WaitEXP">>;
+def wait_vdst : NamedIntOperand<i8, "wait_vdst", "WaitVDST">;
+def wait_exp : NamedIntOperand<i8, "wait_exp", "WaitEXP">;
 
 } // End OperandType = "OPERAND_IMMEDIATE"