[5/x] Lifting/Lowering for cross-context calls

[greenhat]

This PR is stacked on the #358 and should be merged after it

Ref #303

This PR switches the compiler pipeline from Module to Component and adds two stages LowerImportsCrossCtxStage and LiftExportsCrossCtxStage for lifting and lowering components exports/imports for Miden cross-context calls. The implementation so far covers only scalar values as arguments and results, with heap-allocated values support coming later. The rodata is not handled in this PR.

Also, the following changes:

• Introduce CallConv::CrossCtx for Miden cross-context calls ABI;
• Add rust_sdk_cross_ctx_account() and rust_sdk_cross_ctx_note() tests with new Rust cargo-miden projects;

The example of the lifting/lowering MASM code generated can be observed in cross_ctx_account.masm and cross_ctx_note.masm files.

[bitwalker]

Note: Per the Component Translation discussion and our conversation today - it isn't necessary to represent the concept of context switching as an ABI unto itself, rather it is just one property of an ABI - arguably it is a property of a specific call site, not the callee itself, but for our purposes deriving it from ABI information works. I suggest that instead of CrossCtx, we specify three ABIs for use by our Wasm frontend: Wasm, CanonLift, and CanonLower (corresponding to core Wasm ABI and the respective primitives specified in the Canonical ABI).

To elaborate on this:

• Wasm we already have of course, but it doesn't really mean anything, and perhaps could be named something more generic. Essentially, this calling convention mostly indicates that:
  a.) It came from a core Wasm module
  b.) We don't know anything about the "real" types of the arguments/results, i.e. it is mostly opaque from a type system perspective
  c.) The actual calling convention details of the function are encoded in its signature - by which I mean that if a specific convention was expected by the original source language, it has been compiled to Wasm such that those details are already baked in and represented by the core Wasm types that remain. Where it is unclear what the semantics are for a given type, we should consult the core Wasm spec for the call or call_indirect instructions.
• CanonLift and CanonLower correspond to their respective primitives as described in the Canonical ABI spec.
• Functions with the CanonLift ABI represent the actual component boundary, and calling/returning from them involves the actual context switch in the Miden VM. We'll revisit that part in a moment. This ABI is also the only one whose implementation details are not fully specified by the Component Model/Canonical ABI. While the semantics of canon lift are described, the implementation is host-defined, and therefore is ours to dictate. This works out great for us, because this also happens to be where the idiosyncratic details of Miden's context switching must be dealt with, so this gives us flexibility.
• CanonLower is the only ABI other than Wasm that is callable from a Wasm ABI function. The actual details of this ABI are identical to Wasm from the perspective of a caller, however the use of a unique ABI for these functions allows us to identify them as synthetic, and to ensure that calls to a CanonLift function only occur via its CanonLower counterpart. A call to a CanonLower function can be thought of as a call to another Wasm ABI function, but with a layer of lifting/lowering inbetween, so what we end up with is actually a chain of calls composed of all three ABIs, i.e. Wasm (caller) -> CanonLower (wraps Canonical ABI function) -> CanonLift (wraps core function) -> Wasm (callee). A CanonLower function is always paired with a CanonLift function, and is only ever called from a Wasm function. Conversely, CanonLift functions are never called directly from a Wasm function, only from a CanonLower function.
• Neither Wasm nor CanonLower require us to emit any special calling convention-related code, the details of the Canonical ABI itself are materialized in the implementation of the CanonLower and CanonLift functions.
• Both Wasm and CanonLower functions have type signatures consisting solely of core Wasm types.
• CanonLift functions have type signatures consisting solely of Component Model types. For the most part, we don't need to concern ourselves with these types, except in regard to how to pass them across the context switch (i.e. as arguments in the call from CanonLower -> CanonLift, and as results when returning from CanonLift). The lifting/lowering of these types is described here, but the exact way in which we pass them across the CanonLift boundary is up to us.

With all that in mind, the only thing we need to determine if a function invocation should be lowered as exec or call, is to check the ABI of the callee. If it is CanonLift, then we use call, all other function invocations will use exec. As mentioned above, we also will validate all call sites to ensure that calls to CanonLower and CanonLift functions do not originate from invalid callers (e.g. direct from Wasm -> CanonLift). I should note that CanonLower and CanonLift, in terms of ABIs, also encompass all of the other semantics as described in the Canonical ABI spec, so the fact that CanonLift corresponds to a context boundary is only one part of its contract.

