AMDGPU: Add a pass to fix SGPR liveness

[ruiling]

Summary

Sometime, we may get a SGPR-PHI with only one meaningful value from
dominator and undefined in another predecessor. The problem with
such pattern is the SGPR might be overwritten in the undefined
predecessor. The pass added here is to make the SGPR also alive in the
originally dead predecessors.

[critson]

llvm/lib/Target/AMDGPU/SIFixSGPRLiveness.cpp

89

Stray comment?

94

Do these influence SingleOp?

128

Unnecessary continue?

[ruiling]

address review comments

[ruiling]

llvm/lib/Target/AMDGPU/SIFixSGPRLiveness.cpp

94

Thanks for point out this, I have refined this part to be more robust.

[ruiling]

any other concern?

[critson]

I think the code is correct, but this needs a MIR test checking the behaviour of the pass.

[foad]

My understanding is that this pass only fixes specific liveness problems that are created by D132450. It is not a general purpose pass that could handle any possible pattern of SGPR usage.

llvm/lib/Target/AMDGPU/SIFixSGPRLiveness.cpp

17

The implementation does not check that there is a divergent branch, so it would do the same thing for uniform control flow.

[ruiling]

In D132447#3758581, @foad wrote:

My understanding is that this pass only fixes specific liveness problems that are created by D132450. It is not a general purpose pass that could handle any possible pattern of SGPR usage.

Thank you for your comment @foad. But I don't agree the problem was created by D132450. see the attached sgpr-liveness.ll, a hand-made test can also show the problem (The tricky thing is that the existing structurizer still simplify the IR even it is already structurized, thus it helps hiding the issue, but I don't think it is the expected behavior). Yes the current design of the pass only works for this known pattern. That is just because I don't know if there are any other broken patterns.

llvm/lib/Target/AMDGPU/SIFixSGPRLiveness.cpp

17

Sounds true, this is doing things conservative, but I think it may not cause any regression in real case. Practically, such kind of pattern in source program should already been optimized into the single dominate value. The only possible cases that I can think of are either "opt-none" or "created by structurizeCFG". If such pattern is created by structurizeCFG, BB1 will be very likely to be divergent because structurizeCFG was designed to try to skip uniform branching as much as possible. I am not sure if checking the terminator of BB1 against s_cbranch_scc0/scc1 for uniform branch sounds good idea?

[ruiling]

Don't extend for backedge case.
Add mir tests.

[ruiling]

llvm/lib/Target/AMDGPU/SIFixSGPRLiveness.cpp

17

I will try to use the idea at https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/AMDGPU/SILowerI1Copies.cpp#L170 to detect divergent branch.

[ruiling]

Make it work in a few more cases.
Detect divergent branch.

[ruiling]

In D132447#3758581, @foad wrote:

My understanding is that this pass only fixes specific liveness problems that are created by D132450. It is not a general purpose pass that could handle any possible pattern of SGPR usage.

I think the new version should work correctly in all possible shapes of post-structurized CFG I can think of now. Please let me know if you have any further concern.

[ruiling]

abandon this in favor of D133840

llvm/lib/Target/AMDGPU/AMDGPU.h

 FunctionPass *createSILowerI1CopiesPass();
 FunctionPass *createSIShrinkInstructionsPass();
 FunctionPass *createSILoadStoreOptimizerPass();
 FunctionPass *createSIWholeQuadModePass();
 FunctionPass *createSIFixControlFlowLiveIntervalsPass();
 FunctionPass *createSIOptimizeExecMaskingPreRAPass();
 FunctionPass *createSIOptimizeVGPRLiveRangePass();
 FunctionPass *createSIFixSGPRCopiesPass();
+FunctionPass *createSIFixSGPRLivenessPass();
 FunctionPass *createSIMemoryLegalizerPass();
 FunctionPass *createSIInsertWaitcntsPass();
 FunctionPass *createSIPreAllocateWWMRegsPass();
 FunctionPass *createSIFormMemoryClausesPass();
 
 FunctionPass *createSIPostRABundlerPass();
 FunctionPass *createAMDGPUSimplifyLibCallsPass(const TargetMachine *);
 FunctionPass *createAMDGPUUseNativeCallsPass();

 extern char &SIPeepholeSDWAID;
 
 void initializeSIShrinkInstructionsPass(PassRegistry&);
 extern char &SIShrinkInstructionsID;
 
 void initializeSIFixSGPRCopiesPass(PassRegistry &);
 extern char &SIFixSGPRCopiesID;
 
