AMDGPU: Fix legalization of MUBUF instructions in shaders

[nhaehnle]

Summary

The addr64-based legalization is incorrect for MUBUF instructions with idxen
set as well as for BUFFER_LOAD/STORE_FORMAT_* instructions. This affects
e.g. shaders that access buffer textures.

Since we never actually need the addr64-legalization in shaders, this patch
takes the easy route and keys off the calling convention. If this ever
affects (non-OpenGL) compute, the type of legalization needs to be chosen
based on some TSFlag.

Bugzilla: https://bugs.freedesktop.org/show_bug.cgi?id=98664

[• tstellarAMD]

LGTM.

[arsenm]

test/CodeGen/AMDGPU/mubuf-shader-vgpr.ll

49–50 ↗

(On Diff #78185)

This declaration isn't used

llvm/trunk/lib/Target/AMDGPU/SIInstrInfo.cpp

if (!Def)
     return;
 
   // Try to eliminate the copy if it is copying an immediate value.
   if (Def->isMoveImmediate())
     FoldImmediate(*Copy, *Def, OpReg, &MRI);
 }
 
 void SIInstrInfo::legalizeOperands(MachineInstr &MI) const {
-  MachineRegisterInfo &MRI = MI.getParent()->getParent()->getRegInfo();
+  MachineFunction &MF = *MI.getParent()->getParent();
+  MachineRegisterInfo &MRI = MF.getRegInfo();
 
   // Legalize VOP2
   if (isVOP2(MI) || isVOPC(MI)) {
     legalizeOperandsVOP2(MRI, MI);
     return;
   }
 
   // Legalize VOP3
if (MI.getOpcode() == AMDGPU::INSERT_SUBREG) {
     if (DstRC != Src0RC) {
       MachineBasicBlock *MBB = MI.getParent();
       MachineOperand &Op = MI.getOperand(1);
       legalizeGenericOperand(*MBB, MI, DstRC, Op, MRI, MI.getDebugLoc());
     }
     return;
   }
 
-  // Legalize MIMG
-  if (isMIMG(MI)) {
+  // Legalize MIMG and MUBUF/MTBUF for shaders.
+  //
+  // Shaders only generate MUBUF/MTBUF instructions via intrinsics or via
+  // scratch memory access. In both cases, the legalization never involves
+  // conversion to the addr64 form.
+  if (isMIMG(MI) ||
+      (AMDGPU::isShader(MF.getFunction()->getCallingConv()) &&
+       (isMUBUF(MI) || isMTBUF(MI)))) {
     MachineOperand *SRsrc = getNamedOperand(MI, AMDGPU::OpName::srsrc);
     if (SRsrc && !RI.isSGPRClass(MRI.getRegClass(SRsrc->getReg()))) {
       unsigned SGPR = readlaneVGPRToSGPR(SRsrc->getReg(), MI, MRI);
       SRsrc->setReg(SGPR);
     }
 
     MachineOperand *SSamp = getNamedOperand(MI, AMDGPU::OpName::ssamp);
     if (SSamp && !RI.isSGPRClass(MRI.getRegClass(SSamp->getReg()))) {
       unsigned SGPR = readlaneVGPRToSGPR(SSamp->getReg(), MI, MRI);
       SSamp->setReg(SGPR);
     }
     return;
   }
 
-  // Legalize MUBUF* instructions
+  // Legalize MUBUF* instructions by converting to addr64 form.
   // FIXME: If we start using the non-addr64 instructions for compute, we
-  // may need to legalize them here.
+  // may need to legalize them as above. This especially applies to the
+  // buffer_load_format_* variants and variants with idxen (or bothen).
   int SRsrcIdx =
       AMDGPU::getNamedOperandIdx(MI.getOpcode(), AMDGPU::OpName::srsrc);
   if (SRsrcIdx != -1) {
     // We have an MUBUF instruction
     MachineOperand *SRsrc = &MI.getOperand(SRsrcIdx);
     unsigned SRsrcRC = get(MI.getOpcode()).OpInfo[SRsrcIdx].RegClass;
     if (RI.getCommonSubClass(MRI.getRegClass(SRsrc->getReg()),
                                              RI.getRegClass(SRsrcRC))) {

llvm/trunk/test/CodeGen/AMDGPU/mubuf-shader-vgpr.ll

+;RUN: llc < %s -march=amdgcn -mcpu=verde -verify-machineinstrs | FileCheck %s -check-prefix=CHECK
+;RUN: llc < %s -march=amdgcn -mcpu=tonga -verify-machineinstrs | FileCheck %s -check-prefix=CHECK
+
+; Test that buffer_load_format with VGPR resource descriptor is properly
+; legalized.
+
+; CHECK-LABEL: {{^}}test_none:
+; CHECK: buffer_load_format_x v0, off, {{s\[[0-9]+:[0-9]+\]}}, 0{{$}}
+define amdgpu_vs float @test_none(<4 x i32> addrspace(2)* inreg %base, i32 %i) {
+main_body:
+  %ptr = getelementptr <4 x i32>, <4 x i32> addrspace(2)* %base, i32 %i
+  %tmp2 = load <4 x i32>, <4 x i32> addrspace(2)* %ptr, align 32
+  %tmp7 = call float @llvm.amdgcn.buffer.load.format.f32(<4 x i32> %tmp2, i32 0, i32 0, i1 0, i1 0)
+  ret float %tmp7
+}
+
+; CHECK-LABEL: {{^}}test_idxen:
+; CHECK: buffer_load_format_x v0, {{v[0-9]+}}, {{s\[[0-9]+:[0-9]+\]}}, 0 idxen{{$}}
+define amdgpu_vs float @test_idxen(<4 x i32> addrspace(2)* inreg %base, i32 %i) {
+main_body:
+  %ptr = getelementptr <4 x i32>, <4 x i32> addrspace(2)* %base, i32 %i
+  %tmp2 = load <4 x i32>, <4 x i32> addrspace(2)* %ptr, align 32
+  %tmp7 = call float @llvm.amdgcn.buffer.load.format.f32(<4 x i32> %tmp2, i32 undef, i32 0, i1 0, i1 0)
+  ret float %tmp7
+}
+
+; CHECK-LABEL: {{^}}test_offen:
+; CHECK: buffer_load_format_x v0, {{v[0-9]+}}, {{s\[[0-9]+:[0-9]+\]}}, 0 offen{{$}}
+define amdgpu_vs float @test_offen(<4 x i32> addrspace(2)* inreg %base, i32 %i) {
+main_body:
+  %ptr = getelementptr <4 x i32>, <4 x i32> addrspace(2)* %base, i32 %i
+  %tmp2 = load <4 x i32>, <4 x i32> addrspace(2)* %ptr, align 32
+  %tmp7 = call float @llvm.amdgcn.buffer.load.format.f32(<4 x i32> %tmp2, i32 0, i32 undef, i1 0, i1 0)
+  ret float %tmp7
+}
+
+; CHECK-LABEL: {{^}}test_both:
+; CHECK: buffer_load_format_x v0, {{v\[[0-9]+:[0-9]+\]}}, {{s\[[0-9]+:[0-9]+\]}}, 0 idxen offen{{$}}
+define amdgpu_vs float @test_both(<4 x i32> addrspace(2)* inreg %base, i32 %i) {
+main_body:
+  %ptr = getelementptr <4 x i32>, <4 x i32> addrspace(2)* %base, i32 %i
+  %tmp2 = load <4 x i32>, <4 x i32> addrspace(2)* %ptr, align 32
+  %tmp7 = call float @llvm.amdgcn.buffer.load.format.f32(<4 x i32> %tmp2, i32 undef, i32 undef, i1 0, i1 0)
+  ret float %tmp7
+}
+
+declare float @llvm.amdgcn.buffer.load.format.f32(<4 x i32>, i32, i32, i1, i1) nounwind readonly
+
+attributes #0 = { nounwind readnone }