[mlir] Async: Add numWorkerThreads argument to createAsyncParallelForPass

[ezhulenev]

Summary

Add an option to pass the number of worker threads to select the number of async regions for parallel for transformation.

std::unique_ptr<OperationPass<FuncOp>> createAsyncParallelForPass(int numWorkerThreads);

mlir/include/mlir/Dialect/Async/Passes.h

 #define MLIR_DIALECT_ASYNC_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 std::unique_ptr<OperationPass<FuncOp>> createAsyncParallelForPass();
 
+std::unique_ptr<OperationPass<FuncOp>>
+createAsyncParallelForPass(int numWorkerThreads);
+
 std::unique_ptr<OperationPass<FuncOp>> createAsyncRefCountingPass();
 
 std::unique_ptr<OperationPass<FuncOp>> createAsyncRefCountingOptimizationPass();
 
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Async/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ASYNC_PASSES_H_

mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp

 
 private:
   int numConcurrentAsyncExecute;
 };
 
 struct AsyncParallelForPass
     : public AsyncParallelForBase<AsyncParallelForPass> {
   AsyncParallelForPass() = default;
+  AsyncParallelForPass(int numWorkerThreads) {
+    assert(numWorkerThreads >= 1);
+    numConcurrentAsyncExecute = numWorkerThreads;
+  }
   void runOnFunction() override;
 };
 
 } // namespace
 
 LogicalResult
 AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
                                          PatternRewriter &rewriter) const {
void AsyncParallelForPass::runOnFunction() {
 
   if (failed(applyPatternsAndFoldGreedily(getFunction(), std::move(patterns))))
     signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createAsyncParallelForPass() {
   return std::make_unique<AsyncParallelForPass>();
 }
+
+std::unique_ptr<OperationPass<FuncOp>>
+mlir::createAsyncParallelForPass(int numWorkerThreads) {
+  return std::make_unique<AsyncParallelForPass>(numWorkerThreads);
+}