[AMDGPU][AMDGPUIfConverter] First version of AMDGPUIfConversion

Author: jmmartinez

llvm/lib/Target/AMDGPU/AMDGPUIfConverter.cpp

@@ -11,6 +11,9 @@
 #include <llvm/InitializePasses.h>
 
 #include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "SIInstrInfo.h"
+#include "SIMachineFunctionInfo.h"
 
 using namespace llvm;
 
@@ -19,10 +22,8 @@ namespace {
 const char PassName[] = "AMDGPU if conversion";
 
 class AMDGPUIfConverter : public MachineFunctionPass {
-  const TargetInstrInfo *TII = nullptr;
-  const TargetRegisterInfo *TRI = nullptr;
+  const SIInstrInfo *TII = nullptr;
   TargetSchedModel SchedModel;
-  MachineRegisterInfo *MRI = nullptr;
   MachineDominatorTree *DomTree = nullptr;
   MachineBranchProbabilityInfo *MBPI = nullptr;
   MachineLoopInfo *Loops = nullptr;
@@ -41,6 +42,7 @@ class AMDGPUIfConverter : public MachineFunctionPass {
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   bool tryConvertIf(MachineBasicBlock *);
+  bool shouldConvertIf();
 
   StringRef getPassName() const override { return PassName; }
 };
@@ -60,10 +62,11 @@ bool AMDGPUIfConverter::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 
-  const TargetSubtargetInfo &STI = MF.getSubtarget();
+  const auto &STI = MF.getSubtarget<GCNSubtarget>();
+  if (!STI.hasGFX10_3Insts())
+    return false;
+
   TII = STI.getInstrInfo();
-  TRI = STI.getRegisterInfo();
-  MRI = &MF.getRegInfo();
   SchedModel.init(&STI);
   DomTree = &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
   Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
@@ -79,7 +82,193 @@ bool AMDGPUIfConverter::runOnMachineFunction(MachineFunction &MF) {
   return Changed;
 }
 
-bool AMDGPUIfConverter::tryConvertIf(MachineBasicBlock *MBB) { return false; }
+unsigned getReversedVCMPXOpcode(unsigned Opcode) {
+  // TODO: this is a placeholder for the real function
+  switch (Opcode) {
+  case AMDGPU::V_CMPX_LT_I32_nosdst_e64:
+    return AMDGPU::V_CMPX_GE_I32_nosdst_e64;
+  default:
+    errs() << "unhandled: " << Opcode << "\n";
+    llvm_unreachable("unhandled vcmp opcode");
+  }
+}
+
+bool needsPredication(const SIInstrInfo *TII, const MachineInstr &I) {
+  return TII->isVALU(I) || TII->isVMEM(I);
+}
+
+struct ExecPredicate : ifcvt::PredicationStrategy {
+  const SIInstrInfo *TII;
+  const SIRegisterInfo *RegInfo;
+
+  MachineInstr *Cmp = nullptr;
+
+  ExecPredicate(const SIInstrInfo *TII)
+      : TII(TII), RegInfo(&TII->getRegisterInfo()) {}
+
+  bool canConvertIf(MachineBasicBlock *Head, MachineBasicBlock *TBB,
+                    MachineBasicBlock *FBB, MachineBasicBlock *Tail,
+                    ArrayRef<MachineOperand> Cond) override {
+
+    // check that the cmp is just before the branch and that it is promotable to
+    // v_cmpx
+    const unsigned SupportedBranchOpc[]{
+        AMDGPU::S_CBRANCH_SCC0, AMDGPU::S_CBRANCH_SCC1, AMDGPU::S_CBRANCH_VCCNZ,
+        AMDGPU::S_CBRANCH_VCCZ};
+
+    MachineInstr &CBranch = *Head->getFirstInstrTerminator();
+    if (!llvm::is_contained(SupportedBranchOpc, CBranch.getOpcode()))
+      return false;
+
+    auto CmpInstr = std::next(CBranch.getReverseIterator());
+    if (CmpInstr == Head->instr_rend())
+      return false;
+
+    Register SCCorVCC = Cond[1].getReg();
+    bool ModifiesConditionReg = CmpInstr->modifiesRegister(SCCorVCC, RegInfo);
+    if (!ModifiesConditionReg)
+      return false;
+
+    Cmp = &*CmpInstr;
+
+    unsigned CmpOpc = Cmp->getOpcode();
+    if (TII->isSALU(*Cmp))
+      CmpOpc = TII->getVALUOp(*Cmp);
+    if (AMDGPU::getVCMPXOpFromVCMP(CmpOpc) == -1) {
+      errs() << *Cmp << "\n";
+      return false;
+    }
+
+    auto NeedsPredication = [&](const MachineInstr &I) {
+      return needsPredication(TII, I);
+    };
+    auto BlockNeedsPredication = [&](const MachineBasicBlock *MBB) {
+      if (MBB == Tail)
+        return false;
+      auto Insts = llvm::make_range(MBB->begin(), MBB->getFirstTerminator());
+      return llvm::any_of(Insts, NeedsPredication);
+    };
+
+    MachineBasicBlock *Blocks[] = {TBB, FBB};
+
+    if (llvm::none_of(Blocks, BlockNeedsPredication))
+      return false;
+
+    return true;
+  }
+
+  bool canPredicate(const MachineInstr &I) override {
+
+    // TODO: relax this condition, if exec is masked, check that it goes back to
+    // normal
+    // TODO: what about scc or vcc ? Are they taken into acount in the MBB
+    // live-ins ?
+    MCRegister Exec = RegInfo->getExec();
+    bool ModifiesExec = I.modifiesRegister(Exec, RegInfo);
+    if (ModifiesExec)
+      return false;
+
+    if (needsPredication(TII, I))
+      return true;
+
+    bool DontMoveAcrossStore = true;
+    bool IsSpeculatable = I.isDereferenceableInvariantLoad() ||
+                          I.isSafeToMove(DontMoveAcrossStore);
+    if (IsSpeculatable)
+      return true;
+
+    return false;
+  }
+
+  bool predicateBlock(MachineBasicBlock *MBB, ArrayRef<MachineOperand> Cond,
+                      bool Reverse) override {
+    // save exec
+    MachineFunction &MF = *MBB->getParent();
+    SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
+
+    Register ExecBackup = MFI->getSGPRForEXECCopy();
+
+    const DebugLoc &CmpLoc = Cmp->getDebugLoc();
+
+    auto FirstInstruction = MBB->begin();
+    const bool IsSCCLive =
+        false; // asume not since the live-ins are supposed to be empty
+    TII->insertScratchExecCopy(MF, *MBB, FirstInstruction, CmpLoc, ExecBackup,
+                               IsSCCLive);
+
+    // mask exec
+    unsigned CmpOpc = Cmp->getOpcode();
+    if (TII->isSALU(*Cmp))
+      CmpOpc = TII->getVALUOp(*Cmp);
+
+    CmpOpc = AMDGPU::getVCMPXOpFromVCMP(CmpOpc);
+    if (Reverse)
+      CmpOpc = getReversedVCMPXOpcode(CmpOpc);
+
+    // TODO: handle this properly. The second block may kill those registers.
+    Cmp->getOperand(0).setIsKill(false);
+    Cmp->getOperand(1).setIsKill(false);
+
+    auto VCmpX = BuildMI(*MBB, FirstInstruction, CmpLoc, TII->get(CmpOpc));
+    VCmpX->addOperand(Cmp->getOperand(0));
+    VCmpX->addOperand(Cmp->getOperand(1));
+
+    // restore exec
+    TII->restoreExec(MF, *MBB, MBB->end(), DebugLoc(), ExecBackup);
+
+    return true;
+  }
+
+  ~ExecPredicate() override = default;
+};
+
+/// Update the dominator tree after if-conversion erased some blocks.
+void updateDomTree(MachineDominatorTree *DomTree, const SSAIfConv &IfConv,
+                   ArrayRef<MachineBasicBlock *> Removed) {
+  // convertIf can remove TBB, FBB, and Tail can be merged into Head.
+  // TBB and FBB should not dominate any blocks.
+  // Tail children should be transferred to Head.
+  MachineDomTreeNode *HeadNode = DomTree->getNode(IfConv.Head);
+  for (auto *B : Removed) {
+    MachineDomTreeNode *Node = DomTree->getNode(B);
+    assert(Node != HeadNode && "Cannot erase the head node");
+    while (Node->getNumChildren()) {
+      assert(Node->getBlock() == IfConv.Tail && "Unexpected children");
+      DomTree->changeImmediateDominator(Node->back(), HeadNode);
+    }
+    DomTree->eraseNode(B);
+  }
+}
+
+/// Update LoopInfo after if-conversion.
+void updateLoops(MachineLoopInfo *Loops,
+                 ArrayRef<MachineBasicBlock *> Removed) {
+  // If-conversion doesn't change loop structure, and it doesn't mess with back
+  // edges, so updating LoopInfo is simply removing the dead blocks.
+  for (auto *B : Removed)
+    Loops->removeBlock(B);
+}
+
+bool AMDGPUIfConverter::shouldConvertIf() {
+  // TODO: cost model
+  return true;
+}
+
+bool AMDGPUIfConverter::tryConvertIf(MachineBasicBlock *MBB) {
+  ExecPredicate Predicate{TII};
+  bool Changed = false;
+  while (IfConv.canConvertIf(MBB, Predicate) && shouldConvertIf()) {
+    // If-convert MBB and update analyses.
+    SmallVector<MachineBasicBlock *, 4> RemoveBlocks;
+    IfConv.convertIf(RemoveBlocks, Predicate);
+    Changed = true;
+    updateDomTree(DomTree, IfConv, RemoveBlocks);
+    for (MachineBasicBlock *MBB : RemoveBlocks)
+      MBB->eraseFromParent();
+    updateLoops(Loops, RemoveBlocks);
+  }
+  return Changed;
+}
 
 } // namespace
 

llvm/test/CodeGen/AMDGPU/amdgpu-if-cvt.ll

@@ -5,10 +5,9 @@ define amdgpu_kernel void @scalar_cmp(i32 noundef %value, ptr addrspace(8) nocap
 ; GCN-LABEL: scalar_cmp:
 ; GCN:       ; %bb.0: ; %entry
 ; GCN-NEXT:    s_load_dword s0, s[2:3], 0x4c
+; GCN-NEXT:    s_or_saveexec_b32 s105, -1
 ; GCN-NEXT:    s_waitcnt lgkmcnt(0)
-; GCN-NEXT:    s_cmp_lt_i32 s0, 1
-; GCN-NEXT:    s_cbranch_scc1 .LBB0_2
-; GCN-NEXT:  ; %bb.1: ; %if.then
+; GCN-NEXT:    v_cmpx_ge_i32_e64 s0, 1
 ; GCN-NEXT:    s_clause 0x2
 ; GCN-NEXT:    s_load_dword s4, s[2:3], 0x24
 ; GCN-NEXT:    s_load_dword s5, s[2:3], 0x44
@@ -17,7 +16,6 @@ define amdgpu_kernel void @scalar_cmp(i32 noundef %value, ptr addrspace(8) nocap
 ; GCN-NEXT:    v_mov_b32_e32 v0, s4
 ; GCN-NEXT:    v_mov_b32_e32 v1, s5
 ; GCN-NEXT:    buffer_store_dword v0, v1, s[0:3], 0 offen
-; GCN-NEXT:  .LBB0_2: ; %if.end
 ; GCN-NEXT:    s_endpgm
 entry:
   %cmp = icmp sgt i32 %flag, 0