[AMDGPU][GFX9][GFX10] Support base+soffset+offset SMEM atomics.

[kosarev]

Summary

Resolves a part of
https://github.com/llvm/llvm-project/issues/38652

[dp]

LGTM.

llvm/lib/Target/AMDGPU/SMInstructions.td

90

Nice idea! Could other classes for loads, stores, etc. benefit from using OffsetMode?

[kosarev]

llvm/lib/Target/AMDGPU/SMInstructions.td

90

That's the plan, yes, but one step at a time.

llvm/lib/Target/AMDGPU/SMInstructions.td

class SM_Real <SM_Pseudo ps, string opName = ps.Mnemonic>
   // encoding
   bits<7>  sbase;
   bits<7>  sdst;
   bits<32> offset;
   bits<8>  soffset;
   bits<5> cpol;
 }
 
+class OffsetMode<bit hasOffset, bit hasSOffset, string variant,
+                 dag ins, string asm> {
+  bit HasOffset = hasOffset;
+  bit HasSOffset = hasSOffset;
+  string Variant = variant;
+  dag Ins = ins;
+  string Asm = asm;
+}
+
+def IMM_Offset : OffsetMode<1, 0, "_IMM", (ins smem_offset:$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),
+                                 "$soffset$offset">;
+
 class SM_Probe_Pseudo <string opName, string variant, RegisterClass baseClass,
                        dag offsets, string asmOffsets,
                        bit hasOffset, bit hasSOffset>
   : SM_Pseudo<opName, (outs),
               !con((ins i8imm:$sdata, baseClass:$sbase), offsets),
               " $sdata, $sbase, " # asmOffsets> {
   let mayLoad = 0;
   let mayStore = 0;
class SM_Atomic_Pseudo <string opName,
   let IsAtomicRet = isRet;
 
   let AsmMatchConverter = "cvtSMEMAtomic";
 }
 
 class SM_Pseudo_Atomic<string opName,
                        RegisterClass baseClass,
                        RegisterClass dataClass,
-                       bit isImm,
+                       OffsetMode offsets,
                        bit isRet,
-                       string opNameWithSuffix = opName # !if(isImm,
-                                 !if(isRet, "_IMM_RTN", "_IMM"),
-                                 !if(isRet, "_SGPR_RTN", "_SGPR")),
+                       string opNameWithSuffix =
+                         opName # offsets.Variant # !if(isRet, "_RTN", ""),
                        Operand CPolTy = !if(isRet, CPol_GLC1, CPol)> :
   SM_Atomic_Pseudo<opName,
                    !if(isRet, (outs dataClass:$sdst), (outs)),
-                   !if(isImm,
-                       (ins dataClass:$sdata, baseClass:$sbase, smem_offset:$offset, CPolTy:$cpol),
-                       (ins dataClass:$sdata, baseClass:$sbase, SReg_32:$soffset, CPolTy:$cpol)),
-                   !if(isRet, " $sdst", " $sdata") # ", $sbase, " #
-                     !if(isImm, "$offset", "$soffset") # "$cpol",
+                   !con((ins dataClass:$sdata, baseClass:$sbase), offsets.Ins,
+                        (ins CPolTy:$cpol)),
+                   !if(isRet, " $sdst", " $sdata") #
+                     ", $sbase, " # offsets.Asm # "$cpol",
                    isRet>,
   AtomicNoRet <opNameWithSuffix, isRet> {
-  let has_offset = isImm;
-  let has_soffset = !not(isImm);
+  let has_offset = offsets.HasOffset;
+  let has_soffset = offsets.HasSOffset;
   let PseudoInstr = opNameWithSuffix;
 
   let Constraints = !if(isRet, "$sdst = $sdata", "");
   let DisableEncoding = !if(isRet, "$sdata", "");
 }
 
 multiclass SM_Pseudo_Atomics<string opName,
                              RegisterClass baseClass,
                              RegisterClass dataClass> {
-  def _IMM      : SM_Pseudo_Atomic <opName, baseClass, dataClass, 1, 0>;
-  def _SGPR     : SM_Pseudo_Atomic <opName, baseClass, dataClass, 0, 0>;
-  def _IMM_RTN  : SM_Pseudo_Atomic <opName, baseClass, dataClass, 1, 1>;
-  def _SGPR_RTN : SM_Pseudo_Atomic <opName, baseClass, dataClass, 0, 1>;
+  def _IMM      : SM_Pseudo_Atomic <opName, baseClass, dataClass, IMM_Offset, 0>;
+  def _SGPR     : SM_Pseudo_Atomic <opName, baseClass, dataClass, SGPR_Offset, 0>;
+  def _SGPR_IMM : SM_Pseudo_Atomic <opName, baseClass, dataClass, SGPR_IMM_Offset, 0>;
+  def _IMM_RTN  : SM_Pseudo_Atomic <opName, baseClass, dataClass, IMM_Offset, 1>;
+  def _SGPR_RTN : SM_Pseudo_Atomic <opName, baseClass, dataClass, SGPR_Offset, 1>;
+  def _SGPR_IMM_RTN : SM_Pseudo_Atomic <opName, baseClass, dataClass, SGPR_IMM_Offset, 1>;
 }
 
 //===----------------------------------------------------------------------===//
 // Scalar Memory Instructions
 //===----------------------------------------------------------------------===//
 
 // We are using the SReg_32_XM0 and not the SReg_32 register class for 32-bit
 // SMRD instructions, because the SReg_32_XM0 register class does not include M0
