[AMDGPU] Add llvm.amdgcn.raw.atomic.buffer.load intrinsic to support OpAtomicLoad lowering

[mariusz-sikora-at-amd]

Summary

This adds llvm.amdgcn.raw.atomic.buffer.load intrinsic to support
OpAtomicLoad lowering on AMDGPU. Previously this was lowered to
llvm.amdgcn.raw.buffer.load which caused the load in some cases
to be marked as invariant and hoisted in LICM.

[foad]

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.raw.atomic.buffer.load.ll

3

Should test globalisel too.

[arsenm]

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

977

We really ought to fix these bypassing the usual atomic memory model stuff

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.raw.atomic.buffer.load.ll

23

Use named values in tests

53

Needs tests with more types. i64, <2 x i16>, <4 x i16>, and some pointers

llvm/include/llvm/IR/IntrinsicsAMDGPU.td

llvm_i32_ty],      // auxiliary data (imm, cachepolicy     (bit 0 = glc,
                       //                                       bit 1 = slc,
                       //                                       bit 2 = dlc on gfx10+),
                       //                      swizzled buffer (bit 3 = swz))
   [IntrReadMem, IntrWillReturn, ImmArg<ArgIndex<3>>], "", [SDNPMemOperand]>,
   AMDGPURsrcIntrinsic<0>;
 def int_amdgcn_raw_buffer_load_format : AMDGPURawBufferLoad<llvm_anyfloat_ty>;
 def int_amdgcn_raw_buffer_load : AMDGPURawBufferLoad;
 
+class AMDGPURawAtomicBufferLoad<LLVMType data_ty = llvm_any_ty> : Intrinsic <
+  [data_ty],
+  [llvm_v4i32_ty,     // rsrc(SGPR)
+   llvm_i32_ty,       // offset(VGPR/imm, included in bounds checking and swizzling)
+   llvm_i32_ty,       // soffset(SGPR/imm, excluded from bounds checking and swizzling)
+   llvm_i32_ty],      // auxiliary data (imm, cachepolicy     (bit 0 = glc,
+                      //                                       bit 1 = slc,
+                      //                                       bit 2 = dlc on gfx10+),
+                      //                      swizzled buffer (bit 3 = swz))
+  [IntrArgMemOnly, ImmArg<ArgIndex<3>>], "", [SDNPMemOperand]>,
+  AMDGPURsrcIntrinsic<0>;
+def int_amdgcn_raw_atomic_buffer_load : AMDGPURawAtomicBufferLoad;
+
 class AMDGPUStructBufferLoad<LLVMType data_ty = llvm_any_ty> : Intrinsic <
   [data_ty],
   [llvm_v4i32_ty,     // rsrc(SGPR)
    llvm_i32_ty,       // vindex(VGPR)
    llvm_i32_ty,       // offset(VGPR/imm, included in bounds checking and swizzling)
    llvm_i32_ty,       // soffset(SGPR/imm, excluded from bounds checking and swizzling)
    llvm_i32_ty],      // auxiliary data (imm, cachepolicy     (bit 0 = glc,
                       //                                       bit 1 = slc,

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

case Intrinsic::amdgcn_struct_buffer_store:
     return legalizeBufferStore(MI, MRI, B, false, false);
   case Intrinsic::amdgcn_raw_buffer_store_format:
   case Intrinsic::amdgcn_struct_buffer_store_format:
     return legalizeBufferStore(MI, MRI, B, false, true);
   case Intrinsic::amdgcn_raw_tbuffer_store:
   case Intrinsic::amdgcn_struct_tbuffer_store:
     return legalizeBufferStore(MI, MRI, B, true, true);
   case Intrinsic::amdgcn_raw_buffer_load:
+  case Intrinsic::amdgcn_raw_atomic_buffer_load:
   case Intrinsic::amdgcn_struct_buffer_load:
     return legalizeBufferLoad(MI, MRI, B, false, false);
   case Intrinsic::amdgcn_raw_buffer_load_format:
   case Intrinsic::amdgcn_struct_buffer_load_format:
     return legalizeBufferLoad(MI, MRI, B, true, false);
   case Intrinsic::amdgcn_raw_tbuffer_load:
   case Intrinsic::amdgcn_struct_tbuffer_load:
     return legalizeBufferLoad(MI, MRI, B, true, true);

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

case Intrinsic::amdgcn_live_mask: {
       break;
     }
     case Intrinsic::amdgcn_wqm_demote:
     case Intrinsic::amdgcn_kill: {
       OpdsMapping[1] = AMDGPU::getValueMapping(AMDGPU::VCCRegBankID, 1);
       break;
     }
     case Intrinsic::amdgcn_raw_buffer_load:
