WASM Interface

[leizaf]

While I am workin on finishing up #173, I drafted up some stuff here which allows us to wrap the user's infer implementation with custom logic. I don't know if this would actually work though. For the user it would look like:

use carton_wasm_interface::infer;

#[infer]
fn infer(in_: Vec<(String, Tensor)>) -> Vec<(String, Tensor)> {
    // ...
}

The reason this work around is needed is because the .wit interface needs to be implemented in the same place the bindings are generated. Now we can implement stuff like conversions for candle and returning a pointer (and managing it's lifetime) easily.

Let me know if you have any thoughts! I'm hoping it makes developing wasm models more ergonomic.

[VivekPanyam]

I think the high level idea makes sense.

I drafted up some stuff here which allows us to wrap the user's infer implementation with custom logic. I don't know if this would actually work though.

Since we ultimately need to expose multiple functions from user code, what if we structure it as a trait instead of as an attribute macro? That also has the benefit of allowing types to be a little more flexible (e.g. something like impl IntoIterator<Item = (S, Tensor)> where S: Into<String> instead of Vec<(String, Tensor)>). The user then has to register an implementation of this trait (see below).

The reason this work around is needed is because the .wit interface needs to be implemented in the same place the bindings are generated.

What if our "registration" macro uses wit_bindgen::generate! internally? That way the expanded code will generate and implement the bindings in the same place.

For example:

use carton_wasm_interface::{Model, register_model};

struct SomeUserModel;

impl Model for SomeUserModel {
    fn infer(...) -> ... {
        // ...
    }
    // ...
}

register_model!(SomeUserModel);

The macro would expand to something like:

wit_bindgen::generate!({
    world: "model",
    // Either use an absolute path here or pass in the wit file using the `inline` param
    path: "...",
    exports: {
        world: CartonExportedModel
    }
});
struct CartonExportedModel;
impl Guest for CartonExportedModel {
    fn infer(in: Vec<(String, Tensor)>) -> Vec<(String, Tensor)> {
        // We can do type conversions here
        SomeUserModel::infer(...)
    }
}

I don't think we even need a proc macro to implement this and we can just use declarative macros instead.

What do you think?

[leizaf]

I definitely like the trait a lot more. I'm still debating if the wit generate macro should expand in the user's code or ours. If it's in the users code we don't need to deal with initializing the module, if it's in our code, we can be more flexible with implementing stuff, since we wouldn't need to define them in the macro.

[VivekPanyam]

Expanding it in the user's code (as in my comment above) also has the benefit of not requiring the trait to be object safe.

If we wanted to store a user's trait implementation and expose it to the wit_bindgen::generate! macro as a non-generic struct in our code, we'd likely need to wrap the user's trait implementation in a Box<dyn Model>. Calling methods on that requires object safety which removes flexibility for generics in the trait (e.g. impl IntoIterator<...>).

(The approach in my comment above also has the benefit of not requiring dynamic dispatch, but that almost certainly doesn't matter in this case)

[VivekPanyam]

@leifu1128 just wanted to check if you had a chance to work on this more. Thanks!

[leizaf]

Been super bogged down with school work lately, but I graduate in a few weeks and will have a ton of time to work on this afterwards.