AMDGPU: Relax restriction on folding immediates into physregs

[arsenm]

Summary

I never completed the work on the patches referenced by
f8bf7d7f42f28fa18144091022236208e199f331, but this was intended to
avoid folding immediate writes into m0 which the coalescer doesn't
understand very well. Relax this to allow simple SGPR immediates to
fold directly into VGPR copies. This pattern shows up routinely in
current GlobalISel code since nothing is smart enough to emit VGPR
constants yet.

[rampitec]

I do not think this can work with physregs given the CFG.

llvm/lib/Target/AMDGPU/SIFoldOperands.cpp

654

Source must be strictly scalar. You can have a copy of vector to vector (say a to v) and it can be initialized with a different mask.

In fact, how can we be sure even an SGPR was not initialized in a different CFG path? This shall never happen to a virtual register in SSA, but can easily happen to a physreg.

658

This does fit into a single line.

[arsenm]

llvm/lib/Target/AMDGPU/SIFoldOperands.cpp

654

This is folding a constant. The mask doesn't matter. This is also not analyzing physreg defs, it's searching the uses of an SSA virtual register.

[rampitec]

llvm/lib/Target/AMDGPU/SIFoldOperands.cpp

654

This is folding a constant. The mask doesn't matter. This is also not analyzing physreg defs, it's searching the uses of an SSA virtual register.

Ah, right! It is destination physical, not source.

[rampitec]

It can probably create a worse code sometimes. Assume you have a constant which cannot be directly assigned to a register (v_accvgpr_write_b32 only accepts inline literals). If source def will not become dead it will create an extra copy. You probably need some legality checks here.

[rampitec]

In D84153#2180517, @rampitec wrote:

It can probably create a worse code sometimes. Assume you have a constant which cannot be directly assigned to a register (v_accvgpr_write_b32 only accepts inline literals). If source def will not become dead it will create an extra copy. You probably need some legality checks here.

But that's really orthogonal to what you are doing and is a separate change. It equally related to virtual registers.

LGTM, just fix formatting at line 657.

[arsenm]

766cb615a3b96025192707f4670cdf171da84034

llvm/lib/Target/AMDGPU/SIFoldOperands.cpp

if (frameIndexMayFold(TII, *UseMI, UseOpIdx, OpToFold)) {
     return;
   }
 
   bool FoldingImmLike =
       OpToFold.isImm() || OpToFold.isFI() || OpToFold.isGlobal();
 
   if (FoldingImmLike && UseMI->isCopy()) {
     Register DestReg = UseMI->getOperand(0).getReg();
+    Register SrcReg = UseMI->getOperand(1).getReg();
+    assert(SrcReg.isVirtual());
 
-    // Don't fold into a copy to a physical register. Doing so would interfere
-    // with the register coalescer's logic which would avoid redundant
-    // initalizations.
-    if (DestReg.isPhysical())
-      return;
+    const TargetRegisterClass *SrcRC = MRI->getRegClass(SrcReg);
 
-    const TargetRegisterClass *DestRC =  MRI->getRegClass(DestReg);
+    // Don't fold into a copy to a physical register with the same class. Doing
+    // so would interfere with the register coalescer's logic which would avoid
+    // redundant initalizations.
+    if (DestReg.isPhysical() && SrcRC->contains(DestReg))
+      return;
 
-    Register SrcReg = UseMI->getOperand(1).getReg();
-    if (SrcReg.isVirtual()) { // XXX - This can be an assert?
-      const TargetRegisterClass * SrcRC = MRI->getRegClass(SrcReg);
-      if (TRI->isSGPRClass(SrcRC) && TRI->hasVectorRegisters(DestRC)) {
-        MachineRegisterInfo::use_iterator NextUse;
-        SmallVector<FoldCandidate, 4> CopyUses;
-        for (MachineRegisterInfo::use_iterator
+    const TargetRegisterClass *DestRC =
+      TRI->getRegClassForReg(*MRI, DestReg);
              Use = MRI->use_begin(DestReg), E = MRI->use_end();
            Use != E; Use = NextUse) {
-          NextUse = std::next(Use);
-          FoldCandidate FC = FoldCandidate(Use->getParent(),
                                          Use.getOperandNo(), &UseMI->getOperand(1));
-          CopyUses.push_back(FC);
-       }
-        for (auto & F : CopyUses) {
-          foldOperand(*F.OpToFold, F.UseMI, F.UseOpNo,
                     FoldList, CopiesToReplace);
-        }
-      }
-    }
 
     if (DestRC == &AMDGPU::AGPR_32RegClass &&
         TII->isInlineConstant(OpToFold, AMDGPU::OPERAND_REG_INLINE_C_INT32)) {
       UseMI->setDesc(TII->get(AMDGPU::V_ACCVGPR_WRITE_B32));
       UseMI->getOperand(1).ChangeToImmediate(OpToFold.getImm());
       CopiesToReplace.push_back(UseMI);
       return;
     }

llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.kernarg.segment.ptr.ll

define amdgpu_kernel void @opencl_test_implicit_alignment_no_explicit_kernargs_round_up() #3 {
   %implicitarg.ptr = call noalias i8 addrspace(4)* @llvm.amdgcn.implicitarg.ptr()
   %arg.ptr = bitcast i8 addrspace(4)* %implicitarg.ptr to i32 addrspace(4)*
   %val = load volatile i32, i32 addrspace(4)* %arg.ptr
   store volatile i32 %val, i32 addrspace(1)* null
   ret void
 }
 
 ; ALL-LABEL: {{^}}func_kernarg_segment_ptr:
