[mlir] Transform scf.parallel to scf.for + async.execute

Depends On D89958

1. Adds `async.group`/`async.awaitall` to group together multiple async tokens/values
2. Rewrite scf.parallel operation into multiple concurrent async.execute operations over non overlapping subranges of the original loop.

Example:

```
   scf.for (%i, %j) = (%lbi, %lbj) to (%ubi, %ubj) step (%si, %sj) {
     "do_some_compute"(%i, %j): () -> ()
   }
```

Converted to:

```
   %c0 = constant 0 : index
   %c1 = constant 1 : index

   // Compute blocks sizes for each induction variable.
   %num_blocks_i = ... : index
   %num_blocks_j = ... : index
   %block_size_i = ... : index
   %block_size_j = ... : index

   // Create an async group to track async execute ops.
   %group = async.create_group

   scf.for %bi = %c0 to %num_blocks_i step %c1 {
     %block_start_i = ... : index
     %block_end_i   = ... : index

     scf.for %bj = %c0 t0 %num_blocks_j step %c1 {
       %block_start_j = ... : index
       %block_end_j   = ... : index

       // Execute the body of original parallel operation for the current
       // block.
       %token = async.execute {
         scf.for %i = %block_start_i to %block_end_i step %si {
           scf.for %j = %block_start_j to %block_end_j step %sj {
             "do_some_compute"(%i, %j): () -> ()
           }
         }
       }

       // Add produced async token to the group.
       async.add_to_group %token, %group
     }
   }

   // Await completion of all async.execute operations.
   async.await_all %group
```
In this example outer loop launches inner block level loops as separate async
execute operations which will be executed concurrently.

At the end it waits for the completiom of all async execute operations.

Reviewed By: ftynse, mehdi_amini

Differential Revision: https://reviews.llvm.org/D89963

Author: ezhulenev

mlir/include/mlir/Dialect/Async/CMakeLists.txt

@@ -1 +1,7 @@
 add_subdirectory(IR)
+
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name Async)
+add_public_tablegen_target(MLIRAsyncPassIncGen)
+
+add_mlir_doc(Passes -gen-pass-doc AsyncPasses ./)

mlir/include/mlir/Dialect/Async/IR/Async.h

@@ -47,6 +47,12 @@ class ValueType
   Type getValueType();
 };
 
+/// The group type to represent async tokens or values grouped together.
+class GroupType : public Type::TypeBase<GroupType, Type, TypeStorage> {
+public:
+  using Base::Base;
+};
+
 } // namespace async
 } // namespace mlir
 

mlir/include/mlir/Dialect/Async/IR/AsyncBase.td

@@ -56,6 +56,16 @@ class Async_ValueType<Type type>
   Type valueType = type;
 }
 
+def Async_GroupType : DialectType<AsyncDialect,
+    CPred<"$_self.isa<::mlir::async::GroupType>()">, "group type">,
+    BuildableType<"$_builder.getType<::mlir::async::GroupType>()"> {
+  let typeDescription = [{
+    `async.group` represent a set of async tokens or values and allows to
+    execute async operations on all of them together (e.g. wait for the
+    completion of all/any of them).
+  }];
+}
+
 def Async_AnyValueType : DialectType<AsyncDialect,
                            CPred<"$_self.isa<::mlir::async::ValueType>()">,
                                  "async value type">;

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

@@ -81,6 +81,20 @@ def Async_ExecuteOp :
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let verifier = [{ return ::verify(*this); }];
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+    OpBuilderDAG<(ins "TypeRange":$resultTypes, "ValueRange":$dependencies,
+      "ValueRange":$operands,
+      CArg<"function_ref<void(OpBuilder &, Location, ValueRange)>",
+           "nullptr">:$bodyBuilder)>,
+  ];
+
+  let extraClassDeclaration = [{
+    using BodyBuilderFn =
+        function_ref<void(OpBuilder &, Location, ValueRange)>;
+
+  }];
 }
 
 def Async_YieldOp :
@@ -93,12 +107,12 @@ def Async_YieldOp :
 
   let arguments = (ins Variadic<AnyType>:$operands);
 
-  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
+  let assemblyFormat = "($operands^ `:` type($operands))? attr-dict";
 
   let verifier = [{ return ::verify(*this); }];
 }
 
