[llvm] [Uniformity] Avoid marking all PHIs as divergent in join blocks (PR #157808)

[Carl Ritson via llvm-commits]

================
@@ -917,19 +919,67 @@ void GenericUniformityAnalysisImpl<ContextT>::taintAndPushAllDefs(
 /// Mark divergent phi nodes in a join block
 template <typename ContextT>
 void GenericUniformityAnalysisImpl<ContextT>::taintAndPushPhiNodes(
-    const BlockT &JoinBlock) {
+    const BlockT &JoinBlock, const BlockT &DivTermBlock,
+    const DivergenceDescriptorT &DivDesc) {
   LLVM_DEBUG(dbgs() << "taintAndPushPhiNodes in " << Context.print(&JoinBlock)
                     << "\n");
   for (const auto &Phi : JoinBlock.phis()) {
-    // FIXME: The non-undef value is not constant per se; it just happens to be
-    // uniform and may not dominate this PHI. So assuming that the same value
-    // reaches along all incoming edges may itself be undefined behaviour. This
-    // particular interpretation of the undef value was added to
-    // DivergenceAnalysis in the following review:
-    //
-    // https://reviews.llvm.org/D19013
-    if (ContextT::isConstantOrUndefValuePhi(Phi))
+    // Attempt to maintain uniformity for PHIs by considering control
+    // dependencies before marking them.
+    SmallVector<ConstValueRefT> Values;
+    SmallVector<const BlockT *> Blocks;
+    Context.getPhiInputs(Phi, Values, Blocks);
+    assert(Blocks.size() == Values.size());
+
+    // Allow an empty Blocks/Values list to signify getPhiInputs is not
+    // implemented; in which case no uniformity is possible.
+    bool HasSingleValue = !Values.empty();
+    bool UniformOnPath = HasSingleValue;
+
+    std::optional<ConstValueRefT> PhiCommon, PathCommon;
+    for (unsigned I = 0; I < Blocks.size() && (UniformOnPath || HasSingleValue);
+         ++I) {
+      // FIXME: We assume undefs are uniform and/or do not dominate the PHI
+      // in the presence of other constant or uniform values.
+      // This particular interpretation of the undef value was added to
+      // DivergenceAnalysis in the following review:
+      //
+      // https://reviews.llvm.org/D19013
+      if (!Values[I])
+        continue;
+
+      // Track common value for all inputs.
+      if (!PhiCommon)
+        PhiCommon = Values[I];
+      else if (Values[I] != *PhiCommon)
+        HasSingleValue = false;
+
+      // Divergent path does not have uniform value.
+      if (!UniformOnPath)
+        continue;
+
+      // Only consider predecessors on divergent path.
+      if (Blocks[I] != &DivTermBlock &&
+          !DivDesc.BlockLabels.lookup_or(Blocks[I], nullptr))
+        continue;
+
+      // Phi is reached via divergent exit (i.e. respect temporal divergence).
+      if (DivDesc.CycleDivBlocks.contains(Blocks[I])) {
+        UniformOnPath = false;
----------------
perlfu wrote:

> At this point, HasSingleValue might still be true, and then the PHINode will not get marked as divergent.
Yes, this is what the current implementation prior to this patch does.

The key difference between this test and `propagateTemporalDivergence` is whether it is related to the CFG edge or the value definition.
`propagateTemporalDivergence` only marks PHIs if they use a value defined within a loop.
This marks the PHI based on its reachability from a divergent exit.

The test `hidden_doublebreak_diverge` in `hidden_loopdiverge.ll` highlights this.
The PHI `%div.merge.y = phi i32 [ 42, %X ], [ %b, %B ]` only uses values defined outside the function, so the values themselves will not cause temporal divergence.

To borrow Jay's summary:
> Once we go entry->H->G, there is a divergent branch which can choose either G->X->Y or G->L->C->H->B->Y. Therefore the choice of whether we arrive at Y from X or from B is divergent.

The divergent path via B only has a single value (`%b`), so this PHI would not be marked divergent (`UniformOnPath = true`) unless we consider the temporal divergence of the path as well.




https://github.com/llvm/llvm-project/pull/157808