-; ALL: s_mov_b32 [[S_LO:s[0-9]+]], 0{{$}}
-; ALL: s_mov_b32 [[S_HI:s[0-9]+]], 0{{$}}
-; ALL: v_mov_b32_e32 v0, [[S_LO]]{{$}}
-; ALL: v_mov_b32_e32 v1, [[S_HI]]{{$}}
+; ALL: v_mov_b32_e32 v0, 0{{$}}
+; ALL: v_mov_b32_e32 v1, 0{{$}}
 define i8 addrspace(4)* @func_kernarg_segment_ptr() {
   %ptr = call i8 addrspace(4)* @llvm.amdgcn.kernarg.segment.ptr()
   ret i8 addrspace(4)* %ptr
 }
 
 declare i8 addrspace(4)* @llvm.amdgcn.kernarg.segment.ptr() #0
 declare i8 addrspace(4)* @llvm.amdgcn.implicitarg.ptr() #0
 
 attributes #0 = { nounwind readnone }
 attributes #1 = { nounwind }
 attributes #2 = { nounwind "amdgpu-implicitarg-num-bytes"="48" }
 attributes #3 = { nounwind "amdgpu-implicitarg-num-bytes"="38" }

llvm/test/CodeGen/AMDGPU/GlobalISel/zextload.ll