-def Async_AwaitOp : Async_Op<"await", [NoSideEffect]> {
+def Async_AwaitOp : Async_Op<"await"> {
   let summary = "waits for the argument to become ready";
   let description = [{
     The `async.await` operation waits until the argument becomes ready, and for
@@ -133,12 +147,84 @@ def Async_AwaitOp : Async_Op<"await", [NoSideEffect]> {
   }];
 
   let assemblyFormat = [{
-    attr-dict $operand `:` custom<AwaitResultType>(
+    $operand `:` custom<AwaitResultType>(
       type($operand), type($result)
-    )
+    ) attr-dict
   }];
 
   let verifier = [{ return ::verify(*this); }];
 }
 
+def Async_CreateGroupOp : Async_Op<"create_group", [NoSideEffect]> {
+  let summary = "creates an empty async group";
+  let description = [{
+    The `async.create_group` allocates an empty async group. Async tokens or
+    values can be added to this group later.
+
+    Example:
+
+    ```mlir
+    %0 = async.create_group
+    ...
+    async.await_all %0
+    ```
+  }];
+
+  let arguments = (ins );
+  let results = (outs Async_GroupType:$result);
+
+  let assemblyFormat = "attr-dict";
+}
+
+def Async_AddToGroupOp : Async_Op<"add_to_group", []> {
+  let summary = "adds and async token or value to the group";
+  let description = [{
+    The `async.add_to_group` adds an async token or value to the async group.
+    Returns the rank of the added element in the group. This rank is fixed
+    for the group lifetime.
+
+    Example:
+
+    ```mlir
+    %0 = async.create_group
+    %1 = ... : !async.token
+    %2 = async.add_to_group %1, %0 : !async.token
+    ```
+  }];
+
+  let arguments = (ins Async_AnyValueOrTokenType:$operand,
+                       Async_GroupType:$group);
+  let results = (outs Index:$rank);
+
+  let assemblyFormat = "$operand `,` $group `:` type($operand) attr-dict";
+}
+
+def Async_AwaitAllOp : Async_Op<"await_all", []> {
+  let summary = "waits for the all async tokens or values in the group to "
+                "become ready";
+  let description = [{
+    The `async.await_all` operation waits until all the tokens or values in the
+    group become ready.
+
+    Example:
+
+    ```mlir
+    %0 = async.create_group
+
+    %1 = ... : !async.token
+    %2 = async.add_to_group %1, %0 : !async.token
+
+    %3 = ... : !async.token
+    %4 = async.add_to_group %2, %0 : !async.token
+
+    async.await_all %0
+    ```
+  }];
+
+  let arguments = (ins Async_GroupType:$operand);
+  let results = (outs);
+
+  let assemblyFormat = "$operand attr-dict";
+}
+
 #endif // ASYNC_OPS

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

@@ -0,0 +1,32 @@
+//===- Passes.h - Async pass entry points -----------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines prototypes that expose pass constructors.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ASYNC_PASSES_H_
+#define MLIR_DIALECT_ASYNC_PASSES_H_
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+std::unique_ptr<OperationPass<FuncOp>> createAsyncParallelForPass();
+
+//===----------------------------------------------------------------------===//
+// 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/include/mlir/Dialect/Async/Passes.td

@@ -0,0 +1,27 @@
+//===-- Passes.td - Async pass definition file -------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ASYNC_PASSES
+#define MLIR_DIALECT_ASYNC_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def AsyncParallelFor : FunctionPass<"async-parallel-for"> {
+  let summary = "Convert scf.parallel operations to multiple async regions "
+                "executed concurrently for non-overlapping iteration ranges";
+  let constructor = "mlir::createAsyncParallelForPass()";
+  let options = [
+    Option<"numConcurrentAsyncExecute", "num-concurrent-async-execute",
+      "int32_t", /*default=*/"4",
+      "The number of async.execute operations that will be used for concurrent "
+      "loop execution.">
+  ];
+  let dependentDialects = ["async::AsyncDialect", "scf::SCFDialect"];
+}
+
+#endif // MLIR_DIALECT_ASYNC_PASSES

mlir/include/mlir/ExecutionEngine/AsyncRuntime.h

@@ -14,6 +14,8 @@
 #ifndef MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_
 #define MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_
 
+#include <stdint.h>
+
 #ifdef _WIN32
 #ifndef MLIR_ASYNCRUNTIME_EXPORT
 #ifdef mlir_async_runtime_EXPORTS
@@ -37,6 +39,9 @@
 // Runtime implementation of `async.token` data type.
 typedef struct AsyncToken MLIR_AsyncToken;
 
+// Runtime implementation of `async.group` data type.
+typedef struct AsyncGroup MLIR_AsyncGroup;
+
 // Async runtime uses LLVM coroutines to represent asynchronous tasks. Task
 // function is a coroutine handle and a resume function that continue coroutine
 // execution from a suspension point.
@@ -46,6 +51,12 @@ using CoroResume = void (*)(void *); // coroutine resume function
 // Create a new `async.token` in not-ready state.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncToken *mlirAsyncRuntimeCreateToken();
 
+// Create a new `async.group` in empty state.
+extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncGroup *mlirAsyncRuntimeCreateGroup();
+
+extern "C" MLIR_ASYNCRUNTIME_EXPORT int64_t
+mlirAsyncRuntimeAddTokenToGroup(AsyncToken *, AsyncGroup *);
+
 // Switches `async.token` to ready state and runs all awaiters.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
 mlirAsyncRuntimeEmplaceToken(AsyncToken *);
@@ -54,6 +65,10 @@ mlirAsyncRuntimeEmplaceToken(AsyncToken *);
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
 mlirAsyncRuntimeAwaitToken(AsyncToken *);
 
+// Blocks the caller thread until the elements in the group become ready.
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *);
+
 // Executes the task (coro handle + resume function) in one of the threads
 // managed by the runtime.
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void mlirAsyncRuntimeExecute(CoroHandle,
@@ -64,6 +79,11 @@ extern "C" MLIR_ASYNCRUNTIME_EXPORT void mlirAsyncRuntimeExecute(CoroHandle,
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void
 mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *, CoroHandle, CoroResume);
 
+// Executes the task (coro handle + resume function) in one of the threads
+// managed by the runtime after the all members of the group become ready.
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *, CoroHandle, CoroResume);
+
 //===----------------------------------------------------------------------===//
 // Small async runtime support library for testing.
 //===----------------------------------------------------------------------===//

