[InstrProf] Fix BalancedPartitioning when threads are disabled

[ellis]

Summary

In https://reviews.llvm.org/D147812 we introduced the class
BalancedPartitioning which includes some threading code. The tests in
that diff run forever when built with -DLLVM_ENABLE_THREADS=OFF so
some bots were broken.

These tests were skipped in
https://reviews.llvm.org/rGa4845eaf2e9aa18dd900d7cbeff4e5ff52e4b50e
because of this.

This diff disables the threading code if LLVM_ENABLE_THREADS is
disabled so we can re-enable the tests.

[luporl]

LGTM. I was able to run check-all locally without issues with this patch.

Thanks for the quick fix!

[kyulee]

LGTM

llvm/lib/Support/BalancedPartitioning.cpp

 
 void BPFunctionNode::dump(raw_ostream &OS) const {
   OS << formatv("{{ID={0} Utilities={{{1:$[,]}} Bucket={2}}", Id,
                 make_range(UtilityNodes.begin(), UtilityNodes.end()), Bucket);
 }
 
 template <typename Func>
 void BalancedPartitioning::BPThreadPool::async(Func &&F) {
+#if LLVM_ENABLE_THREADS
   // This new thread could spawn more threads, so mark it as active
   ++NumActiveThreads;
   TheThreadPool.async([=]() {
     // Run the task
     F();
 
     // This thread will no longer spawn new threads, so mark it as inactive
     if (--NumActiveThreads == 0) {
       // There are no more active threads, so mark as finished and notify
       {
         std::unique_lock<std::mutex> lock(mtx);
         assert(!IsFinishedSpawning);
         IsFinishedSpawning = true;
       }
       cv.notify_one();
     }
   });
+#else
+  llvm_unreachable("threads are disabled");
+#endif
 }
 
 void BalancedPartitioning::BPThreadPool::wait() {
+#if LLVM_ENABLE_THREADS
   // TODO: We could remove the mutex and condition variable and use
   // std::atomic::wait() instead, but that isn't available until C++20
   {
     std::unique_lock<std::mutex> lock(mtx);
     cv.wait(lock, [&]() { return IsFinishedSpawning; });
     assert(IsFinishedSpawning && NumActiveThreads == 0);
   }
   // Now we can call ThreadPool::wait() since all tasks have been submitted
   TheThreadPool.wait();
+#else
+  llvm_unreachable("threads are disabled");
+#endif
 }
 
 BalancedPartitioning::BalancedPartitioning(
     const BalancedPartitioningConfig &Config)
     : Config(Config) {
   // Pre-computing log2 values
   Log2Cache[0] = 0.0;
   for (unsigned I = 1; I < LOG_CACHE_SIZE; I++)
     Log2Cache[I] = std::log2(I);
 }
 
 void BalancedPartitioning::run(std::vector<BPFunctionNode> &Nodes) const {
   LLVM_DEBUG(
       dbgs() << format(
           "Partitioning %d nodes using depth %d and %d iterations per split\n",
           Nodes.size(), Config.SplitDepth, Config.IterationsPerSplit));
   std::optional<BPThreadPool> TP;
+#if LLVM_ENABLE_THREADS
   if (Config.TaskSplitDepth > 1)
     TP.emplace();
+#endif
 
   // Record the input order
   for (unsigned I = 0; I < Nodes.size(); I++)
     Nodes[I].InputOrderIndex = I;
 
   auto NodesRange = llvm::make_range(Nodes.begin(), Nodes.end());
   auto BisectTask = [=, &TP]() {
     bisect(NodesRange, /*RecDepth=*/0, /*RootBucket=*/1, /*Offset=*/0, TP);

llvm/test/tools/llvm-profdata/show-order.proftext

 # RUN: llvm-profdata order %s | FileCheck %s
-# PR63168: Balanced Partitioning tests currently hang on ARM.
-# UNSUPPORTED: target=arm{{.*}}
 
 # CHECK: a
 # CHECK: b
 # CHECK: c
 # CHECK: x
 
 # Header
 :ir

llvm/unittests/Support/BalancedPartitioningTest.cpp

 #include "gtest/gtest.h"
 
 using testing::Each;
 using testing::Field;
 using testing::Not;
 using testing::UnorderedElementsAre;
 using testing::UnorderedElementsAreArray;
 
-// PR63168: Balanced Partitioning tests currently hang on ARM.
-#if defined(__arm__)
-#define SKIP_UNSUPPORTED_PLATFORM GTEST_SKIP()
-#else
-#define SKIP_UNSUPPORTED_PLATFORM do { } while(0)
-#endif
-
 namespace llvm {
 
 void PrintTo(const BPFunctionNode &Node, std::ostream *OS) {
   raw_os_ostream ROS(*OS);
   Node.dump(ROS);
 }
 
 class BalancedPartitioningTest : public ::testing::Test {
 protected:
   BalancedPartitioningConfig Config;
   BalancedPartitioning Bp;
   BalancedPartitioningTest() : Bp(Config) {}
 
   static std::vector<BPFunctionNode::IDT>
   getIds(std::vector<BPFunctionNode> Nodes) {
     std::vector<BPFunctionNode::IDT> Ids;
     for (auto &N : Nodes)
       Ids.push_back(N.Id);
     return Ids;
   }
 };
 
 TEST_F(BalancedPartitioningTest, Basic) {
-  SKIP_UNSUPPORTED_PLATFORM;
   std::vector<BPFunctionNode> Nodes = {
       BPFunctionNode(0, {1, 2}), BPFunctionNode(2, {3, 4}),
       BPFunctionNode(1, {1, 2}), BPFunctionNode(3, {3, 4}),
       BPFunctionNode(4, {4}),
   };
 
   Bp.run(Nodes);
 
   auto NodeIs = [](BPFunctionNode::IDT Id, std::optional<uint32_t> Bucket) {
     return AllOf(Field("Id", &BPFunctionNode::Id, Id),
                  Field("Bucket", &BPFunctionNode::Bucket, Bucket));
   };
 
   EXPECT_THAT(Nodes,
               UnorderedElementsAre(NodeIs(0, 0), NodeIs(1, 1), NodeIs(2, 2),
                                    NodeIs(3, 3), NodeIs(4, 4)));
 }
 
 TEST_F(BalancedPartitioningTest, Large) {
-  SKIP_UNSUPPORTED_PLATFORM;
   const int ProblemSize = 1000;
   std::vector<BPFunctionNode::UtilityNodeT> AllUNs;
   for (int i = 0; i < ProblemSize; i++)
     AllUNs.emplace_back(i);
 
   std::mt19937 RNG;
   std::vector<BPFunctionNode> Nodes;
   for (int i = 0; i < ProblemSize; i++) {
TEST_F(BalancedPartitioningTest, Large) {
   Bp.run(Nodes);
 
   EXPECT_THAT(
       Nodes, Each(Not(Field("Bucket", &BPFunctionNode::Bucket, std::nullopt))));
   EXPECT_THAT(getIds(Nodes), UnorderedElementsAreArray(OrigIds));
 }
 
 TEST_F(BalancedPartitioningTest, MoveGain) {
-  SKIP_UNSUPPORTED_PLATFORM;
   BalancedPartitioning::SignaturesT Signatures = {
       {10, 10, 10.f, 0.f, true}, // 0
       {10, 10, 0.f, 10.f, true}, // 1
       {10, 10, 0.f, 20.f, true}, // 2
   };
   EXPECT_FLOAT_EQ(Bp.moveGain(BPFunctionNode(0, {}), true, Signatures), 0.f);
   EXPECT_FLOAT_EQ(Bp.moveGain(BPFunctionNode(0, {0, 1}), true, Signatures),
                   10.f);
   EXPECT_FLOAT_EQ(Bp.moveGain(BPFunctionNode(0, {1, 2}), false, Signatures),
                   30.f);
 }
 
 } // end namespace llvm