+void initializeSIFixSGPRLivenessPass(PassRegistry &);
+extern char &SIFixSGPRLivenessID;
+
 void initializeSIFixVGPRCopiesPass(PassRegistry &);
 extern char &SIFixVGPRCopiesID;
 
 void initializeSILowerI1CopiesPass(PassRegistry &);
 extern char &SILowerI1CopiesID;
 
 void initializeSILowerSGPRSpillsPass(PassRegistry &);
 extern char &SILowerSGPRSpillsID;

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAMDGPUTarget() {
   initializeR600PacketizerPass(*PR);
   initializeR600ExpandSpecialInstrsPassPass(*PR);
   initializeR600VectorRegMergerPass(*PR);
   initializeGlobalISel(*PR);
   initializeAMDGPUDAGToDAGISelPass(*PR);
   initializeGCNDPPCombinePass(*PR);
   initializeSILowerI1CopiesPass(*PR);
   initializeSILowerSGPRSpillsPass(*PR);
+  initializeSIFixSGPRLivenessPass(*PR);
   initializeSIFixSGPRCopiesPass(*PR);
   initializeSIFixVGPRCopiesPass(*PR);
   initializeSIFoldOperandsPass(*PR);
   initializeSIPeepholeSDWAPass(*PR);
   initializeSIShrinkInstructionsPass(*PR);
   initializeSIOptimizeExecMaskingPreRAPass(*PR);
   initializeSIOptimizeVGPRLiveRangePass(*PR);
   initializeSILoadStoreOptimizerPass(*PR);
if (EnableEarlyIfConversion)
     addPass(&EarlyIfConverterID);
 
   TargetPassConfig::addILPOpts();
   return false;
 }
 
 bool GCNPassConfig::addInstSelector() {
   AMDGPUPassConfig::addInstSelector();
+  addPass(&SIFixSGPRLivenessID);
   addPass(&SIFixSGPRCopiesID);
   addPass(createSILowerI1CopiesPass());
   return false;
 }
 
 bool GCNPassConfig::addIRTranslator() {
   addPass(new IRTranslator(getOptLevel()));
   return false;

llvm/lib/Target/AMDGPU/CMakeLists.txt

add_llvm_target(AMDGPUCodeGen
   R600OptimizeVectorRegisters.cpp
   R600Packetizer.cpp
   R600RegisterInfo.cpp
   R600Subtarget.cpp
   R600TargetMachine.cpp
   R600TargetTransformInfo.cpp
   SIAnnotateControlFlow.cpp
   SIFixSGPRCopies.cpp
+  SIFixSGPRLiveness.cpp
   SIFixVGPRCopies.cpp
   SIFoldOperands.cpp
   SIFormMemoryClauses.cpp
   SIFrameLowering.cpp
   SIInsertHardClauses.cpp
   SIInsertWaitcnts.cpp
   SIInstrInfo.cpp
   SIISelLowering.cpp

llvm/lib/Target/AMDGPU/SIFixSGPRLiveness.cpp

+//===- SIFixSGPRLiveness.cpp - Fix SGPR Liveness --------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file
+/// Consider a simple example:
+///     BB1: %1:SGPR = COPY 1
+///     | \
+///     |  BB2: %2:SGPR = IMPLICIT_DEF
+///     | /
+///     BB3: %3 = phi (%2, BB2), (%1, BB1)
+///
+/// BB1 ends with a divergent branch. The virtual registers shown in above
+/// example are dead in BB2. After register allocation, they may end up
+/// being put in the same physical register, and the liveness of the register
+/// does not cover the range of BB2. If the corresponding physical register
+/// is reused in BB2, the content of the physical register will be overwritten.
+/// The idea in this pass is to extend the liveness of %1 through BB2.
+///
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
+#include "llvm/ADT/SetOperations.h"
+#include "llvm/CodeGen/MachineDominators.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Target/TargetMachine.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "si-fix-sgpr-liveness"
+
+namespace {
+
+class SIFixSGPRLiveness : public MachineFunctionPass {
+  MachineDominatorTree *MDT;
+
+public:
+  static char ID;
+
+  MachineRegisterInfo *MRI;
+  const SIRegisterInfo *TRI;
+
+  SIFixSGPRLiveness() : MachineFunctionPass(ID) {}
+
+  bool extendSGPRLiveRangeForPHI(MachineFunction &MF);
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  StringRef getPassName() const override { return "SI Fix SGPR Liveness"; }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<MachineDominatorTree>();
+    AU.addPreserved<MachineDominatorTree>();
+    AU.setPreservesCFG();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+};
+
+} // end anonymous namespace
+
+INITIALIZE_PASS_BEGIN(SIFixSGPRLiveness, DEBUG_TYPE, "SI Fix SGPR Liveness",
+                      false, false)
+INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
+INITIALIZE_PASS_END(SIFixSGPRLiveness, DEBUG_TYPE, "SI Fix SGPR Liveness",
+                    false, false)
+
+char SIFixSGPRLiveness::ID = 0;
+
+char &llvm::SIFixSGPRLivenessID = SIFixSGPRLiveness::ID;
+
+FunctionPass *llvm::createSIFixSGPRLivenessPass() {
+  return new SIFixSGPRLiveness();
+}
+
+bool SIFixSGPRLiveness::extendSGPRLiveRangeForPHI(MachineFunction &MF) {
+  bool Changed = false;
+  for (auto &MBB : MF) {
+    for (auto &MI : MBB.phis()) {
+      Register DstReg = MI.getOperand(0).getReg();
+      if (!TRI->isSGPRClass(MRI->getRegClass(DstReg)))
+        continue;
+      SmallVector<std::pair<MachineOperand *, MachineBasicBlock *>, 4> Undefs;
+      MachineOperand *SingleOp = nullptr;
+      MachineBasicBlock *SingleMBB = nullptr;
+      for (unsigned i = 1, e = MI.getNumOperands(); i != e; i += 2) {
+        MachineOperand *SrcOp = &MI.getOperand(i);
+
+        if (SrcOp->isReg() && SrcOp->getReg().isVirtual()) {
+          MachineInstr *Def = MRI->getVRegDef(SrcOp->getReg());
+          if (Def && Def->isImplicitDef()) {
+            Undefs.push_back(
+                std::make_pair(SrcOp, MI.getOperand(i + 1).getMBB()));
+            continue;
+          }
+        }
+
+        if (!SingleOp) {
+          SingleOp = SrcOp;
+          SingleMBB = MI.getOperand(i + 1).getMBB();
+        } else if (!SrcOp->isIdenticalTo(*SingleOp)) {
+          SingleOp = nullptr;
+          break;
+        }
+      }
+
+      if (!SingleOp || Undefs.empty())
+        continue;
+
+      // We have an single non-undef incoming value, make it alive
+      // through all paths from the SingleMBB to MBB.
+      if (all_of(Undefs, [&](auto UD) {
+            return MDT->dominates(SingleMBB, UD.second);
+          })) {
+        assert(SingleOp->isReg());
+        for (auto UndefOp : Undefs) {
+          UndefOp.first->setIsUndef(false);
+          UndefOp.first->setReg(SingleOp->getReg());
+          UndefOp.first->setSubReg(SingleOp->getSubReg());
+        }
+        Changed = true;
+      }
+    }
+  }
+  return Changed;
+}
+
+bool SIFixSGPRLiveness::runOnMachineFunction(MachineFunction &MF) {
+  const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
+  MRI = &MF.getRegInfo();
+  TRI = ST.getRegisterInfo();
+  MDT = &getAnalysis<MachineDominatorTree>();
+
+  return extendSGPRLiveRangeForPHI(MF);
+}

llvm/test/CodeGen/AMDGPU/llc-pipeline.ll

 ; GCN-O0-NEXT:        Safe Stack instrumentation pass
 ; GCN-O0-NEXT:        Insert stack protectors
 ; GCN-O0-NEXT:        Dominator Tree Construction
 ; GCN-O0-NEXT:        Post-Dominator Tree Construction
 ; GCN-O0-NEXT:        Natural Loop Information
 ; GCN-O0-NEXT:        Legacy Divergence Analysis
 ; GCN-O0-NEXT:        AMDGPU DAG->DAG Pattern Instruction Selection
 ; GCN-O0-NEXT:        MachineDominator Tree Construction
+; GCN-O0-NEXT:        SI Fix SGPR Liveness
 ; GCN-O0-NEXT:        SI Fix SGPR copies
 ; GCN-O0-NEXT:        MachinePostDominator Tree Construction
 ; GCN-O0-NEXT:        SI Lower i1 Copies
 ; GCN-O0-NEXT:        Finalize ISel and expand pseudo-instructions
 ; GCN-O0-NEXT:        Local Stack Slot Allocation
 ; GCN-O0-NEXT:        Register Usage Information Propagation
 ; GCN-O0-NEXT:        Eliminate PHI nodes for register allocation
 ; GCN-O0-NEXT:        SI Lower control flow pseudo instructions

 ; GCN-O1-NEXT:        Legacy Divergence Analysis
 ; GCN-O1-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O1-NEXT:        Function Alias Analysis Results
 ; GCN-O1-NEXT:        Branch Probability Analysis
 ; GCN-O1-NEXT:        Lazy Branch Probability Analysis
 ; GCN-O1-NEXT:        Lazy Block Frequency Analysis
 ; GCN-O1-NEXT:        AMDGPU DAG->DAG Pattern Instruction Selection
 ; GCN-O1-NEXT:        MachineDominator Tree Construction
+; GCN-O1-NEXT:        SI Fix SGPR Liveness
 ; GCN-O1-NEXT:        SI Fix SGPR copies
 ; GCN-O1-NEXT:        MachinePostDominator Tree Construction
 ; GCN-O1-NEXT:        SI Lower i1 Copies
 ; GCN-O1-NEXT:        Finalize ISel and expand pseudo-instructions
 ; GCN-O1-NEXT:        Lazy Machine Block Frequency Analysis
 ; GCN-O1-NEXT:        Early Tail Duplication
 ; GCN-O1-NEXT:        Optimize machine instruction PHIs
 ; GCN-O1-NEXT:        Slot index numbering

 ; GCN-O1-OPTS-NEXT:        Legacy Divergence Analysis
 ; GCN-O1-OPTS-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O1-OPTS-NEXT:        Function Alias Analysis Results
 ; GCN-O1-OPTS-NEXT:        Branch Probability Analysis
 ; GCN-O1-OPTS-NEXT:        Lazy Branch Probability Analysis
 ; GCN-O1-OPTS-NEXT:        Lazy Block Frequency Analysis
 ; GCN-O1-OPTS-NEXT:        AMDGPU DAG->DAG Pattern Instruction Selection
 ; GCN-O1-OPTS-NEXT:        MachineDominator Tree Construction
+; GCN-O1-OPTS-NEXT:        SI Fix SGPR Liveness
 ; GCN-O1-OPTS-NEXT:        SI Fix SGPR copies
 ; GCN-O1-OPTS-NEXT:        MachinePostDominator Tree Construction
 ; GCN-O1-OPTS-NEXT:        SI Lower i1 Copies
 ; GCN-O1-OPTS-NEXT:        Finalize ISel and expand pseudo-instructions
 ; GCN-O1-OPTS-NEXT:        Lazy Machine Block Frequency Analysis
 ; GCN-O1-OPTS-NEXT:        Early Tail Duplication
 ; GCN-O1-OPTS-NEXT:        Optimize machine instruction PHIs
 ; GCN-O1-OPTS-NEXT:        Slot index numbering

 ; GCN-O2-NEXT:        Legacy Divergence Analysis
 ; GCN-O2-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O2-NEXT:        Function Alias Analysis Results
 ; GCN-O2-NEXT:        Branch Probability Analysis
 ; GCN-O2-NEXT:        Lazy Branch Probability Analysis
 ; GCN-O2-NEXT:        Lazy Block Frequency Analysis
 ; GCN-O2-NEXT:        AMDGPU DAG->DAG Pattern Instruction Selection
 ; GCN-O2-NEXT:        MachineDominator Tree Construction
+; GCN-O2-NEXT:        SI Fix SGPR Liveness
 ; GCN-O2-NEXT:        SI Fix SGPR copies
 ; GCN-O2-NEXT:        MachinePostDominator Tree Construction
 ; GCN-O2-NEXT:        SI Lower i1 Copies
 ; GCN-O2-NEXT:        Finalize ISel and expand pseudo-instructions
 ; GCN-O2-NEXT:        Lazy Machine Block Frequency Analysis
 ; GCN-O2-NEXT:        Early Tail Duplication
 ; GCN-O2-NEXT:        Optimize machine instruction PHIs
 ; GCN-O2-NEXT:        Slot index numbering

 ; GCN-O3-NEXT:        Legacy Divergence Analysis
 ; GCN-O3-NEXT:        Basic Alias Analysis (stateless AA impl)
 ; GCN-O3-NEXT:        Function Alias Analysis Results
 ; GCN-O3-NEXT:        Branch Probability Analysis
 ; GCN-O3-NEXT:        Lazy Branch Probability Analysis
 ; GCN-O3-NEXT:        Lazy Block Frequency Analysis
 ; GCN-O3-NEXT:        AMDGPU DAG->DAG Pattern Instruction Selection
 ; GCN-O3-NEXT:        MachineDominator Tree Construction
+; GCN-O3-NEXT:        SI Fix SGPR Liveness
 ; GCN-O3-NEXT:        SI Fix SGPR copies
 ; GCN-O3-NEXT:        MachinePostDominator Tree Construction
 ; GCN-O3-NEXT:        SI Lower i1 Copies
 ; GCN-O3-NEXT:        Finalize ISel and expand pseudo-instructions
 ; GCN-O3-NEXT:        Lazy Machine Block Frequency Analysis
 ; GCN-O3-NEXT:        Early Tail Duplication
 ; GCN-O3-NEXT:        Optimize machine instruction PHIs
 ; GCN-O3-NEXT:        Slot index numbering

llvm/test/CodeGen/AMDGPU/sgpr-liveness.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -march=amdgcn -mcpu=gfx1010 -verify-machineinstrs < %s | FileCheck %s -check-prefix=GCN
+;
+; This is a test to show that we need to extend the liveness of SGPR register to
+; stop being overwritten. The example IR has a divergent if-then, and the phi %c2
+; only has meaningful incoming value from %entry. To achieve the goal of keeping
+; more values as uniform, the divergence/uniform analysis will mark %c2 as uniform.
+; For AMDGPU backend, we will assign SGPR for %c2 and its incoming values. Without
+; making the value alive in block %if, it might be overwritten in block %if.
+; Usually this issued does not happend because middle-end optimization tends to
+; simplify %c2 as %c. But this case happens after subsequent structurizer change.
+;
+define amdgpu_ps float @sgpr_live_through_all_paths(float inreg %c, float %v, i32 %x, i32 %y) #0 {
+; GCN-LABEL: sgpr_live_through_all_paths:
+; GCN:       ; %bb.0: ; %entry
+; GCN-NEXT:    v_cmp_lt_i32_e64 s2, v2, v1
+; GCN-NEXT:    s_mov_b32 s1, exec_lo
+; GCN-NEXT:    s_and_b32 s2, s1, s2
+; GCN-NEXT:    s_mov_b32 exec_lo, s2
+; GCN-NEXT:    s_cbranch_execz .LBB0_2
+; GCN-NEXT:  ; %bb.1: ; %if
+; GCN-NEXT:    s_mov_b32 s2, 2.0
+; GCN-NEXT:    v_div_scale_f32 v1, s3, s2, s2, v0
+; GCN-NEXT:    v_rcp_f32_e64 v2, v1
+; GCN-NEXT:    s_mov_b32 s3, 1.0
+; GCN-NEXT:    v_fma_f32 v3, -v1, v2, s3
+; GCN-NEXT:    v_fmac_f32_e64 v2, v3, v2
+; GCN-NEXT:    v_div_scale_f32 v3, vcc_lo, v0, s2, v0
+; GCN-NEXT:    v_mul_f32_e64 v4, v3, v2
+; GCN-NEXT:    v_fma_f32 v5, -v1, v4, v3
+; GCN-NEXT:    v_fmac_f32_e64 v4, v5, v2
+; GCN-NEXT:    v_fma_f32 v1, -v1, v4, v3
+; GCN-NEXT:    v_div_fmas_f32 v1, v1, v2, v4
+; GCN-NEXT:    v_div_fixup_f32 v0, v1, s2, v0
+; GCN-NEXT:  .LBB0_2: ; %end
+; GCN-NEXT:    s_or_b32 exec_lo, exec_lo, s1
+; GCN-NEXT:    v_add_f32_e64 v0, v0, s0
+; GCN-NEXT:    ; return to shader part epilog
+entry:
+  %cc = icmp slt i32 %y, %x
+  br i1 %cc, label %if, label %end
+
+if:
+  %v.if = fdiv float %v, 2.0
+  br label %end
+
+end:
+  %v2 = phi float [ %v.if, %if ], [ %v, %entry ]
+  %c2 = phi float [ undef, %if ], [ %c, %entry ]
+  %r = fadd float %v2, %c2
+  ret float %r
+}
+
+attributes #0 = { nounwind optnone noinline }