mlir/include/mlir/InitAllPasses.h

@@ -16,6 +16,7 @@
 
 #include "mlir/Conversion/Passes.h"
 #include "mlir/Dialect/Affine/Passes.h"
+#include "mlir/Dialect/Async/Passes.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
@@ -47,6 +48,7 @@ inline void registerAllPasses() {
 
   // Dialect passes
   registerAffinePasses();
+  registerAsyncPasses();
   registerGPUPasses();
   registerLinalgPasses();
   LLVM::registerLLVMPasses();

mlir/integration_test/Dialect/Async/CPU/lit.local.cfg

@@ -0,0 +1,5 @@
+import sys
+
+# No JIT on win32.
+if sys.platform == 'win32':
+    config.unsupported = True

mlir/integration_test/Dialect/Async/CPU/test-async-parallel-for-1d.mlir

@@ -0,0 +1,66 @@
+// RUN:   mlir-opt %s -async-parallel-for                                      \
+// RUN:               -convert-async-to-llvm                                   \
+// RUN:               -convert-scf-to-std                                      \
+// RUN:               -convert-std-to-llvm                                     \
+// RUN: | mlir-cpu-runner                                                      \
+// RUN:  -e entry -entry-point-result=void -O0                                 \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_async_runtime%shlibext\
+// RUN: | FileCheck %s --dump-input=always
+
+func @entry() {
+  %c0 = constant 0.0 : f32
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c3 = constant 3 : index
+
+  %lb = constant 0 : index
+  %ub = constant 9 : index
+
+  %A = alloc() : memref<9xf32>
+  %U = memref_cast %A :  memref<9xf32> to memref<*xf32>
+
+  // 1. %i = (0) to (9) step (1)
+  scf.parallel (%i) = (%lb) to (%ub) step (%c1) {
+    %0 = index_cast %i : index to i32
+    %1 = sitofp %0 : i32 to f32
+    store %1, %A[%i] : memref<9xf32>
+  }
+  // CHECK: [0, 1, 2, 3, 4, 5, 6, 7, 8]
+  call @print_memref_f32(%U): (memref<*xf32>) -> ()
+
+  scf.parallel (%i) = (%lb) to (%ub) step (%c1) {
+    store %c0, %A[%i] : memref<9xf32>
+  }
+
+  // 2. %i = (0) to (9) step (2)
+  scf.parallel (%i) = (%lb) to (%ub) step (%c2) {
+    %0 = index_cast %i : index to i32
+    %1 = sitofp %0 : i32 to f32
+    store %1, %A[%i] : memref<9xf32>
+  }
+  // CHECK:  [0, 0, 2, 0, 4, 0, 6, 0, 8]
+  call @print_memref_f32(%U): (memref<*xf32>) -> ()
+
+  scf.parallel (%i) = (%lb) to (%ub) step (%c1) {
+    store %c0, %A[%i] : memref<9xf32>
+  }
+
+  // 3. %i = (-20) to (-11) step (3)
+  %lb0 = constant -20 : index
+  %ub0 = constant -11 : index
+  scf.parallel (%i) = (%lb0) to (%ub0) step (%c3) {
+    %0 = index_cast %i : index to i32
+    %1 = sitofp %0 : i32 to f32
+    %2 = constant 20 : index
+    %3 = addi %i, %2 : index
+    store %1, %A[%3] : memref<9xf32>
+  }
+  // CHECK: [-20, 0, 0, -17, 0, 0, -14, 0, 0]
+  call @print_memref_f32(%U): (memref<*xf32>) -> ()
+
+  dealloc %A : memref<9xf32>
+  return
+}
+
+func @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }