Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow non-canonical type alignment? #258

Open
abrown opened this issue Sep 27, 2023 · 5 comments
Open

Allow non-canonical type alignment? #258

abrown opened this issue Sep 27, 2023 · 5 comments

Comments

@abrown
Copy link
Contributor

abrown commented Sep 27, 2023

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 #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 (WebAssembly/wasi-nn#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
Copy link
Member

lukewagner commented Sep 27, 2023

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
Copy link
Contributor Author

abrown commented Sep 29, 2023

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
Copy link
Member

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
Copy link
Contributor Author

abrown commented Oct 11, 2023

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
Copy link
Member

(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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants