[RFC] ABI Objects for Coroutines

IR & Optimizations
1

Hi All,

This proposal adds new capabilities to the llvm coroutine transformations by providing a mechanism that can be used to separate ABI and user code and at the same time supporting a wider variety of users beyond Swift and C++.

TLDR

This proposal adds an ABI object using a class hierarchy that can be used to implement the base ABIs (Switch, Asyc, and Retcon{Once}), as well as allowing customization of these ABIs through inheritance. Generator lambdas provided to CoroSplit are used to create the customized ABIs according to the new intrinsic llvm.coro.begin.custom. These changes will allow us to improve the separation of the code related to users, ABIs and utilities. Additionally, user specific code that need not be part of the llvm coroutine transforms can be implement within their users plugin libraries, so no need to pollute the upstream code. No code changes are required by any existing users. Instead the changes will allow future development and refactoring to incrementally improve the separation of code for different ABIs, users and shared utilities as well as enabling custom ABIs to be added by plugin libraries.

Problems Addressed

Separation of user code: Users such as Swift and C++ require specific behavior from the CoroSplit pass and its related helpers, especially, CoroFrame. However, that has resulted in code specific to certain users being shared and executed by all the users. For example,

buildCoroutineFrame

// For now, this works for C++ programs only.
buildFrameDebugInfo(F, Shape, FrameData);

Separation of ABI code: To address the most significant cases of the above an ABI enum was added and is used by the many methods, functions, classes, etc. invoked by CoroSplit. Some of these uses perform a delegation of certain operations, such as how function splitting is invoked (below). However, in other places the enum is used for fine-grained differences between the ABIs such as those uses in buildCoroutineFrame (below) and it does not provide a mechanism to separate the code of different users.

splitCoroutine

switch (Shape.ABI) {
case coro::ABI::Switch:
SwitchCoroutineSplitter::split(F, Shape, Clones, TTI);
break;
case coro::ABI::Async:
splitAsyncCoroutine(F, Shape, Clones, TTI);
break;
case coro::ABI::Retcon:
case coro::ABI::RetconOnce:
splitRetconCoroutine(F, Shape, Clones, TTI);
break;
}

buildCoroutineFrame

if (Shape.ABI != coro::ABI::Async && Shape.ABI != coro::ABI::Retcon &&
Shape.ABI != coro::ABI::RetconOnce)
sinkLifetimeStartMarkers(F, Shape, Checker, DT);

Materializable callback: This was added to solve functional issues with our non-C++, non-Swift user. It provides a limited method to control rematerialization of pre-suspend instructions after each suspend point and it is the only supported mechanism to customize CoroSplit. However, it is a function callback so it cannot access any state info being used by the CoroSplit pass. This makes it difficult to use to do anything more than checking the type of instruction and making a determination from that. The callback cannot be used to make an individual decision for each instruction and suspend point, i.e. it is not possible to materialize an inst based on its use or materialize an inst after a subset of suspend points that it may cross.

Plugin Libraries: Other uses such as C++ and Swift would ideally implement their own specific behavior within a plugin library that resides within their repos. Consider for example that llvm.coro.alloca.alloc, llvm.coro.alloca.get and llvm.coro.alloca.free require a fair bit of code within CoroFrame and other helpers to handle them. This includes the classes: CoroAllocaAllocInst, CoroAllocaGetInst & CoroAllocaFreeInst and methods: lowerLocalAllocas, lowerNonLocalAlloca, code within buildCoroutineFrame and various others. However, as far as I can tell these intrinsics are only used by Swift and they are not even documented in Coroutines in LLVM — LLVM 20.0.0git documentation. The problem is that most of the interfaces that are required by users for customization are not accessible to plugin libraries and this means right now user code must be embedded into the llvm transform passes.

Proposed solution

  1. Move helpers and utilities that are used by various ABIs/users into separate sources. For example, Shape can be moved to its own header so other users can use it.

  2. Split and simplify certain methods that can be used to ‘delegate’ certain operations of an ABI. For example, buildCoroutineFrame and the Shape constructor.

  3. For ABI code separation define several ABI objects (SwitchABI, AsyncABI, AnyRetConABI) using a class hierarchy. Each ABI implements a set of virtual methods, using the shared utils, to provide the specific behavior required.

  4. For user code separation, improved materialization, and to support plugin libraries add a constructor to CoroSplit that takes a list of ABI generator lambdas, the ABIs inherit from the above set and provide specific behavior. To select the generator a new intrinsic llvm.coro.begin.custom is used in place of llvm.coro.begin.

  5. For plugin librariers the utility headers are moved into include/llvm/Transform/Coroutines.

Discussion

  • The ABI objects use a delegate style approach, each virtual method performs a high-level operation required by CoroSplit. For example, buildCoroutineFrame, and splitCoroutine are virtual methods in the ABI object.
  • Users requiring specific behavior will: 1) implement their desired behavior by inheriting one of the ABI objects (SwitchABI, AsyncABI, AnyRetConABI), 2) provide a list of lambda generators for their customized ABIs when adding CoroSplit to the CGPM, and 3) specify their customized ABI by using the coro.begin.custom intrinsic instead of llvm.coro.begin.
  • It is not necessary to make any changes to existing user code. There is no need to update clang or swift. Those who work on c++ coroutines or swift may want to take advantage of the customizability afforded by these changes.
  • There is no need to make any changes to user code that uses the existing Materialization Callback. Notice I didn’t have to modify the test case at all.
  • I added an additional test case for using a custom ABI object.
  • It is out of scope of these changes to rework everything. The purpose is to provide a new way to use the coroutine transformations, not necessarily to force all existing code and users to adopt this mechanism.
  • The series of PRs I am proposing will not seek to remove all uses of the Shape.ABI enum. Not all uses may be necessary or practical to remove and some uses may be considered beneficial to have (i.e. a utility method that works with any ABI).
  • Users such as c++ and swift may benefit from creating their own customized ABIs, but it is out of scope for this series of PRs to do so.

Proposed PRs

I propose the following series of changes, many of these changes are good changes on their own and can be merged while discussing the ABI object (PR 6) proposal.

  1. Move suspend crossing info into its own h/cpp [llvm/llvm-project][Coroutines] ABI Object by TylerNowicki · Pull Request #106306 · llvm/llvm-project · GitHub

  2. Move spilling to its own utils [llvm/llvm-project][Coroutines] Move spill related methods to a Spill utils by TylerNowicki · Pull Request #1 · TylerNowicki/llvm-project · GitHub

  3. Separate normalization from buildCoroutineFrame: [llvm/llvm-project][Coroutines] Split buildCoroutineFrame by TylerNowicki · Pull Request #2 · TylerNowicki/llvm-project · GitHub

  4. Move materialization to its own utils [llvm/llvm-project][Coroutines] Move materialization out of CoroFrame to MaterializationUtils.h by TylerNowicki · Pull Request #3 · TylerNowicki/llvm-project · GitHub

  5. Move Shape to its own header [llvm/llvm-project][Coroutines] Move Shape to its own header by TylerNowicki · Pull Request #4 · TylerNowicki/llvm-project · GitHub

  6. ABI Objects for Base, Switch, Async, and AnyRetcon [llvm/llvm-project][Coroutines] ABI Object by TylerNowicki · Pull Request #5 · TylerNowicki/llvm-project · GitHub

  7. Move util headers to include/llvm for testing and plugin libraries [llvm/llvm-project][Coroutines] Move util headers to include/llvm by TylerNowicki · Pull Request #6 · TylerNowicki/llvm-project · GitHub

  8. Support custom ABIs and plugin libraries [llvm/llvm-project][Coroutines] Support Custom ABIs and plugin libraries by TylerNowicki · Pull Request #7 · TylerNowicki/llvm-project · GitHub
    Note: This includes a unit test for custom ABIs as used by plugin libraries.

2 Likes
2

I like it. It makes things more flexible.

3

Here is a PR for the documentation changes.

  1. Documentation for using custom ABIs with a plugin library [Coroutines] Documentation for plugin libraries and custom ABIs by TylerNowicki · Pull Request #9 · TylerNowicki/llvm-project (github.com)