Page MenuHomePhabricator

[mlir:Async] Add intermediate async.coro and async.runtime operations to simplify Async to LLVM lowering
ClosedPublic

Authored by ezhulenev on Jan 18 2021, 10:58 AM.

Details

Summary

[NFC] No new functionality, mostly a cleanup and one more abstraction level between Async and LLVM IR.

Instead of lowering from Async to LLVM coroutines and Async Runtime API in one shot, do it progressively via async.coro and async.runtime operations.

  1. Lower from async to async.runtime/coro (e.g. async.execute to function with coro setup and runtime calls)
  2. Lower from async.runtime/coro to LLVM intrinsics and runtime API calls

Intermediate coro/runtime operations will allow to run transformations on a higher level IR and do not try to match IR based on the LLVM::CallOp properties.

Although async.coro is very close to LLVM coroutines, it is not exactly the same API, instead it is optimized for usability in async lowering, and misses a lot of details that are present in @llvm.coro intrinsic.

Diff Detail

Event Timeline

ezhulenev created this revision.Jan 18 2021, 10:58 AM
ezhulenev requested review of this revision.Jan 18 2021, 10:58 AM
ezhulenev edited the summary of this revision. (Show Details)Jan 18 2021, 11:04 AM
ezhulenev added a reviewer: mehdi_amini.
ezhulenev retitled this revision from [mlir:Async] Add intermediate async.coro and async.runtime to simplify Async to LLVM lowering to [mlir:Async] Add intermediate async.coro and async.runtime operations to simplify Async to LLVM lowering.Jan 18 2021, 11:06 AM

Intermediate coro/runtime operations will allow to run transformations on a higher level IR and do not try to match IR based on the LLVM::CallOp properties.

Although async.coro is very close to LLVM coroutines, it is not exactly the same API, instead it is optimized for usability in async lowering, and misses a lot of details that are present in @llvm.coro intrinsic.

Have you considered "importing" the LLVM intrinsics as MLIR operations? There are numerous examples (nvvm, avx512 to name a few) and automation tools to achieve this. It will remove the awkward manual creation of function calls that persists in the code even after this refactoring. I understand the added value of type safety, but it looks possible to extend the operations so that they accept either !llvm.ptr<i8> or dedicated types !llvm.async.token and implement progressive type rewriting on these. Since these will still be regular operations in MLIR, usability improvements apply equally well.

mlir/include/mlir/Dialect/Async/IR/Async.h
60–84

You can define this in ODS using TypeDef and avoid the boilerplate.

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

It is possible to just omit empty lists for arguments and results.

Do you have code pointers to intrinsic op definitions in the LLVM dialect? I like this idea, however I still would keep 2 levels of IR lowering, async specific async.coro (that omit a lot of irrelevant coro details) that are lowered to LLVM intrinsic ops.

https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td as an example. We also have support in mlir-tblgen to generate the boilerplate for ops given LLVM IR intrinsic definition file, example of usage: https://github.com/llvm/llvm-project/blob/main/mlir/test/mlir-tblgen/llvm-intrinsics.td

ezhulenev updated this revision to Diff 317680.Jan 19 2021, 1:43 PM
ezhulenev marked 2 inline comments as done.

Omit empty ins and outs

mlir/include/mlir/Dialect/Async/IR/Async.h
60–84

Prepared a separate PR: https://reviews.llvm.org/D95000

Will do LLVM intrinsics in a separate change, this one is already pretty large.

ezhulenev updated this revision to Diff 318662.Jan 22 2021, 3:08 PM

Use LLVM coro operations to lower from async.coro

ftynse accepted this revision.Jan 25 2021, 4:54 AM

Nice!

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

Couldn't parse this sentence.

309

Nit: "If it is destroyed"

334

Nit: "an"

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
805–812

Future nit: we probably want to factor this out into something like LLVM::emitSizeof(OpBuilder &, Type)

This revision is now accepted and ready to land.Jan 25 2021, 4:54 AM
ezhulenev updated this revision to Diff 319104.Jan 25 2021, 1:29 PM
ezhulenev marked 3 inline comments as done.

Fix async ops doc

ezhulenev marked an inline comment as done.Jan 25 2021, 1:55 PM
ezhulenev added inline comments.
mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
286

This is mostly a copy of what official coro documention says about coro.end... and I agree it's hard to parse :)

This revision was automatically updated to reflect the committed changes.
ezhulenev marked an inline comment as done.