Page MenuHomePhabricator

[mlir] Convert from Async dialect to LLVM coroutines
ClosedPublic

Authored by ezhulenev on Oct 12 2020, 9:10 PM.

Details

Summary

Lower from Async dialect to LLVM by converting async regions attached to async.execute operations into LLVM coroutines (https://llvm.org/docs/Coroutines.html):

  1. Outline all async regions to functions
  2. Add LLVM coro intrinsics to mark coroutine begin/end
  3. Use MLIR conversion framework to convert all remaining async types and ops to LLVM + Async runtime function calls

All async.await operations inside async regions converted to coroutine suspension points. Await operation outside of a coroutine converted to the blocking wait operations.

Implement simple runtime to support concurrent execution of coroutines.

Diff Detail

Event Timeline

ezhulenev created this revision.Oct 12 2020, 9:10 PM
Herald added a project: Restricted Project. · View Herald Transcript
ezhulenev requested review of this revision.Oct 12 2020, 9:10 PM
ezhulenev edited the summary of this revision. (Show Details)Oct 12 2020, 9:13 PM
rriddle added a comment.EditedOct 12 2020, 9:17 PM

Seems like there are several coding style issues, please fix them:
Naming style violations
Namespaces in .cpp files containing things other than classes

mlir/include/mlir/Transforms/Passes.h
59 ↗(On Diff #297755)

Only passes in the Transforms/ library should be exposed here.

ezhulenev edited the summary of this revision. (Show Details)Oct 12 2020, 9:22 PM
ezhulenev added reviewers: mehdi_amini, csigg, herhut.
ezhulenev updated this revision to Diff 297761.Oct 12 2020, 9:43 PM

Add missing targets to CmakeLists

ezhulenev marked an inline comment as done.

Fix LLVM naming style problems

Fix LLVM style: only classes in anonymous namespaces

Seems like there are several coding style issues, please fix them:
Naming style violations
Namespaces in .cpp files containing things other than classes

I think I've fixed all of them.

herhut added inline comments.Oct 13 2020, 2:50 AM
mlir/include/mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h
2

Nit: fix length of comment

mlir/include/mlir/ExecutionEngine/AsyncRuntime.h
72

Should this be prefixed with mlirAsyncRuntime, as well?

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
82

Maybe auto builder = OpBuilder::atBlockTerminator(module.getBody())?

111

Would it make sense to add these as operations to the LLVM dialect?

262

nit: ve -> be

271

Maybe mention that this is where the actual code will be inserted later.

339

builder.setInsertionPointToStart(cleanupBlock) ?

399

What is coroState?

408

Split block before op?

430

Move this closer to use?

437

builder.setInsertionPointToEnd(splitBlock) ?

442

setInsertionPointToStart and create the isZero here?

451

Could this be a RewritePattern, as well?

465

Could valueType be a lambda?

472

Maybe use SymbolTable::insert to avoid name clashes?

506

Would valueMapping.map(usedAbove, func.getArguments()) work?

511

Set insertion point instead?

536

Converter?

538

Do you need an extra class for this or could you simply configure an existing type converter? That would make it easier to compose this lowering with other lowerings to llvm.

555

Could this use the StandardToLLVM pattern?

599

Maybe put an attribute on the coroutine functions to identify them as such? That would avoid the need for this state.

608

return rewriter.replace... to have early return here?

627

Could you find this from the context? Like looking for the call to coro.handle_id?

Not sure this is better but it would avoid the need to keep state between rewrites, allowing this to be purely pattern based.

653

anon. namespace

676

Maybe use Twine(kAsyncFnPrefix, execute.getParentOfType<FuncOp>().getName()) here to get more meaningful names. The use of SymbolTable would make it unique. Could also make the prefix a suffix to get something like foo_coro_1

ezhulenev marked 20 inline comments as done.

Many small fixes.

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
111

I'll let @ftynse to comment on this, not really familiar with LLVM dialect and if adding intrinsics to it makes sense.

339

Changed all other uses of OpBuilder::??? to setInsertonPoint.

399

Added a documentation

506

Cool, didn't know it works like that.

538

Do you have examples, not sure what do you mean.

599

It still would need to lookup coro machinery to get coroutine tokens/handles to setup suspension point.

608

Tried that, but then erase/replace has to be in two places and I've got a feeling that it got harder to follow the control flow.

627

I also though about this option, but it will be more complex, I think a simple lookup table is easier to follow and less code to write.

ezhulenev added inline comments.Oct 13 2020, 10:04 AM
mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
555

By briefly looking at the pass I think it's mostly about memref lowering, or I didn't find anything?

herhut accepted this revision.Oct 20 2020, 3:41 AM

Sorry, this disappeared from my review queue.

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
479

I still think it would be valuable to derive the name of the async function from the original function it came from to ease debugging. But up to you.

538

Never mind, I did not see that addConversion is private,

555

I was thinking that you could extent the LLVMTypeConverter and then register the one extra conversion. Then your patterns could be combined with the LLVM conversion patterns. The pattern for call operations in the LLVM conversion, as it just uses a type converter, would convert your types, as well.

This revision is now accepted and ready to land.Oct 20 2020, 3:41 AM
ezhulenev updated this revision to Diff 299939.Oct 22 2020, 6:00 AM
ezhulenev marked an inline comment as done.

Add TODO to derive better names for outlined coroutines.

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
479

That's good idea, however it's a bit tricky to do right for multiple nested async.execute ops. I'll get back to this after "automatic async parallel for" transformation pass, because it will be even more useful to track the origin of these coroutines.

This revision was automatically updated to reflect the committed changes.

Reverted in 4986d5eaff359081a867def1c6a2e1147dbb2ad6 ; the build with SHARED_LIBRARY=ON was broken.