[llvm] [SelectionDAGBuilder] Fix non-determanism in `shouldKeepJumpConditionsTogether` (PR #83687)

[via llvm-commits]

https://github.com/goldsteinn updated https://github.com/llvm/llvm-project/pull/83687

>From afc5fbe74695cc3dae474855379bccaafafecae9 Mon Sep 17 00:00:00 2001
From: Noah Goldstein <goldstein.w.n at gmail.com>
Date: Sat, 2 Mar 2024 12:05:05 -0600
Subject: [PATCH] [SelectionDAGBuilder] Fix non-determanism in
 `shouldKeepJumpConditionsTogether`

The issue was we where iteration on `SmallPtrSet` order which is based
on runtime address and thus can vary with the same input. Fix by just
using `SmallMapVector` with dummy value.
---
 .../SelectionDAG/SelectionDAGBuilder.cpp      | 27 ++++++++++---------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index 4f6263cc492fe3..fb057f95ca5be8 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -2450,11 +2450,10 @@ SelectionDAGBuilder::EmitBranchForMergedCondition(const Value *Cond,
 
 // Collect dependencies on V recursively. This is used for the cost analysis in
 // `shouldKeepJumpConditionsTogether`.
-static bool
-collectInstructionDeps(SmallPtrSet<const Instruction *, 8> *Deps,
-                       const Value *V,
-                       SmallPtrSet<const Instruction *, 8> *Necessary = nullptr,
-                       unsigned Depth = 0) {
+static bool collectInstructionDeps(
+    SmallMapVector<const Instruction *, bool, 8> *Deps, const Value *V,
+    SmallMapVector<const Instruction *, bool, 8> *Necessary = nullptr,
+    unsigned Depth = 0) {
   // Return false if we have an incomplete count.
   if (Depth >= SelectionDAG::MaxRecursionDepth)
     return false;
@@ -2471,7 +2470,7 @@ collectInstructionDeps(SmallPtrSet<const Instruction *, 8> *Deps,
   }
 
   // Already added this dep.
-  if (!Deps->insert(I).second)
+  if (!Deps->try_emplace(I, false).second)
     return true;
 
   for (unsigned OpIdx = 0, E = I->getNumOperands(); OpIdx < E; ++OpIdx)
@@ -2526,7 +2525,9 @@ bool SelectionDAGBuilder::shouldKeepJumpConditionsTogether(
     return false;
 
   // Collect "all" instructions that lhs condition is dependent on.
-  SmallPtrSet<const Instruction *, 8> LhsDeps, RhsDeps;
+  // Use map for stable iteration (to avoid non-determanism of iteration of
+  // SmallPtrSet). The `bool` value is just a dummy.
+  SmallMapVector<const Instruction *, bool, 8> LhsDeps, RhsDeps;
   collectInstructionDeps(&LhsDeps, Lhs);
   // Collect "all" instructions that rhs condition is dependent on AND are
   // dependencies of lhs. This gives us an estimate on which instructions we
@@ -2536,7 +2537,7 @@ bool SelectionDAGBuilder::shouldKeepJumpConditionsTogether(
   // Add the compare instruction itself unless its a dependency on the LHS.
   if (const auto *RhsI = dyn_cast<Instruction>(Rhs))
     if (!LhsDeps.contains(RhsI))
-      RhsDeps.insert(RhsI);
+      RhsDeps.try_emplace(RhsI, false);
 
   const auto &TLI = DAG.getTargetLoweringInfo();
   const auto &TTI =
@@ -2564,9 +2565,9 @@ bool SelectionDAGBuilder::shouldKeepJumpConditionsTogether(
   // instructions.
   for (unsigned PruneIters = 0; PruneIters < MaxPruneIters; ++PruneIters) {
     const Instruction *ToDrop = nullptr;
-    for (const auto *Ins : RhsDeps) {
-      if (!ShouldCountInsn(Ins)) {
-        ToDrop = Ins;
+    for (const auto &InsPair : RhsDeps) {
+      if (!ShouldCountInsn(InsPair.first)) {
+        ToDrop = InsPair.first;
         break;
       }
     }
@@ -2575,13 +2576,13 @@ bool SelectionDAGBuilder::shouldKeepJumpConditionsTogether(
     RhsDeps.erase(ToDrop);
   }
 
-  for (const auto *Ins : RhsDeps) {
+  for (const auto &InsPair : RhsDeps) {
     // Finally accumulate latency that we can only attribute to computing the
     // RHS condition. Use latency because we are essentially trying to calculate
     // the cost of the dependency chain.
     // Possible TODO: We could try to estimate ILP and make this more precise.
     CostOfIncluding +=
-        TTI.getInstructionCost(Ins, TargetTransformInfo::TCK_Latency);
+        TTI.getInstructionCost(InsPair.first, TargetTransformInfo::TCK_Latency);
 
     if (CostOfIncluding > CostThresh)
       return false;