AMDGPU/SI: Don't emit multi-dword flat memory ops when they might access scratch

[• tstellarAMD]

Summary

A single flat memory operations that might access the scratch buffer
can only access MaxPrivateElementSize bytes.

[mgrang]

test/CodeGen/AMDGPU/flat-address-space.ll

1

-check-prefixes=CHECK is not needed.

[• tstellarAMD]

Remove unneccessary -check-prefix

[tony-tye]

LGTM

[kzhuravl]

LGTM

lib/Target/AMDGPU/SIISelLowering.cpp

SDValue SITargetLowering::LowerLOAD(SDValue Op, SelectionDAG &DAG) const {
   unsigned AS = Load->getAddressSpace();
   if (!allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), MemVT,
                           AS, Load->getAlignment())) {
     SDValue Ops[2];
     std::tie(Ops[0], Ops[1]) = expandUnalignedLoad(Load, DAG);
     return DAG.getMergeValues(Ops, DL);
   }
 
+  MachineFunction &MF = DAG.getMachineFunction();
+  SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
+  // If there is a possibilty that flat instruction access scratch memory
+  // then we need to use the same legalization rules we use for private.
+  if (AS == AMDGPUAS::FLAT_ADDRESS)
+    AS = MFI->hasFlatScratchInit() ?
+         AMDGPUAS::PRIVATE_ADDRESS : AMDGPUAS::GLOBAL_ADDRESS;
+
   unsigned NumElements = MemVT.getVectorNumElements();
   switch (AS) {
   case AMDGPUAS::CONSTANT_ADDRESS:
     if (isMemOpUniform(Load))
       return SDValue();
     // Non-uniform loads will be selected to MUBUF instructions, so they
     // have the same legalization requires ments as global and private
     // loads.
assert(VT.isVector() &&
          Store->getValue().getValueType().getScalarType() == MVT::i32);
 
   unsigned AS = Store->getAddressSpace();
   if (!allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), VT,
                           AS, Store->getAlignment())) {
     return expandUnalignedStore(Store, DAG);
   }
 
+  MachineFunction &MF = DAG.getMachineFunction();
+  SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
+  // If there is a possibilty that flat instruction access scratch memory
+  // then we need to use the same legalization rules we use for private.
+  if (AS == AMDGPUAS::FLAT_ADDRESS)
+    AS = MFI->hasFlatScratchInit() ?
+         AMDGPUAS::PRIVATE_ADDRESS : AMDGPUAS::GLOBAL_ADDRESS;
+
   unsigned NumElements = VT.getVectorNumElements();
   switch (AS) {
   case AMDGPUAS::GLOBAL_ADDRESS:
   case AMDGPUAS::FLAT_ADDRESS:
     if (NumElements > 4)
       return SplitVectorStore(Op, DAG);
     return SDValue();
   case AMDGPUAS::PRIVATE_ADDRESS: {

test/CodeGen/AMDGPU/flat-address-space.ll

-; RUN: llc -O0 -mtriple=amdgcn-mesa-mesa3d -mcpu=bonaire < %s | FileCheck -check-prefix=CHECK %s
-; RUN: llc -O0 -mtriple=amdgcn-mesa-mesa3d -mcpu=tonga < %s | FileCheck -check-prefix=CHECK %s
+; RUN: llc -O0 -mtriple=amdgcn-mesa-mesa3d -mcpu=bonaire < %s | FileCheck -check-prefixes=CHECK %s
+; RUN: llc -O0 -mtriple=amdgcn-mesa-mesa3d -mcpu=tonga < %s | FileCheck -check-prefixes=CHECK %s
+; RUN: llc -O0 -mtriple=amdgcn-amd-amdhsa -mcpu=fiji < %s | FileCheck -check-prefixes=CHECK,HSA %s
 
 ; Disable optimizations in case there are optimizations added that
 ; specialize away generic pointer accesses.
 
 
 ; These testcases might become useless when there are optimizations to
 ; remove generic pointers.
 

 ; CHECK: flat_store_byte
 define void @flat_scratch_unaligned_store() {
   %scratch = alloca i32
   %fptr = addrspacecast i32* %scratch to i32 addrspace(4)*
   store volatile i32 0, i32 addrspace(4)* %fptr, align 1
   ret void
 }
 
+; CHECK-LABEL: flat_scratch_multidword_load:
+; HSA: flat_load_dword
+; HSA: flat_load_dword
+; FIXME: These tests are broken for os = mesa3d, becasue it doesn't initialize flat_scr
+define void @flat_scratch_multidword_load() {
+  %scratch = alloca <2 x i32>
+  %fptr = addrspacecast <2 x i32>* %scratch to <2 x i32> addrspace(4)*
+  %ld = load volatile <2 x i32>, <2 x i32> addrspace(4)* %fptr
+  ret void
+}
+
+; CHECK-LABEL: flat_scratch_multidword_store:
+; HSA: flat_store_dword
+; HSA: flat_store_dword
+; FIXME: These tests are broken for os = mesa3d, becasue it doesn't initialize flat_scr
+define void @flat_scratch_multidword_store() {
+  %scratch = alloca <2 x i32>
+  %fptr = addrspacecast <2 x i32>* %scratch to <2 x i32> addrspace(4)*
+  store volatile <2 x i32> zeroinitializer, <2 x i32> addrspace(4)* %fptr
+  ret void
+}
+
 attributes #0 = { nounwind }
 attributes #1 = { nounwind convergent }
 attributes #3 = { nounwind readnone }