Page MenuHomePhabricator

Optionally eliminate blocking runtime.await calls by converting functions to coroutines.
ClosedPublic

Authored by bakhtiyarneyman on Jul 21 2021, 4:35 PM.

Details

Summary

Interop parallelism requires needs awaiting on results. Blocking awaits are bad for performance. TFRT supports lightweight resumption on threads, and coroutines are an abstraction than can be used to lower the kernels onto TFRT threads.

Diff Detail

Event Timeline

bakhtiyarneyman requested review of this revision.Jul 21 2021, 4:35 PM

Address clang warnings.

mehdi_amini added inline comments.Jul 23 2021, 7:29 PM
mlir/include/mlir/Dialect/Async/Passes.td
47

Why is this an option and not just the default?

mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
564

This is repeating the function name, can you expand on how it is achieving this instead?

568

Nit: you don't need to specify the size in SmallVector anymore.
(same below in other places)

574

setResultAttrs and getResultAttrs aren't really efficient API. setAllResultAttrs and getAllResultAttrs should be able to do it better.

But also, it seems that what you really want to do here is to say "add a token as the first return for the function", and there is an API for this on FuncOp I believe:

/// Insert a single result of type `resultType` at `resultIndex`.
  void insertResult(unsigned resultIndex, Type resultType,
                    DictionaryAttr resultAttrs) {
578

Nit early-form reduces indentation:

for (Block &block : func.getBlocks()) {
    if (&block != coro.suspend)
       continue;
    Operation *terminator = block.getTerminator();
    ...
586

Where is this invariant verified/enforced before getting into this point?

593
597

Can you make this function (and all the others) static please

610

Please favor for-range loop and avoid indexing when possible:

SmallVector<Value> unwrappedResults;
unwrappedResults.reserve(newCall->getResults().size() - 1);
for (Value result : newCall.getResults().drop_front())
   unwrappedResults.push_back(callBuilder.create<AwaitOp>(loc, result));
612

I don't understand this comment, it may miss some context: dangling pointer where? Is this referring to the comment in front of this function The invocation of this function is safe only when call ops are traversed in reverse order of how they appear in a single block. See funcsToCoroutines.? If so, it seems a bit out-of-place here (we're not careful in these two lines, only in the contract with this function and its caller which is documented on the API itself).

646

walk() is recursive and would traverse all the ops inside the functions.

653–654

(you'll see that LLVM/MLIR uses SmallVector for worklists in general, and because we don't support exceptions we have pop_back_val() :) )

660

You can save double-queries to maps when you know you want the entry inserted regardless:

auto newMapEntry = outlinedFunctions.insert({func, CoroMachinery{});
if (!insertion.second) {
    // This function has already been processed because this is either
    // the corecursive case, or a caller with multiple calls to a newly
    // created corouting. Either way, skip updating the call sites.
    continue;
}
insertion.first->second = rewriteFuncAsCoroutine(func);
666
679

I suspect this can fire on valid IR: can you make this a proper error handling (or if you want to keep an assert here, document the invariant and make it a runtime check and signalPassFailure() at the beginning of runOnOperation.

ezhulenev added inline comments.Jul 25 2021, 8:38 AM
mlir/include/mlir/Dialect/Async/Passes.td
47

I should be the default once it will be tested, I'm +1 to have it as an option for now, so it can be quickly disabled without a need to push complex fixes.

mehdi_amini added inline comments.Jul 25 2021, 12:08 PM
mlir/include/mlir/Dialect/Async/Passes.td
47

OK, then please document it if the option is temporary for bringup :)

bakhtiyarneyman marked 16 inline comments as done.

Addressing comments.

bakhtiyarneyman marked an inline comment as done.Jul 26 2021, 4:34 PM
bakhtiyarneyman edited the summary of this revision. (Show Details)
mehdi_amini added inline comments.Jul 26 2021, 11:10 PM
mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
576

Is this all still useful here? I'd expect the call below to take care of all this.

679

How is this handled now?

ezhulenev added inline comments.Jul 27 2021, 10:44 AM
mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
572

SmallVector<Type> resultTypes = {TokenType::get(ctx)} to skip func.insertResult below?

mehdi_amini added inline comments.Jul 27 2021, 10:46 AM
mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
572

insertResult actually handles also argument attributes remapping.

576

Oh I missed that this is wrapping the results into a "ValueType", forget this comment.

bakhtiyarneyman marked 4 inline comments as done.

Addressing comments.

bakhtiyarneyman marked an inline comment as done.Jul 27 2021, 2:20 PM
bakhtiyarneyman added inline comments.
mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
572

Yes, the reason we are using insertResult is to avoid attributes being shifted.

679

I think it was not pushed properly, PTAL.

ezhulenev accepted this revision.Jul 28 2021, 12:25 PM
This revision is now accepted and ready to land.Jul 28 2021, 12:25 PM