; GFX6-NEXT:    s_setpc_b64 s[30:31]
   ret i64 %ext
 }
 
 define i96 @zextload_global_i32_to_i96(i32 addrspace(1)* %ptr) {
 ; GFX9-LABEL: zextload_global_i32_to_i96:
 ; GFX9:       ; %bb.0:
 ; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX9-NEXT:    global_load_dword v0, v[0:1], off
-; GFX9-NEXT:    s_mov_b32 s4, 0
 ; GFX9-NEXT:    v_mov_b32_e32 v1, 0
-; GFX9-NEXT:    v_mov_b32_e32 v2, s4
+; GFX9-NEXT:    v_mov_b32_e32 v2, 0
 ; GFX9-NEXT:    s_waitcnt vmcnt(0)
 ; GFX9-NEXT:    s_setpc_b64 s[30:31]
 ;
 ; GFX8-LABEL: zextload_global_i32_to_i96:
 ; GFX8:       ; %bb.0:
 ; GFX8-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX8-NEXT:    flat_load_dword v0, v[0:1]
-; GFX8-NEXT:    s_mov_b32 s4, 0
 ; GFX8-NEXT:    v_mov_b32_e32 v1, 0
-; GFX8-NEXT:    v_mov_b32_e32 v2, s4
+; GFX8-NEXT:    v_mov_b32_e32 v2, 0
 ; GFX8-NEXT:    s_waitcnt vmcnt(0) lgkmcnt(0)
 ; GFX8-NEXT:    s_setpc_b64 s[30:31]
 ;
 ; GFX6-LABEL: zextload_global_i32_to_i96:
 ; GFX6:       ; %bb.0:
 ; GFX6-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX6-NEXT:    s_mov_b32 s6, 0
 ; GFX6-NEXT:    s_mov_b32 s7, 0xf000
 ; GFX6-NEXT:    s_mov_b64 s[4:5], 0
 ; GFX6-NEXT:    buffer_load_dword v0, v[0:1], s[4:7], 0 addr64
-; GFX6-NEXT:    s_mov_b32 s4, 0
 ; GFX6-NEXT:    v_mov_b32_e32 v1, 0
-; GFX6-NEXT:    v_mov_b32_e32 v2, s4
+; GFX6-NEXT:    v_mov_b32_e32 v2, 0
 ; GFX6-NEXT:    s_waitcnt vmcnt(0)
 ; GFX6-NEXT:    s_setpc_b64 s[30:31]
   %load = load i32, i32 addrspace(1)* %ptr
   %ext = zext i32 %load to i96
   ret i96 %ext
 }
 
 define i128 @zextload_global_i32_to_i128(i32 addrspace(1)* %ptr) {
 ; GFX9-LABEL: zextload_global_i32_to_i128:
 ; GFX9:       ; %bb.0:
 ; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX9-NEXT:    global_load_dword v0, v[0:1], off
-; GFX9-NEXT:    s_mov_b32 s4, 0
-; GFX9-NEXT:    s_mov_b32 s5, 0
 ; GFX9-NEXT:    v_mov_b32_e32 v1, 0
-; GFX9-NEXT:    v_mov_b32_e32 v2, s4
-; GFX9-NEXT:    v_mov_b32_e32 v3, s5
+; GFX9-NEXT:    v_mov_b32_e32 v2, 0
+; GFX9-NEXT:    v_mov_b32_e32 v3, 0
 ; GFX9-NEXT:    s_waitcnt vmcnt(0)
 ; GFX9-NEXT:    s_setpc_b64 s[30:31]
 ;
 ; GFX8-LABEL: zextload_global_i32_to_i128:
 ; GFX8:       ; %bb.0:
 ; GFX8-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX8-NEXT:    flat_load_dword v0, v[0:1]
-; GFX8-NEXT:    s_mov_b32 s4, 0
-; GFX8-NEXT:    s_mov_b32 s5, 0
 ; GFX8-NEXT:    v_mov_b32_e32 v1, 0
-; GFX8-NEXT:    v_mov_b32_e32 v2, s4
-; GFX8-NEXT:    v_mov_b32_e32 v3, s5
+; GFX8-NEXT:    v_mov_b32_e32 v2, 0
+; GFX8-NEXT:    v_mov_b32_e32 v3, 0
 ; GFX8-NEXT:    s_waitcnt vmcnt(0) lgkmcnt(0)
 ; GFX8-NEXT:    s_setpc_b64 s[30:31]
 ;
 ; GFX6-LABEL: zextload_global_i32_to_i128:
 ; GFX6:       ; %bb.0:
 ; GFX6-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX6-NEXT:    s_mov_b32 s6, 0
 ; GFX6-NEXT:    s_mov_b32 s7, 0xf000
 ; GFX6-NEXT:    s_mov_b64 s[4:5], 0
 ; GFX6-NEXT:    buffer_load_dword v0, v[0:1], s[4:7], 0 addr64
-; GFX6-NEXT:    s_mov_b32 s4, 0
-; GFX6-NEXT:    s_mov_b32 s5, 0
 ; GFX6-NEXT:    v_mov_b32_e32 v1, 0
-; GFX6-NEXT:    v_mov_b32_e32 v2, s4
-; GFX6-NEXT:    v_mov_b32_e32 v3, s5
+; GFX6-NEXT:    v_mov_b32_e32 v2, 0
+; GFX6-NEXT:    v_mov_b32_e32 v3, 0
 ; GFX6-NEXT:    s_waitcnt vmcnt(0)
 ; GFX6-NEXT:    s_setpc_b64 s[30:31]
   %load = load i32, i32 addrspace(1)* %ptr
   %ext = zext i32 %load to i128
   ret i128 %ext
 }

llvm/test/CodeGen/AMDGPU/fold-imm-copy.mir

 body:             |
   bb.0:
     %0:sreg_64 = S_MOV_B64 -8
     %1:sgpr_32 = COPY %0.sub0
     %2:sgpr_32 = COPY %0.sub1
     S_ENDPGM 0, implicit %1, implicit %2
 
 ...
+
+# GCN-LABEL: name: no_fold_imm_into_m0{{$}}
+# GCN: %0:sreg_32 = S_MOV_B32 -8
+# GCN-NEXT: $m0 = COPY %0
+
+---
+name: no_fold_imm_into_m0
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    %0:sreg_32 = S_MOV_B32 -8
+    $m0 = COPY %0
+    S_ENDPGM 0, implicit $m0
+
+...
+
+# GCN-LABEL: name: fold_sgpr_imm_to_vgpr_copy{{$}}
+# GCN: $vgpr0 = V_MOV_B32_e32 -8, implicit $exec
+---
+name: fold_sgpr_imm_to_vgpr_copy
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    %0:sreg_32 = S_MOV_B32 -8
+    $vgpr0 = COPY %0
+    S_ENDPGM 0, implicit $vgpr0
+
+...