[AMDGPU] Require explicit immediate offsets for SGPR+IMM SMEM instructions. (#79131)

As otherwise SGPR+IMM instructions are not distinguishable to SGPR-only
ones in AsmParser, leading to ambiguities.

GFX12 doesn't have special SGPR-only variants, so we still allow
optional immediate offsets for the subtarget.

Also rename the offset operand classes while there.

Part of <https://github.com/llvm/llvm-project/issues/69256>.

2 files changed, 32 insertions, 23 deletions

diff --git a/llvm/lib/Target/AMDGPU/SIInstrInfo.td b/llvm/lib/Target/AMDGPU/SIInstrInfo.td
index 4ab840a4..903a047 100644
--- a/llvm/lib/Target/AMDGPU/SIInstrInfo.td
+++ b/llvm/lib/Target/AMDGPU/SIInstrInfo.td
@@ -994,9 +994,9 @@ def SDWAVopcDst : BoolRC {
   let PrintMethod = "printVOPDst";
 }
 
-class NamedIntOperand<ValueType Type, string Prefix, string Name = NAME,
-                      string ConvertMethod = "nullptr">
-    : CustomOperand<Type, 1, Name> {
+class NamedIntOperand<ValueType Type, string Prefix, bit Optional = 1,
+                      string Name = NAME, string ConvertMethod = "nullptr">
+    : CustomOperand<Type, Optional, Name> {
   let ParserMethod =
     "[this](OperandVector &Operands) -> ParseStatus { "#
     "return parseIntWithPrefix(\""#Prefix#"\", Operands, "#
@@ -1039,9 +1039,9 @@ class ArrayOperand0<string Id, string Name = NAME>
 
 let ImmTy = "ImmTyOffset" in
 def flat_offset : CustomOperand<i32, 1, "FlatOffset">;
-def offset : NamedIntOperand<i32, "offset", "Offset">;
-def offset0 : NamedIntOperand<i8, "offset0", "Offset0">;
-def offset1 : NamedIntOperand<i8, "offset1", "Offset1">;
+def offset : NamedIntOperand<i32, "offset", 1, "Offset">;
+def offset0 : NamedIntOperand<i8, "offset0", 1, "Offset0">;
+def offset1 : NamedIntOperand<i8, "offset1", 1, "Offset1">;
 
 def gds : NamedBitOperand<"gds", "GDS">;
 
@@ -1092,25 +1092,25 @@ def dpp8 : CustomOperand<i32, 0, "DPP8">;
 def dpp_ctrl : CustomOperand<i32, 0, "DPPCtrl">;
 
 let DefaultValue = "0xf" in {
-def row_mask : NamedIntOperand<i32, "row_mask", "DppRowMask">;
-def bank_mask : NamedIntOperand<i32, "bank_mask", "DppBankMask">;
+def row_mask : NamedIntOperand<i32, "row_mask", 1, "DppRowMask">;
+def bank_mask : NamedIntOperand<i32, "bank_mask", 1, "DppBankMask">;
 }
-def bound_ctrl : NamedIntOperand<i1, "bound_ctrl", "DppBoundCtrl",
+def bound_ctrl : NamedIntOperand<i1, "bound_ctrl", 1, "DppBoundCtrl",
     "[this] (int64_t &BC) -> bool { return convertDppBoundCtrl(BC); }">;
-def FI : NamedIntOperand<i32, "fi", "DppFI">;
+def FI : NamedIntOperand<i32, "fi", 1, "DppFI">;
 
 def blgp : CustomOperand<i32, 1, "BLGP">;
-def cbsz : NamedIntOperand<i32, "cbsz", "CBSZ">;
-def abid : NamedIntOperand<i32, "abid", "ABID">;
+def cbsz : NamedIntOperand<i32, "cbsz", 1, "CBSZ">;
+def abid : NamedIntOperand<i32, "abid", 1, "ABID">;
 
 def hwreg : CustomOperand<i32, 0, "Hwreg">;
 
 def exp_tgt : CustomOperand<i32, 0, "ExpTgt">;
 
-def wait_vdst : NamedIntOperand<i8, "wait_vdst", "WaitVDST">;
-def wait_exp : NamedIntOperand<i8, "wait_exp", "WaitEXP">;
-def wait_va_vdst : NamedIntOperand<i8, "wait_va_vdst", "WaitVAVDst">;
-def wait_va_vsrc : NamedIntOperand<i8, "wait_vm_vsrc", "WaitVMVSrc">;
+def wait_vdst : NamedIntOperand<i8, "wait_vdst", 1, "WaitVDST">;
+def wait_exp : NamedIntOperand<i8, "wait_exp", 1, "WaitEXP">;
+def wait_va_vdst : NamedIntOperand<i8, "wait_va_vdst", 1, "WaitVAVDst">;
+def wait_va_vsrc : NamedIntOperand<i8, "wait_vm_vsrc", 1, "WaitVMVSrc">;
 
 class KImmFPOperand<ValueType vt> : ImmOperand<vt> {
   let OperandNamespace = "AMDGPU";
diff --git a/llvm/lib/Target/AMDGPU/SMInstructions.td b/llvm/lib/Target/AMDGPU/SMInstructions.td
index 9a27d22..f082de3 100644
--- a/llvm/lib/Target/AMDGPU/SMInstructions.td
+++ b/llvm/lib/Target/AMDGPU/SMInstructions.td
@@ -10,8 +10,13 @@ def smrd_offset_8 : ImmOperand<i32, "SMRDOffset8", 1>;
 
 let EncoderMethod = "getSMEMOffsetEncoding",
     DecoderMethod = "decodeSMEMOffset" in {
-def smem_offset : ImmOperand<i32, "SMEMOffset", 1>;
-def smem_offset_mod : NamedIntOperand<i32, "offset", "SMEMOffsetMod">;
+def SMEMOffset : ImmOperand<i32, "SMEMOffset", 1>;
+def SMEMOffsetMod : NamedIntOperand<i32, "offset", 0>;
+def OptSMEMOffsetMod : NamedIntOperand<i32, "offset"> {
+  let ImmTy = SMEMOffsetMod.ImmTy;
+  let PredicateMethod = SMEMOffsetMod.PredicateMethod;
+  let PrintMethod = SMEMOffsetMod.PrintMethod;
+}
 }
 
 //===----------------------------------------------------------------------===//
@@ -87,11 +92,14 @@ class OffsetMode<bit hasOffset, bit hasSOffset, string variant,
   string Asm = asm;
 }
 
-def IMM_Offset : OffsetMode<1, 0, "_IMM", (ins smem_offset:$offset), "$offset">;
+def IMM_Offset : OffsetMode<1, 0, "_IMM", (ins SMEMOffset:$offset), "$offset">;
 def SGPR_Offset : OffsetMode<0, 1, "_SGPR", (ins SReg_32:$soffset), "$soffset">;
 def SGPR_IMM_Offset : OffsetMode<1, 1, "_SGPR_IMM",
-                                 (ins SReg_32:$soffset, smem_offset_mod:$offset),
+                                 (ins SReg_32:$soffset, SMEMOffsetMod:$offset),
                                  "$soffset$offset">;
+def SGPR_IMM_OptOffset : OffsetMode<1, 1, "_SGPR_IMM",
+                                    (ins SReg_32:$soffset, OptSMEMOffsetMod:$offset),
+                                    "$soffset$offset">;
 
 class SM_Probe_Pseudo <string opName, RegisterClass baseClass, OffsetMode offsets>
   : SM_Pseudo<opName, (outs),
@@ -201,6 +209,7 @@ multiclass SM_Pseudo_Probe<RegisterClass baseClass> {
   def _IMM  : SM_Probe_Pseudo <opName, baseClass, IMM_Offset>;
   def _SGPR : SM_Probe_Pseudo <opName, baseClass, SGPR_Offset>;
   def _SGPR_IMM : SM_Probe_Pseudo <opName, baseClass, SGPR_IMM_Offset>;
+  def _SGPR_OPT_IMM : SM_Probe_Pseudo <opName, baseClass, SGPR_IMM_OptOffset>;
 }
 
 class SM_WaveId_Pseudo<string opName, SDPatternOperator node> : SM_Pseudo<
@@ -214,7 +223,7 @@ class SM_WaveId_Pseudo<string opName, SDPatternOperator node> : SM_Pseudo<
 
 class SM_Prefetch_Pseudo <string opName, RegisterClass baseClass, bit hasSBase>
   : SM_Pseudo<opName, (outs), !con(!if(hasSBase, (ins baseClass:$sbase), (ins)),
-                                   (ins smem_offset:$offset, SReg_32:$soffset, i8imm:$sdata)),
+                                   (ins SMEMOffset:$offset, SReg_32:$soffset, i8imm:$sdata)),
               !if(hasSBase, " $sbase,", "") # " $offset, $soffset, $sdata"> {
   // Mark prefetches as both load and store to prevent reordering with loads
   // and stores. This is also needed for pattern to match prefetch intrinsic.
@@ -1410,7 +1419,7 @@ class SMEM_Real_Load_gfx12<bits<6> op, string ps, string opName, OffsetMode offs
 multiclass SM_Real_Loads_gfx12<bits<6> op, string ps = NAME> {
   defvar opName = !tolower(NAME);
   def _IMM_gfx12 : SMEM_Real_Load_gfx12<op, ps, opName, IMM_Offset>;
-  def _SGPR_IMM_gfx12 : SMEM_Real_Load_gfx12<op, ps, opName, SGPR_IMM_Offset>;
+  def _SGPR_IMM_gfx12 : SMEM_Real_Load_gfx12<op, ps, opName, SGPR_IMM_OptOffset>;
 }
 
 defm S_LOAD_B32  : SM_Real_Loads_gfx12<0x00, "S_LOAD_DWORD">;
@@ -1448,7 +1457,7 @@ def S_PREFETCH_DATA_PC_REL_gfx12 : SMEM_Real_Prefetch_gfx12<0x28, S_PREFETCH_DAT
 multiclass SMEM_Real_Probe_gfx12<bits<6> op> {
   defvar ps = NAME;
   def _IMM_gfx12      : SMEM_Real_Prefetch_gfx12<op, !cast<SM_Probe_Pseudo>(ps#_IMM)>;
-  def _SGPR_IMM_gfx12 : SMEM_Real_Prefetch_gfx12<op, !cast<SM_Probe_Pseudo>(ps#_SGPR_IMM)>;
+  def _SGPR_IMM_gfx12 : SMEM_Real_Prefetch_gfx12<op, !cast<SM_Probe_Pseudo>(ps#_SGPR_OPT_IMM)>;
 }
 
 defm S_ATC_PROBE        : SMEM_Real_Probe_gfx12<0x22>;