[mlir] Use thread-pool's notion of thread count instead of requerying system.

[stellaraccident]

Summary

The computed number of hardware threads can change over the life of the process based on affinity changes. Since we need a data structure that is at least as large as the maximum parallelism, it is important to use the value that was actually latched for the thread pool we will be dispatching work to.

Also adds an assert specifically for if it doesn't line up (I was getting a crash on an index into the vector).

[mehdi_amini]

mlir/lib/Transforms/Inliner.cpp

708

[mehdi_amini]

mlir/lib/Transforms/Inliner.cpp

680

Patch LGTM, but I don't quite get this comment right now.

[stellaraccident]

Remove comment.

[mehdi_amini]

mlir/lib/Transforms/Inliner.cpp

680

(I meant that I don't get the existing comment in the codebase)

[stellaraccident]

mlir/lib/Transforms/Inliner.cpp

680

I was trying to indicate that there is an action at a distance constraint, but I couldn't come up with a better way to say it so just removed.

[stellaraccident]

Comment.

[Mogball]

Good catch. I think this is the only use of compute_thread_count in MLIR core.

[stellaraccident]

mlir/lib/Transforms/Inliner.cpp

680

How about:

"We must maintain a fixed pool of pass managers which is at least as large as the maximum parallelism of the failableParallelForEach below."

I don't understand the instrumentation/main thread connection myself.

[stellaraccident]

Comment.

mlir/lib/Transforms/Inliner.cpp

680

I went ahead and kept the Note. I don't know quite what it is trying to convey but seems important to understand in a followup.

[mehdi_amini]

mlir/lib/Transforms/Inliner.cpp

680

Yeah it's the part about "The number of pass managers here needs to remain constant" that puzzles me, since we resize it right below.

708

(did you miss this one?)

[stellaraccident]

mlir/lib/Transforms/Inliner.cpp

708

It shows as changed in my snapshot.

[stellaraccident]

Filed issue: https://github.com/llvm/llvm-project/issues/52862 for clarification. I don't like not understanding action-at-a-distance things like this and maybe the author can recall.

mlir/lib/Transforms/Inliner.cpp

if (!region->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>())
       continue;
     nodesToVisit.push_back(node);
   }
   if (nodesToVisit.empty())
     return success();
 
   // Optimize each of the nodes within the SCC in parallel.
   if (failed(optimizeSCCAsync(nodesToVisit, context)))
-      return failure();
 
   // Recompute the uses held by each of the nodes.
   for (CallGraphNode *node : nodesToVisit)
     useList.recomputeUses(node, cg);
   return success();
 }
 
 LogicalResult
 InlinerPass::optimizeSCCAsync(MutableArrayRef<CallGraphNode *> nodesToVisit,
                               MLIRContext *ctx) {
   // Ensure that there are enough pipeline maps for the optimizer to run in
   // parallel. Note: The number of pass managers here needs to remain constant
   // to prevent issues with pass instrumentations that rely on having the same
   // pass manager for the main thread.
-  size_t numThreads = llvm::hardware_concurrency().compute_thread_count();
+  // Note that this lining up is dependent on failableParallelForEach using
+  // the context thread pool under the covers (the thread count must be
+  // consistent).
+  llvm::ThreadPool &threadPool = ctx->getThreadPool();
+  size_t numThreads = threadPool.getThreadCount();
   if (opPipelines.size() < numThreads) {
     // Reserve before resizing so that we can use a reference to the first
     // element.
     opPipelines.reserve(numThreads);
     opPipelines.resize(numThreads, opPipelines.front());
   }
 
   // Ensure an analysis manager has been constructed for each of the nodes.
   // This prevents thread races when running the nested pipelines.
   for (CallGraphNode *node : nodesToVisit)
     getAnalysisManager().nest(node->getCallableRegion()->getParentOp());
 
   // An atomic failure variable for the async executors.
   std::vector<std::atomic<bool>> activePMs(opPipelines.size());
   std::fill(activePMs.begin(), activePMs.end(), false);
   return failableParallelForEach(ctx, nodesToVisit, [&](CallGraphNode *node) {
     // Find a pass manager for this operation.
     auto it = llvm::find_if(activePMs, [](std::atomic<bool> &isActive) {
       bool expectedInactive = false;
       return isActive.compare_exchange_strong(expectedInactive, true);
     });
+    assert(it != activePMs.end() &&
+           "could not find active pass manager for thread");
     unsigned pmIndex = it - activePMs.begin();
 
     // Optimize this callable node.
     LogicalResult result = optimizeCallable(node, opPipelines[pmIndex]);
 
     // Reset the active bit for this pass manager.
     activePMs[pmIndex].store(false);
     return result;