WebAssembly / component-model

Repository for design and specification of the Component Model
Other
914 stars 78 forks source link

Allow non-canonical type alignment? #258

Open abrown opened 10 months ago

abrown commented 10 months ago

Over in wasi-nn, we have been wondering whether there might be a way to specify an alignment for a type larger than the one specified by the canonical ABI. This is similar to what @titzer brings up in https://github.com/WebAssembly/component-model/issues/194 but I thought it was different enough to merit its own issue.

wasi-nn currently specifies the tensor data as list<u8> (WIT). The initial rationale (long ago, back in the WITX days) was a performance consideration: ML frameworks can deal with tensor data as byte blobs and if we can just "pass the bytes through" we can avoid some extra copies. This means that when we, as a wasi-nn guest user, receive tensor data back we must convert it somehow from list<u8> to list<f64>, e.g. This conversion is an unsafe cast in some of the bindings — what if the list<u8> is not aligned to cast to list<f64>? There has been an outstanding issue (https://github.com/WebAssembly/wasi-nn/issues/27) to figure this out. One solution would be to always copy the data; another to conditionally copy the data only if unaligned; etc.

The solution I'm asking about here is adding some way to inform the runtime of the alignment we would like when allocating space for the tensor data. Is this possible? Is it different enough from what @titzer is asking? I suspect we would only want to allow allocations for alignments larger than what the canonical ABI specifies, but maybe not.

lukewagner commented 10 months ago

In theory, yes, it makes sense and seems complementary to what we were talking about in #194, and maybe belongs as part of the same feature.

Just to help me understand the motivation (and see if there's a short-term workaround): what's the reason for using a list<u8> instead of, say, a list<u64> (which would have essentially the same Canonical ABI as list<u8> except with the additional alignment check)? That would mean having separately-named "overloads" for list<u16>-vs-list<u32>-vs-list<u64>; is that the problem?

abrown commented 10 months ago

Well, the list<u64> would "overallocate" when used with smaller types and certain unaligned sizes (e.g. passing 42 bytes of u8 would result in 6 extra bytes allocated... with zeroes?). This is not really a huge problem but I was hoping for better. And having separate functions for each list<*> type is a bit wordy, but again not a huge problem. So, yes, that is the essence of the problem. There are a bunch of possible solutions here, some hackier than others, and what I'm hoping to figure out from this issue is whether adding some extra syntax to WIT is a feasible option.

lukewagner commented 10 months ago

Yeah, I was asking about having a separate list<T> for each of u8, u16, u32 and u64. That would be the easy fix while we work on a more a full byte-layout-control mechanism.

abrown commented 9 months ago

Yeah, I don't think I'm in a terrible hurry so I can see two options: (a) use list<T> for all the types temporarily and then switch over or (b) just wait for the byte-layout-control mechanism. For (b), are you thinking about more than just an optional annotation to the type definitions, like <defvaltype> ('align=' <align>)? where align ::= 1 | 2 | 4 | 8 | 16? It would need some way to say "don't allow alignment less than the minimum the canonical ABI specifies" or something like that, but I didn't see this addition as too complicated. Or maybe it is complicated by all the places it would need to be implemented in?

[edit: I'm kind of asking, do you just want me to submit a PR for this kind of thing or do you think more thought/discussion is needed?]

lukewagner commented 9 months ago

(sorry for the delay; back from the wasm CG meeting) That's a good question; thanks for asking. Given that Preview 2 is almost done and there's not really time to make a significant addition like this, I think it's probably a good idea to go with approach (a), so that you can produce a nice working Preview 2 interface without waiting on anything else, and then we can include this issue as a use case for a possible Preview 3-or-after feature addition to the C-M/WIT that would improve the quality/ergonomics of the existing working interface.