[AMDGPU] Fix-up cases where writelane has 2 SGPR operands

[dstuttard]

Summary

Even though writelane doesn't have the same constraints as other valu
instructions it still can't violate the >1 SGPR operand constraint

Due to later register propagation (e.g. fixing up vgpr operands via
readfirstlane) changing writelane to only have a single SGPR is tricky.

This implementation puts a new check after SIFixSGPRCopies that prevents
multiple SGPRs being used in any writelane instructions.

The algorithm used is to check for trivial copy prop of suitable constants into
one of the SGPR operands and perform that if possible. If this isn't possible
put an explicit copy of Src1 SGPR into M0 and use that instead (this is
allowable for writelane as the constraint is for SGPR read-port and not
constant-bus access).

[rampitec]

lib/Target/AMDGPU/SIFixSGPRCopies.cpp

457 ↗

(On Diff #164878)

Why cannot you use loop in the runOnMachineFunction()?

482 ↗

(On Diff #164878)

You do not need vector here.

for (auto MO : {&Src0, &Src1})

[dstuttard]

lib/Target/AMDGPU/SIFixSGPRCopies.cpp

457 ↗

(On Diff #164878)

I think we need all the sgpr to vgpr moves to have been completed before applying this fix since in some cases it might not be necessary.
I guess there's an argument for this to be done in a separate pass, or a later pass, then it could go into runOnMachineFunction - any suggestions?

482 ↗

(On Diff #164878)

Good point, I'll make this one.

[rampitec]

lib/Target/AMDGPU/SIFixSGPRCopies.cpp

457 ↗

(On Diff #164878)

Given the semantics of writelane it is hard to believe its sources will be moved to VALU. Also if that is going to happen in general, it should have already happened by the time iterator would reach the instruction.

[dstuttard]

Made suggested changes

[dstuttard]

lib/Target/AMDGPU/SIFixSGPRCopies.cpp

457 ↗

(On Diff #164878)

Yes - I think you're right. I've changed it part of the main runOnMachineFunction.

[rampitec]

lib/Target/AMDGPU/SIFixSGPRCopies.cpp

751 ↗

(On Diff #165104)

missing break.

[dstuttard]

Added missing break

[dstuttard]

lib/Target/AMDGPU/SIFixSGPRCopies.cpp

751 ↗

(On Diff #165104)

Doh

[rampitec]

LGTM

[arsenm]

lib/Target/AMDGPU/Utils/AMDGPUMCUtils.cpp

23–25 ↗

(On Diff #165215)

This seems more like something for TII

[arsenm]

Should have a special check in the verifier

[dstuttard]

Moved foldToImm into SIInstrInfo as suggested
Implemented check in verifyInstruction and checked that it worked when the fix was removed

[arsenm]

I don't actually understand why this code is where it is? Why is SIFixSGPRCopies doing this? To clarify is this just an optimization? My initial reaction was that it was a fix, but looking at it again it seems like an optimization to me

lib/Target/AMDGPU/SIFixSGPRCopies.cpp

744–746 ↗

(On Diff #165274)

This should be a COPY?

[dstuttard]

Updated a mov to a copy as per review comment

[dstuttard]

In D51932#1236779, @arsenm wrote:

I don't actually understand why this code is where it is? Why is SIFixSGPRCopies doing this? To clarify is this just an optimization? My initial reaction was that it was a fix, but looking at it again it seems like an optimization to me

It isn't an optimization - it's a bug. We encountered this in graphics shaders - hence the requirement for the fix.

SIFixSGPRCopies does feel like a strange place to put this fix. It has to be somewhere late enough to catch the issue since it's a transformation that happens after isel that causes the problem in this case. Have you got any other suggestions as to where it could go instead? FixSGPR copies isn't perfect, but doesn't seem too unreasonable as it is an SGPR related fix-up.

[nhaehnle]

LGTM, to be honest. Matt?

[dstuttard]

@arsenm Matt - good to go?

[dstuttard]

ping

[dstuttard]

ping

[arsenm]

LGTM

[dstuttard]

GFX10 support has gone in since this change was approved - gfx10 allows 2 sgprs
on the constant bus. Implementation updated to allow for this.

Also updated a MIR test that had an incorrect WRITELANE instruction with 2 SGPR
accesses (one of which was VCC_LO).

[arsenm]

lib/Target/AMDGPU/SIFixSGPRCopies.cpp

798–799 ↗

(On Diff #220667)

I think there's a missing word here; " as the lane selector and doesn't"

800 ↗

(On Diff #220667)

Missing period

816–818 ↗

(On Diff #220667)

Is this really necessary? Will SIFoldOperands not get this for some reason?

lib/Target/AMDGPU/SIInstrInfo.cpp

3438 ↗

(On Diff #220667)

Extra empty line

3440 ↗

(On Diff #220667)

Should use Register

6494 ↗

(On Diff #220667)

I would expect this to handle only a single vreg def

[dstuttard]

Made some changes based on review

[dstuttard]

lib/Target/AMDGPU/SIFixSGPRCopies.cpp

798–799 ↗

(On Diff #220667)

I think this might be a difference in language - so I've changed it to make it less confusing

816–818 ↗

(On Diff #220667)

Yes, SIFoldOperands can replace the operand with the immediate, but the problem is that this code is forced to change one of the SGPR operands to M0 if it doesn't know if the other can be immediate.
If we just detect that an immediate will be folded, but don't do it, then the verification step will fail.

So, the simplest thing to do is detect trivial cases and change the operand to an immediate.
In some cases it is possible that it will not detect an immediate, replace the second SGPR with M0 and then the SIFoldOperands will replace the other SGPR operand with an immediate. This case is rare and functionally correct, just slightly less efficient.

I've reverted moving the foldToImm function from the SDWA peephole pass and implemented a simpler detection of immediates.

lib/Target/AMDGPU/SIInstrInfo.cpp

6494 ↗

(On Diff #220667)

I reverted moving the foldToImm function to SIInstrInfo - it was simpler to write a new one that attempted to do less.

[arsenm]

lib/Target/AMDGPU/SIFixSGPRCopies.cpp

820 ↗

(On Diff #220882)

No auto

821–822 ↗

(On Diff #220882)

The hasOneDef check is suspicious. You should be able to check getVRegDef and just a null check. This is missing a guard for virtual registers

[dstuttard]

Updates in light of review comments

[dstuttard]

Made suggested changes - is this more in line with what you were thinking Matt?

[dstuttard]

Matt - are you now happy for me to submit this? (It is tagged as approved, but since you've made some extra comments I'm waiting for you to agree with the latest changes).

[dstuttard]

ping

[nhaehnle]

I think this should be good to go.

[arsenm]

LGTM

[arsenm]

lib/Target/AMDGPU/SIFixSGPRCopies.cpp

815–816 ↗

(On Diff #221267)

I'm working on a patch to stop reserving m0; I suspect this will avoid the need for the special case propagation

llvm/lib/Target/AMDGPU/SIFixSGPRCopies.cpp

for (MachineBasicBlock::iterator I = MBB.begin(), E = MBB.end();
         if (TRI->isSGPRClass(DstRC) &&
             (TRI->hasVectorRegisters(Src0RC) ||
              TRI->hasVectorRegisters(Src1RC))) {
           LLVM_DEBUG(dbgs() << " Fixing INSERT_SUBREG: " << MI);
           TII->moveToVALU(MI, MDT);
         }
         break;
       }
+      case AMDGPU::V_WRITELANE_B32: {
+        // Some architectures allow more than one constant bus access without
+        // SGPR restriction
+        if (ST.getConstantBusLimit(MI.getOpcode()) != 1)
+          break;
+
+        // Writelane is special in that it can use SGPR and M0 (which would
+        // normally count as using the constant bus twice - but in this case it
+        // is allowed since the lane selector doesn't count as a use of the
+        // constant bus). However, it is still required to abide by the 1 SGPR
+        // rule. Apply a fix here as we might have multiple SGPRs after
+        // legalizing VGPRs to SGPRs
+        int Src0Idx =
+            AMDGPU::getNamedOperandIdx(MI.getOpcode(), AMDGPU::OpName::src0);
+        int Src1Idx =
+            AMDGPU::getNamedOperandIdx(MI.getOpcode(), AMDGPU::OpName::src1);
+        MachineOperand &Src0 = MI.getOperand(Src0Idx);
+        MachineOperand &Src1 = MI.getOperand(Src1Idx);
+
+        // Check to see if the instruction violates the 1 SGPR rule
+        if ((Src0.isReg() && TRI->isSGPRReg(*MRI, Src0.getReg()) &&
+             Src0.getReg() != AMDGPU::M0) &&
+            (Src1.isReg() && TRI->isSGPRReg(*MRI, Src1.getReg()) &&
+             Src1.getReg() != AMDGPU::M0)) {
+
+          // Check for trivially easy constant prop into one of the operands
+          // If this is the case then perform the operation now to resolve SGPR
+          // issue. If we don't do that here we will always insert a mov to m0
+          // that can't be resolved in later operand folding pass
+          bool Resolved = false;
+          for (MachineOperand *MO : {&Src0, &Src1}) {
+            if (Register::isVirtualRegister(MO->getReg())) {
+              MachineInstr *DefMI = MRI->getVRegDef(MO->getReg());
+              if (DefMI && TII->isFoldableCopy(*DefMI)) {
+                const MachineOperand &Def = DefMI->getOperand(0);
+                if (Def.isReg() &&
+                    MO->getReg() == Def.getReg() &&
+                    MO->getSubReg() == Def.getSubReg()) {
+                  const MachineOperand &Copied = DefMI->getOperand(1);
+                  if (Copied.isImm() &&
+                      TII->isInlineConstant(APInt(64, Copied.getImm(), true))) {
+                    MO->ChangeToImmediate(Copied.getImm());
+                    Resolved = true;
+                    break;
+                  }
+                }
+              }
+            }
+          }
+
+          if (!Resolved) {
+            // Haven't managed to resolve by replacing an SGPR with an immediate
+            // Move src1 to be in M0
+            BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
+                    TII->get(AMDGPU::COPY), AMDGPU::M0)
+                .add(Src1);
+            Src1.ChangeToRegister(AMDGPU::M0, false);
+          }
+        }
+        break;
+      }
       }
     }
   }
 
   if (MF.getTarget().getOptLevel() > CodeGenOpt::None && EnableM0Merge)
     hoistAndMergeSGPRInits(AMDGPU::M0, *MRI, TRI, *MDT, TII);
 
   return true;

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

if (isVOP3(MI) && LiteralCount) {
       }
       if (LiteralCount > 1) {
         ErrInfo = "VOP3 instruction uses more than one literal";
         return false;
       }
     }
   }
 
+  // Special case for writelane - this can break the multiple constant bus rule,
+  // but still can't use more than one SGPR register
+  if (Desc.getOpcode() == AMDGPU::V_WRITELANE_B32) {
+    unsigned SGPRCount = 0;
+    Register SGPRUsed = AMDGPU::NoRegister;
+
+    for (int OpIdx : {Src0Idx, Src1Idx, Src2Idx}) {
+      if (OpIdx == -1)
+        break;
+
+      const MachineOperand &MO = MI.getOperand(OpIdx);
+
+      if (usesConstantBus(MRI, MO, MI.getDesc().OpInfo[OpIdx])) {
+        if (MO.isReg() && MO.getReg() != AMDGPU::M0) {
+          if (MO.getReg() != SGPRUsed)
+            ++SGPRCount;
+          SGPRUsed = MO.getReg();
+        }
+      }
+      if (SGPRCount > ST.getConstantBusLimit(Opcode)) {
+        ErrInfo = "WRITELANE instruction violates constant bus restriction";
+        return false;
+      }
+    }
+  }
+
   // Verify misc. restrictions on specific instructions.
   if (Desc.getOpcode() == AMDGPU::V_DIV_SCALE_F32 ||
       Desc.getOpcode() == AMDGPU::V_DIV_SCALE_F64) {
     const MachineOperand &Src0 = MI.getOperand(Src0Idx);
     const MachineOperand &Src1 = MI.getOperand(Src1Idx);
     const MachineOperand &Src2 = MI.getOperand(Src2Idx);
     if (Src0.isReg() && Src1.isReg() && Src2.isReg()) {
       if (!compareMachineOp(Src0, Src1) &&

llvm/test/CodeGen/AMDGPU/inserted-wait-states.mir

bb.1:
     S_BRANCH %bb.2
 
   bb.2:
     $vgpr0,implicit $vcc = V_ADD_I32_e32 $vgpr1, $vgpr2, implicit $vcc, implicit $exec
     $sgpr4 = V_READLANE_B32 $vgpr4, $vcc_lo
     S_BRANCH %bb.3
 
   bb.3:
+    $m0 = S_MOV_B32 $sgpr4
     $vgpr0,implicit $vcc = V_ADD_I32_e32 $vgpr1, $vgpr2, implicit $vcc, implicit $exec
-    $vgpr4 = V_WRITELANE_B32 $sgpr4, $vcc_lo, $vgpr4
+    $vgpr4 = V_WRITELANE_B32 $m0, $vcc_lo, $vgpr4
     S_ENDPGM 0
 
 ...
 
 ...
 ---
 
 # GCN-LABEL: name: rfe

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.writelane.ll

-; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=gfx700 -verify-machineinstrs < %s | FileCheck %s
-; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=gfx802 -verify-machineinstrs < %s | FileCheck %s
+; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=gfx700 -verify-machineinstrs < %s | FileCheck -check-prefixes=CHECK,CI,CIGFX9 %s
+; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=gfx802 -verify-machineinstrs < %s | FileCheck -check-prefixes=CHECK,GFX9,CIGFX9 %s
+; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=gfx1010 -verify-machineinstrs < %s | FileCheck -check-prefixes=CHECK,GFX10 %s
 
 declare i32 @llvm.amdgcn.writelane(i32, i32, i32) #0
 
 ; CHECK-LABEL: {{^}}test_writelane_sreg:
-; CHECK: v_writelane_b32 v{{[0-9]+}}, s{{[0-9]+}}, s{{[0-9]+}}
+; CIGFX9: v_writelane_b32 v{{[0-9]+}}, s{{[0-9]+}}, m0
+; GFX10: v_writelane_b32 v{{[0-9]+}}, s{{[0-9]+}}, s{{[0-9]+}}
 define amdgpu_kernel void @test_writelane_sreg(i32 addrspace(1)* %out, i32 %src0, i32 %src1) #1 {
   %oldval = load i32, i32 addrspace(1)* %out
   %writelane = call i32 @llvm.amdgcn.writelane(i32 %src0, i32 %src1, i32 %oldval)
   store i32 %writelane, i32 addrspace(1)* %out, align 4
   ret void
 }
 
 ; CHECK-LABEL: {{^}}test_writelane_imm_sreg:
define amdgpu_kernel void @test_writelane_vreg_lane(i32 addrspace(1)* %out, <2 x i32> addrspace(1)* %in) #1 {
   %args = load <2 x i32>, <2 x i32> addrspace(1)* %gep.in
   %oldval = load i32, i32 addrspace(1)* %out
   %lane = extractelement <2 x i32> %args, i32 1
   %writelane = call i32 @llvm.amdgcn.writelane(i32 12, i32 %lane, i32 %oldval)
   store i32 %writelane, i32 addrspace(1)* %out, align 4
   ret void
 }
 
-; TODO: m0 should be folded.
 ; CHECK-LABEL: {{^}}test_writelane_m0_sreg:
 ; CHECK: s_mov_b32 m0, -1
 ; CHECK: s_mov_b32 [[COPY_M0:s[0-9]+]], m0
-; CHECK: v_writelane_b32 v{{[0-9]+}}, [[COPY_M0]], s{{[0-9]+}}
+; CIGFX9: v_writelane_b32 v{{[0-9]+}}, [[COPY_M0]], m0
+; GFX10: v_writelane_b32 v{{[0-9]+}}, [[COPY_M0]], s{{[0-9]+}}
 define amdgpu_kernel void @test_writelane_m0_sreg(i32 addrspace(1)* %out, i32 %src1) #1 {
   %oldval = load i32, i32 addrspace(1)* %out
   %m0 = call i32 asm "s_mov_b32 m0, -1", "={m0}"()
   %writelane = call i32 @llvm.amdgcn.writelane(i32 %m0, i32 %src1, i32 %oldval)
   store i32 %writelane, i32 addrspace(1)* %out, align 4
   ret void
 }
 
 ; CHECK-LABEL: {{^}}test_writelane_imm:
 ; CHECK: v_writelane_b32 v{{[0-9]+}}, s{{[0-9]+}}, 32
 define amdgpu_kernel void @test_writelane_imm(i32 addrspace(1)* %out, i32 %src0) #1 {
   %oldval = load i32, i32 addrspace(1)* %out
   %writelane = call i32 @llvm.amdgcn.writelane(i32 %src0, i32 32, i32 %oldval) #0
   store i32 %writelane, i32 addrspace(1)* %out, align 4
   ret void
 }
 
 ; CHECK-LABEL: {{^}}test_writelane_sreg_oldval:
 ; CHECK: v_mov_b32_e32 [[OLDVAL:v[0-9]+]], s{{[0-9]+}}
-; CHECK: v_writelane_b32 [[OLDVAL]], s{{[0-9]+}}, s{{[0-9]+}}
+; CIGFX9: v_writelane_b32 [[OLDVAL]], s{{[0-9]+}}, m0
+; GFX10: v_writelane_b32 [[OLDVAL]], s{{[0-9]+}}, s{{[0-9]+}}
 define amdgpu_kernel void @test_writelane_sreg_oldval(i32 inreg %oldval, i32 addrspace(1)* %out, i32 %src0, i32 %src1) #1 {
   %writelane = call i32 @llvm.amdgcn.writelane(i32 %src0, i32 %src1, i32 %oldval)
   store i32 %writelane, i32 addrspace(1)* %out, align 4
   ret void
 }
 
 ; CHECK-LABEL: {{^}}test_writelane_imm_oldval:
 ; CHECK: v_mov_b32_e32 [[OLDVAL:v[0-9]+]], 42
-; CHECK: v_writelane_b32 [[OLDVAL]], s{{[0-9]+}}, s{{[0-9]+}}
+; CIGFX9: v_writelane_b32 [[OLDVAL]], s{{[0-9]+}}, m0
+; GFX10: v_writelane_b32 [[OLDVAL]], s{{[0-9]+}}, s{{[0-9]+}}
 define amdgpu_kernel void @test_writelane_imm_oldval(i32 addrspace(1)* %out, i32 %src0, i32 %src1) #1 {
   %writelane = call i32 @llvm.amdgcn.writelane(i32 %src0, i32 %src1, i32 42)
   store i32 %writelane, i32 addrspace(1)* %out, align 4
   ret void
 }
 
 declare i32 @llvm.amdgcn.workitem.id.x() #2
 
 attributes #0 = { nounwind readnone convergent }
 attributes #1 = { nounwind }
 attributes #2 = { nounwind readnone }