class SMEM_Atomic_Real_vi <bits<8> op, SM_Atomic_Pseudo ps>
 
   let cpol{CPolBit.GLC} = ps.glc;
   let Inst{12-6} = !if(ps.glc, sdst{6-0}, sdata{6-0});
 }
 
 multiclass SM_Real_Atomics_vi<bits<8> op, string ps> {
   def _IMM_vi       : SMEM_Atomic_Real_vi <op, !cast<SM_Atomic_Pseudo>(ps#_IMM)>;
   def _SGPR_vi      : SMEM_Atomic_Real_vi <op, !cast<SM_Atomic_Pseudo>(ps#_SGPR)>;
+  def _SGPR_alt_gfx9
+    : SMEM_Atomic_Real_vi <op, !cast<SM_Atomic_Pseudo>(ps#_SGPR)>,
+      SMEM_Real_SGPR_alt_gfx9;
+  let IsGFX9SpecificEncoding = true in
+  def _SGPR_IMM_gfx9
+    : SMEM_Atomic_Real_vi <op, !cast<SM_Atomic_Pseudo>(ps#_SGPR_IMM)>;
   def _IMM_RTN_vi   : SMEM_Atomic_Real_vi <op, !cast<SM_Atomic_Pseudo>(ps#_IMM_RTN)>;
   def _SGPR_RTN_vi  : SMEM_Atomic_Real_vi <op, !cast<SM_Atomic_Pseudo>(ps#_SGPR_RTN)>;
+  def _SGPR_RTN_alt_gfx9
+    : SMEM_Atomic_Real_vi <op, !cast<SM_Atomic_Pseudo>(ps#_SGPR_RTN)>,
+      SMEM_Real_SGPR_alt_gfx9;
+  let IsGFX9SpecificEncoding = true in
+  def _SGPR_IMM_RTN_gfx9
+    : SMEM_Atomic_Real_vi <op, !cast<SM_Atomic_Pseudo>(ps#_SGPR_IMM_RTN)>;
 }
 
 defm S_BUFFER_ATOMIC_SWAP         : SM_Real_Atomics_vi <0x40, "S_BUFFER_ATOMIC_SWAP">;
 defm S_BUFFER_ATOMIC_CMPSWAP      : SM_Real_Atomics_vi <0x41, "S_BUFFER_ATOMIC_CMPSWAP">;
 defm S_BUFFER_ATOMIC_ADD          : SM_Real_Atomics_vi <0x42, "S_BUFFER_ATOMIC_ADD">;
 defm S_BUFFER_ATOMIC_SUB          : SM_Real_Atomics_vi <0x43, "S_BUFFER_ATOMIC_SUB">;
 defm S_BUFFER_ATOMIC_SMIN         : SM_Real_Atomics_vi <0x44, "S_BUFFER_ATOMIC_SMIN">;
 defm S_BUFFER_ATOMIC_UMIN         : SM_Real_Atomics_vi <0x45, "S_BUFFER_ATOMIC_UMIN">;
class SMEM_Atomic_Real_gfx10 <bits<8> op, SM_Atomic_Pseudo ps>
 
   let Inst{14} = !if(ps.has_dlc, cpol{CPolBit.DLC}, 0);
   let Inst{12-6} = !if(ps.glc, sdst{6-0}, sdata{6-0});
 }
 
 multiclass SM_Real_Atomics_gfx10<bits<8> op, string ps> {
   def _IMM_gfx10       : SMEM_Atomic_Real_gfx10 <op, !cast<SM_Atomic_Pseudo>(ps#_IMM)>;
   def _SGPR_gfx10      : SMEM_Atomic_Real_gfx10 <op, !cast<SM_Atomic_Pseudo>(ps#_SGPR)>;
+  def _SGPR_IMM_gfx10  : SMEM_Atomic_Real_gfx10 <op, !cast<SM_Atomic_Pseudo>(ps#_SGPR_IMM)>;
   def _IMM_RTN_gfx10   : SMEM_Atomic_Real_gfx10 <op, !cast<SM_Atomic_Pseudo>(ps#_IMM_RTN)>;
   def _SGPR_RTN_gfx10  : SMEM_Atomic_Real_gfx10 <op, !cast<SM_Atomic_Pseudo>(ps#_SGPR_RTN)>;
+  def _SGPR_IMM_RTN_gfx10 : SMEM_Atomic_Real_gfx10 <op, !cast<SM_Atomic_Pseudo>(ps#_SGPR_IMM_RTN)>;
 }
 
 let SubtargetPredicate = HasScalarAtomics in {
 
 defm S_BUFFER_ATOMIC_SWAP         : SM_Real_Atomics_gfx10 <0x40, "S_BUFFER_ATOMIC_SWAP">;
 defm S_BUFFER_ATOMIC_CMPSWAP      : SM_Real_Atomics_gfx10 <0x41, "S_BUFFER_ATOMIC_CMPSWAP">;
 defm S_BUFFER_ATOMIC_ADD          : SM_Real_Atomics_gfx10 <0x42, "S_BUFFER_ATOMIC_ADD">;
 defm S_BUFFER_ATOMIC_SUB          : SM_Real_Atomics_gfx10 <0x43, "S_BUFFER_ATOMIC_SUB">;

llvm/test/MC/AMDGPU/gfx10_asm_smem.s

 // GFX10: s_dcache_discard_x2 s[2:3], s7 offset:0x64 ; encoding: [0x01,0x00,0xa4,0xf4,0x64,0x00,0x00,0x0e]
 
 s_atomic_add s5, s[2:3], s101
 // GFX10: encoding: [0x41,0x01,0x08,0xf6,0x00,0x00,0x00,0xca]
 
 s_atomic_add s5, s[2:3], 0x64
 // GFX10: encoding: [0x41,0x01,0x08,0xf6,0x64,0x00,0x00,0xfa]
 
+s_atomic_add s5, s[2:3], s7 offset:0x64
+// GFX10: encoding: [0x41,0x01,0x08,0xf6,0x64,0x00,0x00,0x0e]
+
 s_atomic_add_x2 s[10:11], s[2:3], s101
 // GFX10: encoding: [0x81,0x02,0x88,0xf6,0x00,0x00,0x00,0xca]
 
 s_atomic_and s5, s[2:3], s101
 // GFX10: encoding: [0x41,0x01,0x20,0xf6,0x00,0x00,0x00,0xca]
 
 s_atomic_and_x2 s[10:11], s[2:3], 0x64
 // GFX10: encoding: [0x81,0x02,0xa0,0xf6,0x64,0x00,0x00,0xfa]

 // GFX10: encoding: [0x42,0x01,0x28,0xf5,0x00,0x00,0x00,0x04]
 
 s_atomic_add s5, s[2:3], s101 glc
 // GFX10: encoding: [0x41,0x01,0x09,0xf6,0x00,0x00,0x00,0xca]
 
 s_atomic_add s5, s[2:3], 0x64 glc
 // GFX10: encoding: [0x41,0x01,0x09,0xf6,0x64,0x00,0x00,0xfa]
 
+s_atomic_add s5, s[2:3], s7 offset:0x64 glc
+// GFX10: encoding: [0x41,0x01,0x09,0xf6,0x64,0x00,0x00,0x0e]
+
 s_atomic_add_x2 s[10:11], s[2:3], s101 glc
 // GFX10: encoding: [0x81,0x02,0x89,0xf6,0x00,0x00,0x00,0xca]
 
 s_atomic_and s5, s[2:3], s101 glc
 // GFX10: encoding: [0x41,0x01,0x21,0xf6,0x00,0x00,0x00,0xca]
 
 s_atomic_and_x2 s[10:11], s[2:3], 0x64 glc
 // GFX10: encoding: [0x81,0x02,0xa1,0xf6,0x64,0x00,0x00,0xfa]

llvm/test/MC/AMDGPU/gfx9_asm_smem.s

 // CHECK: [0x41,0x01,0x08,0xc2,0x7b,0x00,0x00,0x00]
 
 s_atomic_add s5, s[2:3], m0
 // CHECK: [0x41,0x01,0x08,0xc2,0x7c,0x00,0x00,0x00]
 
 s_atomic_add s5, s[2:3], 0x0
 // CHECK: [0x41,0x01,0x0a,0xc2,0x00,0x00,0x00,0x00]
 
+s_atomic_add s5, s[2:3], s7 offset:0x12345
+// CHECK: [0x41,0x41,0x0a,0xc2,0x45,0x23,0x01,0x0e]
+
 s_atomic_add s5, s[2:3], s0 glc
 // CHECK: [0x41,0x01,0x09,0xc2,0x00,0x00,0x00,0x00]
 
+s_atomic_add s5, s[2:3], s7 offset:0x12345 glc
+// CHECK: [0x41,0x41,0x0b,0xc2,0x45,0x23,0x01,0x0e]
+
 s_atomic_sub s5, s[2:3], s0
 // CHECK: [0x41,0x01,0x0c,0xc2,0x00,0x00,0x00,0x00]
 
 s_atomic_sub s101, s[2:3], s0
 // CHECK: [0x41,0x19,0x0c,0xc2,0x00,0x00,0x00,0x00]
 
 s_atomic_sub flat_scratch_lo, s[2:3], s0
 // CHECK: [0x81,0x19,0x0c,0xc2,0x00,0x00,0x00,0x00]

llvm/test/MC/Disassembler/AMDGPU/gfx10_dasm_all.txt

 0x02,0x04,0xe8,0x85
 
 # GFX10: s_cselect_b64 vcc, s[2:3], s[4:5]       ; encoding: [0x02,0x04,0xea,0x85]
 0x02,0x04,0xea,0x85
 
 # GFX10: s_atomic_add s5, s[2:3], 0x64           ; encoding: [0x41,0x01,0x08,0xf6,0x64,0x00,0x00,0xfa]
 0x41,0x01,0x08,0xf6,0x64,0x00,0x00,0xfa
 
+# GFX10: s_atomic_add s5, s[2:3], s7 offset:0x64 ; encoding: [0x41,0x01,0x08,0xf6,0x64,0x00,0x00,0x0e]
+0x41,0x01,0x08,0xf6,0x64,0x00,0x00,0x0e
+
 # GFX10: s_atomic_add s5, s[2:3], 0x64 dlc       ; encoding: [0x41,0x41,0x08,0xf6,0x64,0x00,0x00,0xfa]
 0x41,0x41,0x08,0xf6,0x64,0x00,0x00,0xfa
 
 # GFX10: s_atomic_add s5, s[2:3], 0x64 glc       ; encoding: [0x41,0x01,0x09,0xf6,0x64,0x00,0x00,0xfa]
 0x41,0x01,0x09,0xf6,0x64,0x00,0x00,0xfa
 
+# GFX10: s_atomic_add s5, s[2:3], s7 offset:0x64 glc ; encoding: [0x41,0x01,0x09,0xf6,0x64,0x00,0x00,0x0e]
+0x41,0x01,0x09,0xf6,0x64,0x00,0x00,0x0e
+
 # GFX10: s_atomic_add s5, s[2:3], s101           ; encoding: [0x41,0x01,0x08,0xf6,0x00,0x00,0x00,0xca]
 0x41,0x01,0x08,0xf6,0x00,0x00,0x00,0xca
 
 # GFX10: s_atomic_add s5, s[2:3], s101 dlc       ; encoding: [0x41,0x41,0x08,0xf6,0x00,0x00,0x00,0xca]
 0x41,0x41,0x08,0xf6,0x00,0x00,0x00,0xca
 
 # GFX10: s_atomic_add s5, s[2:3], s101 glc       ; encoding: [0x41,0x01,0x09,0xf6,0x00,0x00,0x00,0xca]
 0x41,0x01,0x09,0xf6,0x00,0x00,0x00,0xca

llvm/test/MC/Disassembler/AMDGPU/smem_gfx9.txt

 
 #===------------------------------------------------------------------------===#
 # s_atomic
 #===------------------------------------------------------------------------===#
 
 # GFX9: s_atomic_add s5, s[2:3], s101    ; encoding: [0x41,0x01,0x08,0xc2,0x65,0x00,0x00,0x00]
 0x41,0x01,0x08,0xc2,0x65,0x00,0x00,0x00
 
+# The SGRP variants can alternatively be encoded with imm=0, soffset_en=1
+# and the offset register encoded in the soffset field with the offset
+# field being disregarded.
+# GFX9: s_atomic_add s5, s[2:3], s101    ; encoding: [0x41,0x41,0x08,0xc2,0x00,0x00,0x00,0xca]
+0x41,0x41,0x08,0xc2,0x2e,0x00,0x00,0xca
+
+# The SGRP variants can alternatively be encoded with imm=0, soffset_en=1
+# and the offset register encoded in the soffset field with the offset
+# field being disregarded.
+# GFX9: s_atomic_add s5, s[2:3], s101 glc ; encoding: [0x41,0x41,0x09,0xc2,0x00,0x00,0x00,0xca]
+0x41,0x41,0x09,0xc2,0x2e,0x00,0x00,0xca
+
 # GFX9: s_atomic_add_x2 s[10:11], s[2:3], s101    ; encoding: [0x81,0x02,0x88,0xc2,0x65,0x00,0x00,0x00]
 0x81,0x02,0x88,0xc2,0x65,0x00,0x00,0x00
 
+# GFX9: s_atomic_add s5, s[2:3], s7 offset:0x12345 ; encoding: [0x41,0x41,0x0a,0xc2,0x45,0x23,0x01,0x0e]
+0x41,0x41,0x0a,0xc2,0x45,0x23,0x01,0x0e
+
+# GFX9: s_atomic_add s5, s[2:3], s7 offset:0x12345 glc ; encoding: [0x41,0x41,0x0b,0xc2,0x45,0x23,0x01,0x0e]
+0x41,0x41,0x0b,0xc2,0x45,0x23,0x01,0x0e
+
 # GFX9: s_atomic_and s101, s[2:3], s0    ; encoding: [0x41,0x19,0x20,0xc2,0x00,0x00,0x00,0x00]
 0x41,0x19,0x20,0xc2,0x00,0x00,0x00,0x00
 
 # GFX9: s_atomic_and_x2 s[12:13], s[2:3], s0    ; encoding: [0x01,0x03,0xa0,0xc2,0x00,0x00,0x00,0x00]
 0x01,0x03,0xa0,0xc2,0x00,0x00,0x00,0x00
 
 # GFX9: s_atomic_cmpswap s[10:11], s[2:3], s101    ; encoding: [0x81,0x02,0x04,0xc2,0x65,0x00,0x00,0x00]
 0x81,0x02,0x04,0xc2,0x65,0x00,0x00,0x00