[mlir] Add microbenchmark for linalg+async-parallel-for

[aartbik]

Nice to see this in action! Right in time also for some parallel work in the sparse compiler!

mlir/integration_test/Dialect/Async/CPU/microbench-linalg-async-parallel-for.mlir

75

Nit: if you pick another name like "SERIAL" for the sequential case, the output nicely lines up

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

234

I am guessing it is okay that merely pulling in the async runtime lib also implies starting a threadpool?

[mehdi_amini]

Are the changes in mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp bug fixes? Can you add IR tests for this (and commit separately so there is a proper commit description)?

mlir/integration_test/Dialect/Async/CPU/microbench-linalg-async-parallel-for.mlir

75

Aren't they identical? Can't we just use "CHECK" and have one?

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

234

I'd expect this to only start the thread pool if there is something to execute (when this function is invoked for the first time)?

[aartbik]

mlir/integration_test/Dialect/Async/CPU/microbench-linalg-async-parallel-for.mlir

75

Interesting. I may have set the wrong example in other test cases, where I duplicate the CHECK-XXX even if the cases are identical to make it easier to see which of the XXX variants was responsible for a failure. But I guess you could indeed use one CHECK and just look at the command line of the failing report.....

[nicolasvasilache]

Nice, hooray for composition :) !

[ezhulenev]

Address PR comments

[ezhulenev]

Will submit changes to AsyncToLLVM.cpp in a separate PR.

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

234

Yes, this will be instantiated only of thread pool is required for async execution. Also this runtime is primarily for testing, e.g. in TF it will be replaced with a TFRT integration.

[jpienaar]

This seems to be failing on some configurations https://buildkite.com/mlir/mlir-core/builds/9510#4c77cda4-9dce-4b4d-8ad7-c220a892865b , known issue?

[ezhulenev]

No, will take a look in ~30 min.

[ezhulenev]

Fixed in: https://reviews.llvm.org/D91912

[mehdi_amini]

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

234

I believe that on the medium-to-long-term we will want a proper runtime here, in the scope of the low-level bits of TFRT. So I rather not consider this "just for testing" if it means that we're don't have a path to productionize this or if it becomes too hacky. This should all make sense and have an integration story other than TF/TFRT upstream.

mlir/integration_test/Dialect/Async/CPU/microbench-linalg-async-parallel-for.mlir

+// RUN:   mlir-opt %s                                                          \
+// RUN:               -linalg-tile-to-parallel-loops="linalg-tile-sizes=256"   \
+// RUN:               -async-parallel-for="num-concurrent-async-execute=4"     \
+// RUN:               -async-ref-counting                                      \
+// RUN:               -convert-async-to-llvm                                   \
+// RUN:               -lower-affine                                            \
+// RUN:               -convert-linalg-to-loops                                 \
+// RUN:               -convert-scf-to-std                                      \
+// RUN:               -std-expand                                              \
+// RUN:               -convert-vector-to-llvm                                  \
+// RUN:               -convert-std-to-llvm                                     \
+// RUN: | mlir-cpu-runner                                                      \
+// RUN:  -e entry -entry-point-result=void -O3                                 \
+// 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 --check-prefix=CHECK-ASYNC
+
+// RUN:   mlir-opt %s                                                          \
+// RUN:               -convert-linalg-to-loops                                 \
+// RUN:               -convert-scf-to-std                                      \
+// RUN:               -convert-vector-to-llvm                                  \
+// RUN:               -convert-std-to-llvm                                     \
+// RUN: | mlir-cpu-runner                                                      \
+// RUN:  -e entry -entry-point-result=void -O3                                 \
+// 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 --check-prefix=CHECK-SEQ
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+
+func @linalg_generic(%lhs: memref<?x?xf32>,
+                     %rhs: memref<?x?xf32>,
+                     %sum: memref<?x?xf32>) {
+  linalg.generic {
+    indexing_maps = [#map0, #map0, #map0],
+    iterator_types = ["parallel", "parallel"]
+  }
+    ins(%lhs, %rhs : memref<?x?xf32>, memref<?x?xf32>)
+    outs(%sum : memref<?x?xf32>)
+  {
+    ^bb0(%lhs_in: f32, %rhs_in: f32, %sum_out: f32):
+      %0 = addf %lhs_in, %rhs_in : f32
+      linalg.yield %0 : f32
+  }
+
+  return
+}
+
+func @entry() {
+  %f1 = constant 1.0 : f32
+  %f4 = constant 4.0 : f32
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %cM = constant 1000 : index
+
+  //
+  // Sanity check for the function under test.
+  //
+
+  %LHS10 = alloc() {alignment = 64} : memref<1x10xf32>
+  %RHS10 = alloc() {alignment = 64} : memref<1x10xf32>
+  %DST10 = alloc() {alignment = 64} : memref<1x10xf32>
+
+  linalg.fill(%LHS10, %f1) : memref<1x10xf32>, f32
+  linalg.fill(%RHS10, %f1) : memref<1x10xf32>, f32
+
+  %LHS = memref_cast %LHS10 : memref<1x10xf32> to memref<?x?xf32>
+  %RHS = memref_cast %RHS10 : memref<1x10xf32> to memref<?x?xf32>
+  %DST = memref_cast %DST10 : memref<1x10xf32> to memref<?x?xf32>
+
+  call @linalg_generic(%LHS, %RHS, %DST)
+    : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> ()
+
+  // CHECK-ASYNC: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2]
+  // CHECK-SEQ: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2]
+  %U = memref_cast %DST10 :  memref<1x10xf32> to memref<*xf32>
+  call @print_memref_f32(%U): (memref<*xf32>) -> ()
+
+  dealloc %LHS10: memref<1x10xf32>
+  dealloc %RHS10: memref<1x10xf32>
+  dealloc %DST10: memref<1x10xf32>
+
+  //
+  // Allocate data for microbenchmarks.
+  //
+
+  %LHS4096 = alloc() {alignment = 64} : memref<4096x4096xf32>
+  %RHS4096 = alloc() {alignment = 64} : memref<4096x4096xf32>
+  %DST4096 = alloc() {alignment = 64} : memref<4096x4096xf32>
+
+  %LHS0 = memref_cast %LHS4096 : memref<4096x4096xf32> to memref<?x?xf32>
+  %RHS0 = memref_cast %RHS4096 : memref<4096x4096xf32> to memref<?x?xf32>
+  %DST0 = memref_cast %DST4096 : memref<4096x4096xf32> to memref<?x?xf32>
+
+  //
+  // Warm up.
+  //
+
+  call @linalg_generic(%LHS0, %RHS0, %DST0)
+    : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> ()
+
+  //
+  // Measure execution time.
+  //
+
+  %t0 = call @rtclock() : () -> f64
+  scf.for %i = %c0 to %cM step %c1 {
+    call @linalg_generic(%LHS0, %RHS0, %DST0)
+      : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> ()
+  }
+  %t1 = call @rtclock() : () -> f64
+  %t4096 = subf %t1, %t0 : f64
+
+  // Print timings.
+  vector.print %t4096 : f64
+
+  // Free.
+  dealloc %LHS4096: memref<4096x4096xf32>
+  dealloc %RHS4096: memref<4096x4096xf32>
+  dealloc %DST4096: memref<4096x4096xf32>
+
+  return
+}
+
+func private @rtclock() -> f64
+
+func private @print_memref_f32(memref<*xf32>)
+  attributes { llvm.emit_c_interface }

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     AsyncRuntimeTypeConverter converter;
 
     SmallVector<Type, 5> resultTypes;
     converter.convertTypes(op->getResultTypes(), resultTypes);
 
     CallOp call = cast<CallOp>(op);
     rewriter.replaceOpWithNewOp<CallOp>(op, resultTypes, call.callee(),
-                                        call.getOperands());
+                                        operands);
 
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Async reference counting ops lowering (`async.add_ref` and `async.drop_ref`
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     const bool isInCoroutine = outlined != outlinedFunctions.end();
 
     Location loc = op->getLoc();
 
     // Inside regular function we convert await operation to the blocking
     // async API await function call.
     if (!isInCoroutine)
       rewriter.create<CallOp>(loc, TypeRange(), blockingAwaitFuncName,
-                              ValueRange(op->getOperand(0)));
+                              ValueRange(operands[0]));
 
     // Inside the coroutine we convert await operation into coroutine suspension
     // point, and resume execution asynchronously.
     if (isInCoroutine) {
       const CoroMachinery &coro = outlined->getSecond();
 
       OpBuilder builder(op);
       MLIRContext *ctx = op->getContext();
 
       // A pointer to coroutine resume intrinsic wrapper.
       auto resumeFnTy = AsyncAPI::resumeFunctionType(ctx);
       auto resumePtr = builder.create<LLVM::AddressOfOp>(
           loc, resumeFnTy.getPointerTo(), kResume);
 
       // Save the coroutine state: @llvm.coro.save
       auto coroSave = builder.create<LLVM::CallOp>(
           loc, LLVM::LLVMTokenType::get(ctx),
           builder.getSymbolRefAttr(kCoroSave), ValueRange(coro.coroHandle));
 
       // Call async runtime API to resume a coroutine in the managed thread when
       // the async await argument becomes ready.
-      SmallVector<Value, 3> awaitAndExecuteArgs = {
-          await.getOperand(), coro.coroHandle, resumePtr.res()};
+      SmallVector<Value, 3> awaitAndExecuteArgs = {operands[0], coro.coroHandle,
+                                                   resumePtr.res()};
       builder.create<CallOp>(loc, TypeRange(), coroAwaitFuncName,
                              awaitAndExecuteArgs);
 
       // Split the entry block before the await operation.
       addSuspensionPoint(coro, coroSave.getResult(0), op);
     }
 
     // Original operation was replaced by function call or suspension point.

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

 #include <cassert>
 #include <condition_variable>
 #include <functional>
 #include <iostream>
 #include <mutex>
 #include <thread>
 #include <vector>
 
+#include "llvm/Support/ThreadPool.h"
+
 //===----------------------------------------------------------------------===//
 // Async runtime API.
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 // Forward declare class defined below.
 class RefCounted;

 mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *group) {
   std::unique_lock<std::mutex> lock(group->mu);
   if (group->pendingTokens != 0)
     group->cv.wait(lock, [group] { return group->pendingTokens == 0; });
 }
 
 extern "C" void mlirAsyncRuntimeExecute(CoroHandle handle, CoroResume resume) {
 #if LLVM_ENABLE_THREADS
-  std::thread thread([handle, resume]() { (*resume)(handle); });
-  thread.detach();
+  static llvm::ThreadPool *threadPool = new llvm::ThreadPool();
+  threadPool->async([handle, resume]() { (*resume)(handle); });
 #else
   (*resume)(handle);
 #endif
 }
 
 extern "C" void mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *token,
                                                      CoroHandle handle,
                                                      CoroResume resume) {