[AMDGPU] Eliminate likely-spurious execz checks via intrinsic argument

Currently, we introduce branches to skip conditionally executed
instructions if the EXEC mask is zero and only eliminate them if the
scheduling model says that executing the skipped instructions is cheaper
than taking the branch instruction.

This patch adds a heuristic to SIAnnotateControlFlow to determine if the
lanes of a wavefront are likely to have dynamically varying values for
the branch condition.
This information is passed through new arguments/operands of the
amdgcn.if/else intrinsics and the SI_IF/ELSE pseudo instructions to
SILowerControlFlow, where the execz branch is inserted with
corresponding branch probabilities.
This causes SIPreEmitPeephole to eliminate the corresponding execz
branch if it is legal to do so.

This is an alternative to PR llvm#117567, using a simpler heuristic and
passing the LikelyVarying information through new arguments for the
amdgcn.if/else intrinsics and the SI_IF/ELSE pseudo instructions instead
of abusing branch weight metadata.

Most test changes are caused by the new arguments for the amdgcn.if/else
intrinsics and the SI_IF/ELSE pseudo instructions; the LikelyVarying
argument is set to false/0 in these existing tests.
New tests for the functionality are in
conditional-mem-no-cbranch-execz.ll and
annotate-likely-varying-branches.ll.

For SWDEV-483228.

Author: ritter-x2a

llvm/include/llvm/IR/IntrinsicsAMDGPU.td

@@ -3407,11 +3407,11 @@ def int_amdgcn_ashr_pk_u8_i32 : ClangBuiltin<"__builtin_amdgcn_ashr_pk_u8_i32">,
 // having side effects, which is sufficient to prevent optimizations without
 // having to mark them as convergent.
 def int_amdgcn_if : Intrinsic<[llvm_i1_ty, llvm_anyint_ty],
-  [llvm_i1_ty], [IntrWillReturn, IntrNoCallback, IntrNoFree]
+  [llvm_i1_ty, llvm_i1_ty], [ImmArg<ArgIndex<1>>, IntrWillReturn, IntrNoCallback, IntrNoFree]
 >;
 
 def int_amdgcn_else : Intrinsic<[llvm_i1_ty, llvm_anyint_ty],
-  [llvm_anyint_ty], [IntrWillReturn, IntrNoCallback, IntrNoFree]
+  [llvm_anyint_ty, llvm_i1_ty], [ImmArg<ArgIndex<1>>, IntrWillReturn, IntrNoCallback, IntrNoFree]
 >;
 
 def int_amdgcn_if_break : Intrinsic<[llvm_anyint_ty],

llvm/lib/Target/AMDGPU/AMDGPUInstrInfo.td

@@ -42,12 +42,12 @@ def AMDGPUFmasOp : SDTypeProfile<1, 4,
 def ImmOp : SDTypeProfile<0, 1, [SDTCisInt<0>]>;
 def AMDGPUKillSDT : SDTypeProfile<0, 1, [SDTCisInt<0>]>;
 
-def AMDGPUIfOp : SDTypeProfile<1, 2,
-  [SDTCisVT<0, i1>, SDTCisVT<1, i1>, SDTCisVT<2, OtherVT>]
+def AMDGPUIfOp : SDTypeProfile<1, 3,
+  [SDTCisVT<0, i1>, SDTCisVT<1, i1>, SDTCisVT<2, i1>, SDTCisVT<3, OtherVT>]
 >;
 
-def AMDGPUElseOp : SDTypeProfile<1, 2,
-  [SDTCisVT<0, i1>, SDTCisVT<1, i1>, SDTCisVT<2, OtherVT>]
+def AMDGPUElseOp : SDTypeProfile<1, 3,
+  [SDTCisVT<0, i1>, SDTCisVT<1, i1>, SDTCisVT<2, i1>, SDTCisVT<3, OtherVT>]
 >;
 
 def AMDGPULoopOp : SDTypeProfile<0, 2,

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

@@ -7264,6 +7264,7 @@ bool AMDGPULegalizerInfo::legalizeIntrinsic(LegalizerHelper &Helper,
 
       Register Def = MI.getOperand(1).getReg();
       Register Use = MI.getOperand(3).getReg();
+      auto LikelyVarying = MI.getOperand(4).getImm();
 
       MachineBasicBlock *CondBrTarget = BrCond->getOperand(1).getMBB();
 
@@ -7275,11 +7276,13 @@ bool AMDGPULegalizerInfo::legalizeIntrinsic(LegalizerHelper &Helper,
         B.buildInstr(AMDGPU::SI_IF)
           .addDef(Def)
           .addUse(Use)
+          .addImm(LikelyVarying)
           .addMBB(UncondBrTarget);
       } else {
         B.buildInstr(AMDGPU::SI_ELSE)
             .addDef(Def)
             .addUse(Use)
+            .addImm(LikelyVarying)
             .addMBB(UncondBrTarget);
       }
 

llvm/lib/Target/AMDGPU/SIAnnotateControlFlow.cpp

@@ -22,6 +22,7 @@
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
+#include "llvm/IR/IntrinsicsR600.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
@@ -36,6 +37,24 @@ namespace {
 using StackEntry = std::pair<BasicBlock *, Value *>;
 using StackVector = SmallVector<StackEntry, 16>;
 
+class LikelyVaryingHeuristic {
+public:
+  LikelyVaryingHeuristic(const Function &F, const GCNSubtarget &ST) {
+    IsSingleLaneExecution = ST.isSingleLaneExecution(F);
+  }
+
+  /// Check if \p V is likely to be have dynamically varying values among the
+  /// workitems in each wavefront.
+  bool isLikelyVarying(const Value *V);
+
+private:
+  bool IsSingleLaneExecution = false;
+
+  bool isRelevantSourceOfDivergence(const Value *V) const;
+
+  ValueMap<const Value *, bool> LikelyVaryingCache;
+};
+
 class SIAnnotateControlFlow {
 private:
   Function *F;
@@ -62,6 +81,8 @@ class SIAnnotateControlFlow {
 
   LoopInfo *LI;
 
+  LikelyVaryingHeuristic LVHeuristic;
+
   void initialize(const GCNSubtarget &ST);
 
   bool isUniform(BranchInst *T);
@@ -99,7 +120,7 @@ class SIAnnotateControlFlow {
 public:
   SIAnnotateControlFlow(Function &F, const GCNSubtarget &ST, DominatorTree &DT,
                         LoopInfo &LI, UniformityInfo &UA)
-      : F(&F), UA(&UA), DT(&DT), LI(&LI) {
+      : F(&F), UA(&UA), DT(&DT), LI(&LI), LVHeuristic(F, ST) {
     initialize(ST);
   }
 
@@ -186,9 +207,14 @@ bool SIAnnotateControlFlow::openIf(BranchInst *Term) {
   if (isUniform(Term))
     return false;
 
+  // Check if it's likely that at least one lane will always follow the
+  // then-branch, i.e., the then-branch is never skipped completly.
+  Value *IsLikelyVarying =
+      LVHeuristic.isLikelyVarying(Term->getCondition()) ? BoolTrue : BoolFalse;
+
   IRBuilder<> IRB(Term);
   Value *IfCall = IRB.CreateCall(getDecl(If, Intrinsic::amdgcn_if, IntMask),
-                                 {Term->getCondition()});
+                                 {Term->getCondition(), IsLikelyVarying});
   Value *Cond = IRB.CreateExtractValue(IfCall, {0});
   Value *Mask = IRB.CreateExtractValue(IfCall, {1});
   Term->setCondition(Cond);
@@ -202,9 +228,16 @@ bool SIAnnotateControlFlow::insertElse(BranchInst *Term) {
     return false;
   }
 
+  Value *IncomingMask = popSaved();
+  // Check if it's likely that at least one lane will always follow the
+  // else-branch, i.e., the else-branch is never skipped completly.
+  Value *IsLikelyVarying =
+      LVHeuristic.isLikelyVarying(IncomingMask) ? BoolTrue : BoolFalse;
+
   IRBuilder<> IRB(Term);
-  Value *ElseCall = IRB.CreateCall(
-      getDecl(Else, Intrinsic::amdgcn_else, {IntMask, IntMask}), {popSaved()});
+  Value *ElseCall =
+      IRB.CreateCall(getDecl(Else, Intrinsic::amdgcn_else, {IntMask, IntMask}),
+                     {IncomingMask, IsLikelyVarying});
   Value *Cond = IRB.CreateExtractValue(ElseCall, {0});
   Value *Mask = IRB.CreateExtractValue(ElseCall, {1});
   Term->setCondition(Cond);
@@ -385,6 +418,62 @@ bool SIAnnotateControlFlow::run() {
   return Changed;
 }
 
+bool LikelyVaryingHeuristic::isRelevantSourceOfDivergence(
+    const Value *V) const {
+  auto *II = dyn_cast<IntrinsicInst>(V);
+  if (!II)
+    return false;
+
+  switch (II->getIntrinsicID()) {
+  case Intrinsic::amdgcn_workitem_id_z:
+  case Intrinsic::r600_read_tidig_z:
+  case Intrinsic::amdgcn_workitem_id_y:
+  case Intrinsic::r600_read_tidig_y:
+  case Intrinsic::amdgcn_workitem_id_x:
+  case Intrinsic::r600_read_tidig_x:
+  case Intrinsic::amdgcn_mbcnt_hi:
+  case Intrinsic::amdgcn_mbcnt_lo:
+    return true;
+  default:
+    return false;
+  }
+}
+
+bool LikelyVaryingHeuristic::isLikelyVarying(const Value *V) {
+  if (IsSingleLaneExecution)
+    return false;
+
+  if (isRelevantSourceOfDivergence(V))
+    return true;
+
+  auto *I = dyn_cast<Instruction>(V);
+  if (!I)
+    return false;
+
+  // ExtractValueInst and IntrinsicInst enable looking through the
+  // amdgcn_if/else intrinsics inserted by SIAnnotateControlFlow.
+  // This condition excludes PHINodes, which prevents infinite recursion.
+  if (!isa<BinaryOperator>(I) && !isa<UnaryOperator>(I) && !isa<CastInst>(I) &&
+      !isa<CmpInst>(I) && !isa<ExtractValueInst>(I) && !isa<IntrinsicInst>(I))
+    return false;
+
+  // Have we already checked V?
+  auto CacheEntry = LikelyVaryingCache.find(V);
+  if (CacheEntry != LikelyVaryingCache.end())
+    return CacheEntry->second;
+
+  // Does it use a likely varying Value?
+  bool Result = false;
+  for (const auto &Use : I->operands()) {
+    Result |= isLikelyVarying(Use);
+    if (Result)
+      break;
+  }
+
+  LikelyVaryingCache.insert({V, Result});
+  return Result;
+}
+
 PreservedAnalyses SIAnnotateControlFlowPass::run(Function &F,
                                                  FunctionAnalysisManager &FAM) {
   const GCNSubtarget &ST = TM.getSubtarget<GCNSubtarget>(F);

llvm/lib/Target/AMDGPU/SIInstructions.td

@@ -417,8 +417,8 @@ def IGLP_OPT : SPseudoInstSI<(outs), (ins i32imm:$mask),
 let isTerminator = 1, isNotDuplicable = 1 in {
 
 def SI_IF: CFPseudoInstSI <
-  (outs SReg_1:$dst), (ins SReg_1:$vcc, brtarget:$target),
-  [(set i1:$dst, (AMDGPUif i1:$vcc, bb:$target))], 1, 1> {
+  (outs SReg_1:$dst), (ins SReg_1:$vcc, i1imm:$likelyvarying, brtarget:$target),
+  [(set i1:$dst, (AMDGPUif i1:$vcc, (i1 timm:$likelyvarying), bb:$target))], 1, 1> {
   let Constraints = "";
   let Size = 12;
   let hasSideEffects = 1;
@@ -427,7 +427,7 @@ def SI_IF: CFPseudoInstSI <
 
 def SI_ELSE : CFPseudoInstSI <
   (outs SReg_1:$dst),
-  (ins SReg_1:$src, brtarget:$target), [], 1, 1> {
+  (ins SReg_1:$src, i1imm:$likelyvarying, brtarget:$target), [], 1, 1> {
   let Size = 12;
   let hasSideEffects = 1;
   let IsNeverUniform = 1;
@@ -1049,8 +1049,8 @@ def : GCNPat<
 >;
 
 def : GCNPat<
-  (AMDGPUelse i1:$src, bb:$target),
-  (SI_ELSE $src, $target)
+  (AMDGPUelse i1:$src, i1:$likelyvarying, bb:$target),
+  (SI_ELSE $src, $likelyvarying, $target)
 >;
 
 def : Pat <

llvm/lib/Target/AMDGPU/SILowerControlFlow.cpp

@@ -55,8 +55,11 @@
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/LiveVariables.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/Support/BranchProbability.h"
 #include "llvm/Target/TargetMachine.h"
 
 using namespace llvm;
@@ -221,9 +224,11 @@ void SILowerControlFlow::emitIf(MachineInstr &MI) {
   MachineOperand& Cond = MI.getOperand(1);
   assert(Cond.getSubReg() == AMDGPU::NoSubRegister);
 
-  MachineOperand &ImpDefSCC = MI.getOperand(4);
+  MachineOperand &ImpDefSCC = MI.getOperand(5);
   assert(ImpDefSCC.getReg() == AMDGPU::SCC && ImpDefSCC.isDef());
 
+  bool LikelyVarying = MI.getOperand(2).getImm();
+
   // If there is only one use of save exec register and that use is SI_END_CF,
   // we can optimize SI_IF by returning the full saved exec mask instead of
   // just cleared bits.
@@ -281,7 +286,17 @@ void SILowerControlFlow::emitIf(MachineInstr &MI) {
   // Insert the S_CBRANCH_EXECZ instruction which will be optimized later
   // during SIPreEmitPeephole.
   MachineInstr *NewBr = BuildMI(MBB, I, DL, TII->get(AMDGPU::S_CBRANCH_EXECZ))
-                            .add(MI.getOperand(2));
+                            .add(MI.getOperand(3));
+
+  if (LikelyVarying) {
+    MachineBasicBlock *ExeczDest = MI.getOperand(3).getMBB();
+    auto **E = MBB.succ_end();
+    for (auto **SI = MBB.succ_begin(); SI != E; ++SI) {
+      if (*SI == ExeczDest)
+        MBB.setSuccProbability(SI, BranchProbability::getZero());
+    }
+    MBB.normalizeSuccProbs();
+  }
 
   if (!LIS) {
     MI.eraseFromParent();
@@ -329,7 +344,9 @@ void SILowerControlFlow::emitElse(MachineInstr &MI) {
   if (LV)
     LV->replaceKillInstruction(SrcReg, MI, *OrSaveExec);
 
-  MachineBasicBlock *DestBB = MI.getOperand(2).getMBB();
+  bool LikelyVarying = MI.getOperand(2).getImm();
+
+  MachineBasicBlock *DestBB = MI.getOperand(3).getMBB();
 
   MachineBasicBlock::iterator ElsePt(MI);
 
@@ -352,6 +369,15 @@ void SILowerControlFlow::emitElse(MachineInstr &MI) {
       BuildMI(MBB, ElsePt, DL, TII->get(AMDGPU::S_CBRANCH_EXECZ))
           .addMBB(DestBB);
 
+  if (LikelyVarying) {
+    auto **E = MBB.succ_end();
+    for (auto **SI = MBB.succ_begin(); SI != E; ++SI) {
+      if (*SI == DestBB)
+        MBB.setSuccProbability(SI, BranchProbability::getZero());
+    }
+    MBB.normalizeSuccProbs();
+  }
+
   if (!LIS) {
     MI.eraseFromParent();
     return;

llvm/lib/Target/AMDGPU/SIOptimizeVGPRLiveRange.cpp

@@ -184,7 +184,7 @@ MachineBasicBlock *
 SIOptimizeVGPRLiveRange::getElseTarget(MachineBasicBlock *MBB) const {
   for (auto &BR : MBB->terminators()) {
     if (BR.getOpcode() == AMDGPU::SI_ELSE)
-      return BR.getOperand(2).getMBB();
+      return BR.getOperand(3).getMBB();
   }
   return nullptr;
 }
@@ -682,7 +682,7 @@ bool SIOptimizeVGPRLiveRange::run(MachineFunction &MF) {
     for (auto &MI : MBB.terminators()) {
       // Detect the if-else blocks
       if (MI.getOpcode() == AMDGPU::SI_IF) {
-        MachineBasicBlock *IfTarget = MI.getOperand(2).getMBB();
+        MachineBasicBlock *IfTarget = MI.getOperand(3).getMBB();
         auto *Endif = getElseTarget(IfTarget);
         if (!Endif)
           continue;

llvm/lib/Target/AMDGPU/SIPreEmitPeephole.cpp

@@ -14,6 +14,7 @@
 #include "AMDGPU.h"
 #include "GCNSubtarget.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/TargetSchedule.h"
 #include "llvm/Support/BranchProbability.h"
@@ -22,6 +23,8 @@ using namespace llvm;
 
 #define DEBUG_TYPE "si-pre-emit-peephole"
 
+STATISTIC(NumCBranchExeczElim, "Number of s_cbranch_execz eliminated.");
+
 namespace {
 
 class SIPreEmitPeephole : public MachineFunctionPass {
@@ -404,6 +407,7 @@ bool SIPreEmitPeephole::removeExeczBranch(MachineInstr &MI,
     return false;
 
   LLVM_DEBUG(dbgs() << "Removing the execz branch: " << MI);
+  ++NumCBranchExeczElim;
   MI.eraseFromParent();
   SrcMBB.removeSuccessor(TrueMBB);
 

llvm/test/Analysis/DotMachineCFG/AMDGPU/irreducible.mir

@@ -14,10 +14,10 @@
 # MCFG-NEXT: Node{{[0-9A-Za-z]*}} -> Node{{[0-9A-Za-z]*}};
 # MCFG-NEXT: Node{{[0-9A-Za-z]*}} [shape=record,label="{bb.3:\l|\l successors: %bb.4(0x80000000)\l\l  %4:vgpr_32 = PHI %5:vgpr_32, %bb.1, %7:vgpr_32, %bb.2\l}"];
 # MCFG-NEXT: Node{{[0-9A-Za-z]*}} -> Node{{[0-9A-Za-z]*}};
-# MCFG-NEXT: Node{{[0-9A-Za-z]*}} [shape=record,label="{bb.4:\l|\l successors: %bb.2(0x40000000), %bb.5(0x40000000)\l\l  %8:vgpr_32 = V_AND_B32_e32 3, %1:vgpr_32, implicit $exec\l  %9:sreg_64 = V_CMP_EQ_U32_e64 %8:vgpr_32, 2, implicit $exec\l  %10:sreg_64 = SI_IF killed %9:sreg_64, %bb.2, implicit-def dead $exec,\l... implicit-def dead $scc, implicit $exec\l}"];
+# MCFG-NEXT: Node{{[0-9A-Za-z]*}} [shape=record,label="{bb.4:\l|\l successors: %bb.2(0x40000000), %bb.5(0x40000000)\l\l  %8:vgpr_32 = V_AND_B32_e32 3, %1:vgpr_32, implicit $exec\l  %9:sreg_64 = V_CMP_EQ_U32_e64 %8:vgpr_32, 2, implicit $exec\l  %10:sreg_64 = SI_IF killed %9:sreg_64, 0, %bb.2, implicit-def dead $exec,\l... implicit-def dead $scc, implicit $exec\l}"];
 # MCFG-NEXT: Node{{[0-9A-Za-z]*}} -> Node{{[0-9A-Za-z]*}};
 # MCFG-NEXT: Node{{[0-9A-Za-z]*}} -> Node{{[0-9A-Za-z]*}};
-# MCFG-NEXT: Node{{[0-9A-Za-z]*}} [shape=record,label="{bb.5:\l|\l successors: %bb.1(0x40000000), %bb.6(0x40000000)\l\l  %11:sreg_64 = V_CMP_EQ_U32_e64 %8:vgpr_32, 1, implicit $exec\l  %12:sreg_64 = SI_IF killed %11:sreg_64, %bb.1, implicit-def dead $exec,\l... implicit-def dead $scc, implicit $exec\l}"];
+# MCFG-NEXT: Node{{[0-9A-Za-z]*}} [shape=record,label="{bb.5:\l|\l successors: %bb.1(0x40000000), %bb.6(0x40000000)\l\l  %11:sreg_64 = V_CMP_EQ_U32_e64 %8:vgpr_32, 1, implicit $exec\l  %12:sreg_64 = SI_IF killed %11:sreg_64, 0, %bb.1, implicit-def dead $exec,\l... implicit-def dead $scc, implicit $exec\l}"];
 # MCFG-NEXT: Node{{[0-9A-Za-z]*}} -> Node{{[0-9A-Za-z]*}};
 # MCFG-NEXT: Node{{[0-9A-Za-z]*}} -> Node{{[0-9A-Za-z]*}};
 # MCFG-NEXT: Node{{[0-9A-Za-z]*}} [shape=record,label="{bb.6:\l|\l\l S_ENDPGM 0\l}"];
@@ -74,12 +74,12 @@ body:             |
 
     %50:vgpr_32 = V_AND_B32_e32 3, %2, implicit $exec
     %51:sreg_64 = V_CMP_EQ_U32_e64 %50, 2, implicit $exec
-    %52:sreg_64 = SI_IF killed %51:sreg_64, %bb.2, implicit-def dead $exec, implicit-def dead $scc, implicit $exec
+    %52:sreg_64 = SI_IF killed %51:sreg_64, 0, %bb.2, implicit-def dead $exec, implicit-def dead $scc, implicit $exec
 
   bb.5:
     successors: %bb.1, %bb.6
     %61:sreg_64 = V_CMP_EQ_U32_e64 %50, 1, implicit $exec
-    %62:sreg_64 = SI_IF killed %61:sreg_64, %bb.1, implicit-def dead $exec, implicit-def dead $scc, implicit $exec
+    %62:sreg_64 = SI_IF killed %61:sreg_64, 0, %bb.1, implicit-def dead $exec, implicit-def dead $scc, implicit $exec
 
   bb.6:
     S_ENDPGM 0

llvm/test/Analysis/UniformityAnalysis/AMDGPU/MIR/control-flow-intrinsics.mir

@@ -10,7 +10,7 @@ body:             |
     ; CHECK-NOT: DIVERGENT: %1
     %1:sreg_64(s64) = G_IMPLICIT_DEF
     ; CHECK: DIVERGENT: {{.*}} SI_IF
-    %2:sreg_64 = SI_IF %1, %bb.2, implicit-def $exec, implicit-def $scc, implicit $exec
+    %2:sreg_64 = SI_IF %1, 0, %bb.2, implicit-def $exec, implicit-def $scc, implicit $exec
 
   bb.1:
     SI_RETURN
@@ -30,7 +30,7 @@ body:             |
     ; CHECK-NOT: DIVERGENT: %1
     %1:sreg_64(s64) = G_IMPLICIT_DEF
     ; CHECK: DIVERGENT: {{.*}} SI_ELSE
-    %2:sreg_64 = SI_ELSE %1, %bb.2, implicit-def $exec, implicit-def $scc, implicit $exec
+    %2:sreg_64 = SI_ELSE %1, 0, %bb.2, implicit-def $exec, implicit-def $scc, implicit $exec
 
   bb.1:
     SI_RETURN