[greenhat]

Thanks for the detailed explanation, I changed it in eb2a0ca
I agree, but this part is hard for me to grasp:

• CanonLift functions have type signatures consisting solely of Component Model types. For the most part, we don't need to concern ourselves with these types, except in regard to how to pass them across the context switch (i.e. as arguments in the call from CanonLower -> CanonLift, and as results when returning from CanonLift). The lifting/lowering of these types is described here, but the exact way in which we pass them across the CanonLift boundary is up to us.

I cannot brake from the mental model that CanonLift synthetic function has two signatures. The canon lif function description has CM type signatures, but the generated CanonLift function during the account code compilation has a "lowered" Miden cross-context ABI signature. For example, for the CM signature (list<felt>) -> list<felt> we will generate a synthetic function that uses advice provider, pass hashes and will have the "lowered" Miden CCABI (word) -> word signature. We need to store the CM signature, but I think the "lowered" Miden CCABI signature should also be stored. I understand that we can calculate(flatten) "lowered" signature from the CM signature on-demand, but it seems suboptimal not mentioning that the real synthetic function signature, the one with which it will be called by the corresponding CanonLower synthetic function in the note code, is the "lowered" one.

[bitwalker]

The canon lift function description has CM type signatures, but the generated CanonLift function during the account code compilation has a "lowered" Miden cross-context ABI signature.

It only has a single signature, the "high-level" type signature, there is no low-level signature for the CanonLift function at all. This is because it is called only from a corresponding CanonLower function. The mechanics of the calling convention (when calling from CanonLower -> CanonLift) for those high-level types is up to us, i.e. how do we pass/return e.g. a list? However we see fit, because the CanonLift calling convention is host-defined. This calling convention is also where the context switch occurs, and the mechanics of that are also ours to define.

For example, for the CM signature (list) -> list we will generate a synthetic function that uses advice provider, pass hashes and will have the "lowered" Miden CCABI (word) -> word signature.

Ah, I see where the confusion is coming from. You are assuming that we need to represent the calling convention implementation details in the function signature, resulting in needing two signatures (one to preserve the original information, one to encode the implementation details). This is not necessary though. The implementation details of the calling convention are always derived from the original signature, so it is not necessary for us to encode them as a second signature, we just derive the latter as needed (essentially only during codegen). To be clear here, there is no "Miden ABI" or "Miden calling convention" per se - there are constraints on how we can execute procedure calls in the Miden VM, and so any calling convention must be lowered to MASM in such a way that it works within those constraints. The "lowering" of calling conventions into Miden Assembly is responsible for emitting code that preserves the original calling convention semantics, but adheres to Miden's constraints at the VM level.

There are two places where we might want to do the actual work of preparing a function so that it is "Miden-compatible", by which I mean it can be lowered to Miden Assembly directly following the normal type layout rules, without violating any of the procedure invocation constraints:

1. During code generation, which is where the vast majority of these details will be handled, in particular all of the code related to laying out types in memory, interacting with the advice provider, operand stack management, etc.
2. As an IR transformation prior to codegen, so that we can ensure that code motion affecting function arguments is accounted for when we run the spills transformation. For example, if we have a CanonLift function whose signature is list<felt> -> list<felt>, we must rewrite the function (and any call sites) as if the signature was digest -> digest, as heap-allocated types cannot be passed across a context switch, so we instead store the lists in the advice provider, and pass the key in the advice map from which the actual list can be retrieved. The prologue of the function will then use that digest to extract the list into memory, and replace any uses of the original function argument with the list value once it has been read out of the advice provider. It should be noted that this transformation is implied/derived from the original signature, so we do not need to have the actual function definition to rewrite call sites, nor do we need to worry about call sites when rewriting the function.

The important thing to note about 2 is that while the actual IR is rewritten, the type signature is left unmodified! This is because the rewrite performed here doesn't change the semantics of the function at all, it is purely an artifact of representing those semantics in Miden Assembly.

I should also note that we do, in a sense, store both the original signature, and the "effective" signature here, because in HIR2 a Function has a signature attribute (a Signature value, corresponding to the original signature, and also providing things like ABI details), but we can also directly interrogate the types of its arguments and results as represented in the IR at any given point in time. However in almost all cases, the two will be the same. The only time we'd allow them to diverge is during codegen, in order to facilitate the rewrite I mentioned above.

[greenhat]

The canon lift function description has CM type signatures, but the generated CanonLift function during the account code compilation has a "lowered" Miden cross-context ABI signature.

It only has a single signature, the "high-level" type signature, there is no low-level signature for the CanonLift function at all. This is because it is called only from a corresponding CanonLower function. The mechanics of the calling convention (when calling from CanonLower -> CanonLift) for those high-level types is up to us, i.e. how do we pass/return e.g. a list? However we see fit, because the CanonLift calling convention is host-defined. This calling convention is also where the context switch occurs, and the mechanics of that are also ours to define.

For example, for the CM signature (list) -> list we will generate a synthetic function that uses advice provider, pass hashes and will have the "lowered" Miden CCABI (word) -> word signature.

Ah, I see where the confusion is coming from. You are assuming that we need to represent the calling convention implementation details in the function signature, resulting in needing two signatures (one to preserve the original information, one to encode the implementation details). This is not necessary though. The implementation details of the calling convention are always derived from the original signature, so it is not necessary for us to encode them as a second signature, we just derive the latter as needed (essentially only during codegen). To be clear here, there is no "Miden ABI" or "Miden calling convention" per se - there are constraints on how we can execute procedure calls in the Miden VM, and so any calling convention must be lowered to MASM in such a way that it works within those constraints. The "lowering" of calling conventions into Miden Assembly is responsible for emitting code that preserves the original calling convention semantics, but adheres to Miden's constraints at the VM level.

This helps and clarifies a lot. Thank you for the detailed explanation! It started to click for me. The high-level signature is surely enough to derive the low-level "implementation detail" one.

There are two places where we might want to do the actual work of preparing a function so that it is "Miden-compatible", by which I mean it can be lowered to Miden Assembly directly following the normal type layout rules, without violating any of the procedure invocation constraints:

1. During code generation, which is where the vast majority of these details will be handled, in particular all of the code related to laying out types in memory, interacting with the advice provider, operand stack management, etc.
2. As an IR transformation prior to codegen, so that we can ensure that code motion affecting function arguments is accounted for when we run the spills transformation. For example, if we have a CanonLift function whose signature is list<felt> -> list<felt>, we must rewrite the function (and any call sites) as if the signature was digest -> digest, as heap-allocated types cannot be passed across a context switch, so we instead store the lists in the advice provider, and pass the key in the advice map from which the actual list can be retrieved. The prologue of the function will then use that digest to extract the list into memory, and replace any uses of the original function argument with the list value once it has been read out of the advice provider. It should be noted that this transformation is implied/derived from the original signature, so we do not need to have the actual function definition to rewrite call sites, nor do we need to worry about call sites when rewriting the function.

The important thing to note about 2 is that while the actual IR is rewritten, the type signature is left unmodified! This is because the rewrite performed here doesn't change the semantics of the function at all, it is purely an artifact of representing those semantics in Miden Assembly.

This part confuses me. From the previous comment, I understood that CanonLift and CanonLower synthetic functions are generated in the frontend, and now it appears that they are "implemented" in the IR transformation phase or the codegen. This leads me to think that the frontend generates empty synthetic functions which later "implemented" with the "Miden-compatible" lowering you described above either in the IR transformation phase or the codegen. Am I on the right track?

[bitwalker]

There are a couple different things here being accidentally conflated:

• A CanonLower function declaration presents as just another Wasm function, to Wasm callers, from a semantics point of view. However, what it really doing is adapting a function with the CanonLift calling convention to the Wasm convention. An alternate way of thinking about it, is that we could allow Wasm functions to directly call CanonLift functions, but then we would need to emit the equivalent of the CanonLower implementation at every call site. This would generate a lot more code for no benefit.
• The implementation of a CanonLower function emitted by the frontend must handle translating core Wasm types into their Canonical ABI representation. This could be represented with something like cast, but may need to be more explicit (depending on the type, I'm not sure). It then calls the corresponding CanonLift function with arguments in the Canonical ABI representation. Note that, at this point, no low-level calling convention details are being represented - only the type representation from core Wasm to Canonical ABI.
• The CanonLift declaration is only callable from a CanonLower function, which can be thought of as a Canonical ABI caller (since internally it has converted types to the Canonical ABI). Much like how CanonLower presents as a Wasm function from the perspective of a Wasm caller CanonLift presents as a Canonical ABI function to a Canonical ABI caller, by internally adapting the underlying Wasm calling convention function to the Canonical ABI.
• The implementation of a CanonLift function emitted by the frontend must handle translating Canonical ABI types into their core Wasm representation. Again, this could be represented with cast, but likely requires more explicit lowering code for certain types. Once complete, it then calls the underlying Wasm function. Again, no low-level calling convention details are being represented here - just the type representation change.

So in short:

• The frontend synthesizes the CanonLift/CanonLower functions with code that effects the type representation change
• The backend handles the low-level calling convention details (e.g. how Canonical ABI types are passed for the CanonLift convention, how to handle context switches, operand stack management, laying out types in memory, etc.)

[bitwalker]

I should also note that, we can't represent the type representation change purely using cast (except where the cast is obviously applicable, e.g. between integral types and such), because the IR type system doesn't distinguish between WebAssembly and Canonical ABI types, and shouldn't. Casts are essentially representation-agnostic, expressing only a conversion between IR types, and the default Miden ABI layout of those types may not match what is specified by the Canonical ABI.

So by emitting (in the frontend) code to handle the representation change, we ensure that the code correctly lays things out according to the Canonical ABI rules, and then during codegen, we will know that the IR types corresponding to the function arguments/results are in Canonical ABI representation (as it is implied by the calling convention), and we will emit the low-level details of transferring arguments/results in that representation from caller to callee. In other words, the combination of IR type and calling convention imply a specific data layout when used as a function argument or result.

The Wasm aspect of all this is a bit weird. Basically, the Component Model is able to avoid explicitly representing the type representation change in the form of Wasm instructions because the host runtime speaks Wasm and can handle it automatically (and therefore more efficiently). Typically though, adapting types from one type system to another would require the frontend emit code to do that work somewhere, and then the low-level details are left up to the backend. We're in kind of an awkward position, as we must act as both frontend and host runtime: First, we materialize functions that perform the work that would normally have been emitted by a frontend on more traditional targets, to effect the type representation change. Second, we act as the host, by enforcing the shared-nothing boundary via context switching, and managing the movement of data across that boundary. Where it gets confusing is that, in the functions we synthesize to do this, we have more traditional calling convention details in the mix as well. So it gets a bit fuzzy here, in terms of being precise about what parts of the apparent calling convention are handled by the frontend, and what is handled by the backend. Hopefully my commentary here has made that distinction clearer, but I certainly can't blame you if it is still confusing.

[greenhat]

Now, I feel that it finally clicked for me! Thank you for the detailed explanation!

[bitwalker]

This doesn't account for situations in which someone constructs IR that exports a function from the component with an ABI that is something other than Wasm-oriented, e.g. SystemV. To be clear, this isn't something our tools will produce, but someone else lowering to HIR may not do so through the Wasm frontend, and instead emit HIR components directly with whatever calling conventions they find appropriate.

Basically, I think we only ever synthesize these "lifting" functions when they correspond to a canon lift declaration, and only then. This is one of the reasons why synthesizing these functions in the frontend, where that information is present, is preferable.

[greenhat]

I'm not sure I understand it. Canonical here means the Miden ABI. That's confusing. It's another type we have for describing the component imports/exports - https://github.com/0xPolygonMiden/compiler/blob/greenhat/i303-cross-ctx-lower/hir-type/src/lib.rs#L597-L608, and it'll be deleted once we migrate to the HIR2. We'll generate the synthetic functions in the frontend only for the canon lift/lower ones.

[bitwalker]

Yeah I'm saying we do not need to represent the notion of a "Miden ABI" at all, it is implied that lowering a function to Miden Assembly by definition must adhere to the "Miden ABI". From a call site perspective, it is sufficient to know only the original callee signature (including its calling convention), because the low-level details of how to invoke the resulting Miden Assembly procedure can be completely derived from that signature.

Backing up a bit, and to help clarify conversations about this going forward, here's how I define the distinction between ABIs and calling conventions:

• An ABI describes how a given set of types are laid out in memory, and how procedure calls in general are executed by the machine, such that two applications can even begin to talk to one another. It is a very low-level, machine-dependent specification. For our purposes, ABI consists of:
  • The specification for how byte-addressable memory is laid out in Miden's word-addressable memory, and on the Miden operand stack
  • The bitwise representation of types in the IR type system, i.e. how they are laid out in byte-addressable memory
  • The specification of what constitutes valid immediates in IR type system terms
  • The maximum addressable depth of the operand stack, and how overflow is handled by the VM
  • Operand stack effects of the various procedure invocation instructions: i.e. exec has none, call implicitly hides/restores the overflow portion of the operand stack and switches context, dynexec consumes the first element on the stack as a procedure reference, etc.
  • Program lifecycle, e.g. rodata initialization, global constructors, etc.
• A calling convention encodes the information necessary to understand how to call a given function, i.e. how it expects to receive its arguments, and how you can expect to receive its results. A calling convention is generally specified in terms of a specific ABI. For example, the C calling convention can mean different things depending on the target ABI/architecture. In general, this includes:
  • The set of types specified by that convention
  • The way in which specific types are expected to be passed at the machine level (i.e. is a struct passed by value, or by reference, is there a threshold at which that happens). Is there a distinction between integers and floats, etc.
  • How to handle extension of integer values whose type is smaller than the machine word (i.e. are they sign-extended, or zero-extended). This isn't as important in a stack machine, but in a register machine it is essential, since different machine instructions can read the same register as different bit-widths, and failing to sign/zero-extend a value can result in silent corruption of computed results.
  • The range of possible function effects (e.g. are calls asynchronous, will they return, can they throw exceptions) and how those are handled
• In addition to those more general properties, any calling convention we support must also specify:
  • The procedure call protocol for passing arguments via the advice provider when the call involves a context switch, or when the callee arguments would overflow the addressable stack.
  • Whether the call involves a context switch (can also be considered a function effect)

So, with that in mind, what I'm trying to say is that the notion of a "Miden ABI" is implicit - we do not have any other target than the Miden VM, so by definition all of our calling conventions are/must be specified in terms of how they are represented at the Miden VM level.

So consider the Wasm calling convention, it specifies what types are allowed in function signatures using that convention (i.e. any type described in the core Wasm spec). We have further defined how that convention is to be translated into Miden Assembly. When the IR is constructed, we will validate that there are no convention violations (i.e. that there aren't argument types with no corresponding Wasm representation). Only when we move to the code generation stage will we actually need to deal with calling convention details such as spilling of excess arguments, layout of types on the operand stack/in memory, etc.

Consider another example: the CanonLift convention. This one is maybe more interesting, because not only does the convention imply a context switch, it also permits a much broader range of types. Because of the context switch, we are required to handle certain types (i.e. heap-allocated types) a certain way, by passing them via the advice provider. We must also account for spilling of excess arguments, but again, due to the context switch we are required to handle spills a specific way as well (via the advice provider). Beyond those details, we have not yet defined specifically how the various allowed types permitted by the convention will be passed at the VM level, but that is entirely our perogative. For the most part, I suspect that we will treat most types the same we do in Wasm, as that already accounts for most (all?) scalar types allowed.

[greenhat]

Whoa! That's gold! I'm certainly guilty of conflating calling convention and ABI many times.

[bitwalker]

Yeah me too, I've been playing fast and loose with the terminology in a lot of these conversations and now we're paying the price 😅

[bitwalker]

Overall, I think we could probably merge this as-is - but, once you start integrating HIR2, I think you'll find that much of this PR changes. Some of that I brought up in my comments, but obviously a fair bit is just inherently due to the differences in the IRs.

So I guess my question is - do we merge this, and then refactor it once HIR2 is merged, or merge HIR2, and then rework/recreate this PR on top of the new IR? On the one hand, the more code that has changed, the more work it is to integrate; but on the other, maybe you will find it easier to do the integration with all of this landed. What's your preference?

I think the only change I'd like to make to this PR, assuming it is practical to do so, is the calling convention changes I mentioned. However, if you think it is easier to do that change as part of the larger HIR2 integration, then I'm down to just merge this as-is, and go that route. Ultimately, whether they happen before or as part of integration doesn't make much of a difference, mostly just trying to avoid making things harder.

I'll get #332 wrapped up tomorrow (Tuesday), unfortunately ran out of time today with PR reviews, etc., so if you don't catch me with an update tonight before I sign off, I'll check in on this first thing in the morning.

[greenhat]

Thanks for the thorough review! I addressed all your notes and changed the CallConv. I'd like to merge it as is and take care of the rest in the HIR2 migration.

[bitwalker]

Looks good for now, I left a few more notes that I think will help clarify some of the confusing aspects of ABI vs calling convention, let me know if you still aren't sure what to make of that portion of my feedback.

Hopefully I can have all the HIR2 stuff ready for you by tomorrow, at the very least I there will be updates for you to review, and we can start to discuss any open questions of integrating it into the frontend.