+    case Intrinsic::amdgcn_raw_atomic_buffer_load:
     case Intrinsic::amdgcn_raw_tbuffer_load: {
       // FIXME: Should make intrinsic ID the last operand of the instruction,
       // then this would be the same as store
       OpdsMapping[0] = getVGPROpMapping(MI.getOperand(0).getReg(), MRI, *TRI);
       OpdsMapping[2] = getSGPROpMapping(MI.getOperand(2).getReg(), MRI, *TRI);
       OpdsMapping[3] = getVGPROpMapping(MI.getOperand(3).getReg(), MRI, *TRI);
       OpdsMapping[4] = getSGPROpMapping(MI.getOperand(4).getReg(), MRI, *TRI);
       break;

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

if (const AMDGPU::RsrcIntrinsic *RsrcIntr =
     if (RsrcIntr->IsImage) {
       Info.ptrVal = MFI->getImagePSV(TM);
       Info.align.reset();
     } else {
       Info.ptrVal = MFI->getBufferPSV(TM);
     }
 
     Info.flags |= MachineMemOperand::MODereferenceable;
-    if (ME.onlyReadsMemory()) {
+    if (ME.onlyReadsMemory() ||
+        IntrID == Intrinsic::amdgcn_raw_atomic_buffer_load) {
       unsigned MaxNumLanes = 4;
 
       if (RsrcIntr->IsImage) {
         const AMDGPU::ImageDimIntrinsicInfo *Intr
           = AMDGPU::getImageDimIntrinsicInfo(IntrID);
         const AMDGPU::MIMGBaseOpcodeInfo *BaseOpcode =
           AMDGPU::getMIMGBaseOpcodeInfo(Intr->BaseOpcode);
 
case Intrinsic::amdgcn_buffer_load_format: {
     if (LoadVT.getScalarType() == MVT::i8 ||
         LoadVT.getScalarType() == MVT::i16)
       return handleByteShortBufferLoads(DAG, LoadVT, DL, Ops, M);
 
     return getMemIntrinsicNode(Opc, DL, Op->getVTList(), Ops, IntVT,
                                M->getMemOperand(), DAG);
   }
   case Intrinsic::amdgcn_raw_buffer_load:
+  case Intrinsic::amdgcn_raw_atomic_buffer_load:
   case Intrinsic::amdgcn_raw_buffer_load_format: {
     const bool IsFormat = IntrID == Intrinsic::amdgcn_raw_buffer_load_format;
 
     auto Offsets = splitBufferOffsets(Op.getOperand(3), DAG);
     SDValue Ops[] = {
       Op.getOperand(0), // Chain
       Op.getOperand(2), // rsrc
       DAG.getConstant(0, DL, MVT::i32), // vindex

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.raw.atomic.buffer.load.ll

+;RUN: llc < %s -march=amdgcn -mcpu=verde -verify-machineinstrs | FileCheck %s -check-prefix=CHECK -check-prefix=SI
+;RUN: llc < %s -march=amdgcn -mcpu=tonga -verify-machineinstrs | FileCheck %s -check-prefix=CHECK -check-prefix=SI
+;RUN: llc < %s -march=amdgcn -mcpu=gfx1010 -verify-machineinstrs | FileCheck %s -check-prefix=CHECK -check-prefix=GFX10
+
+;CHECK-LABEL: {{^}}raw_atomic_buffer_load
+;CHECK-LABEL: BB0_1: ; %bb1
+;CHECK-NEXT: ; =>This Inner Loop Header: Depth=1
+;CHECK-NEXT: s_waitcnt lgkmcnt(0)
+;CHECK-NEXT: buffer_load_dword v1, off, s[0:3], 0 offset:4 glc
+;CHECK-NEXT: s_waitcnt vmcnt(0)
+;SI-NEXT: v_cmp_ne_u32_e32 vcc, v1, v0
+;GFX10-NEXT: v_cmp_ne_u32_e32 vcc_lo, v1, v0
+;SI-NEXT: s_or_b64 s[4:5], vcc, s[4:5]
+;GFX10-NEXT: s_or_b32 s4, vcc_lo, s4
+;SI-NEXT: s_andn2_b64 exec, exec, s[4:5]
+;GFX10-NEXT: s_andn2_b32 exec_lo, exec_lo, s4
+;CHECK-NEXT: s_cbranch_execnz .LBB0_1
+define amdgpu_kernel void @raw_atomic_buffer_load(<4 x i32> %addr) {
+bb:
+  %tmp0 = tail call i32 @llvm.amdgcn.workitem.id.x()
+  br label %bb1
+bb1:
+  %0 = call i32 @llvm.amdgcn.raw.atomic.buffer.load.i32(<4 x i32> %addr, i32 4, i32 0, i32 1)
+  %1 = icmp eq i32 %0, %tmp0
+  br i1 %1, label %bb1, label %bb2
+bb2:
+  ret void
+}
+
+;CHECK-LABEL: {{^}}raw_nonatomic_buffer_load
+;CHECK: ; =>This Inner Loop Header: Depth=1
+;SI-NEXT: s_and_b64 s[2:3], exec, vcc
+;GFX10-NEXT: s_and_b32 s1, exec_lo, vcc_lo
+;SI-NEXT: s_or_b64 s[0:1], s[2:3], s[0:1]
+;GFX10-NEXT: s_or_b32 s0, s1, s0
+;SI-NEXT: s_andn2_b64 exec, exec, s[0:1]
+;GFX10-NEXT: s_andn2_b32 exec_lo, exec_lo, s0
+;CHECK-NEXT: s_cbranch_execnz .LBB1_1
+define amdgpu_kernel void @raw_nonatomic_buffer_load(<4 x i32> %addr) {
+bb:
+  %tmp0 = tail call i32 @llvm.amdgcn.workitem.id.x()
+  br label %bb1
+bb1:
+  %0 = call i32 @llvm.amdgcn.raw.buffer.load.i32(<4 x i32> %addr, i32 4, i32 0, i32 1)
+  %1 = icmp eq i32 %0, %tmp0
+  br i1 %1, label %bb1, label %bb2
+bb2:
+  ret void
+}
+
+; Function Attrs: nounwind readonly
+declare i32 @llvm.amdgcn.raw.atomic.buffer.load.i32(<4 x i32>, i32, i32, i32 immarg)
+declare i32 @llvm.amdgcn.raw.buffer.load.i32(<4 x i32>, i32, i32, i32 immarg)
+declare void @llvm.amdgcn.raw.buffer.store.i32(i32, <4 x i32>, i32, i32, i32 immarg)
+declare i32 @llvm.amdgcn.workitem.id.x()
+