[WASI-NN] Add support for a PyTorch backend for wasi-nn

[rahulchaphalkar]

This change adds a PyTorch backend for wasi-nn.
tch crate is used for Libtorch bindings. I have added an image classification example to demonstrate its usage, which uses a torchscript model.
This backend is currently gated behind a wasi-nn feature flag --features pytorch as due to dynamic linking, a Libtorch v2.4.0 installation on the system (specified by LIBTORCH=/path/to/libtorch) is needed for building.

[abrown]

This is a good start. The main thing to fix is the handling of the input and output tensors.

[abrown]

The cargo vet situation is a bit much:

    cargo vet diff zstd-safe 5.0.1+zstd.1.5.2 5.0.2+zstd.1.5.2
                                                          gyscos         zstd                        2 files changed, 4 insertions(+), 4 deletions(-)
    cargo vet diff zstd 0.11.1+zstd.1.5.2 0.11.2+zstd.1.5.2
                                                          gyscos         zip                         3 files changed, 5 insertions(+), 5 deletions(-)
    cargo vet diff num-complex 0.4.2 0.4.6                cuviper        ndarray                     6 files changed, 188 insertions(+), 48 deletions(-)
      NOTE: this project trusts Josh Stone (cuviper) - consider cargo vet trust num-complex or cargo vet trust --all cuviper
    cargo vet inspect constant_time_eq 0.1.5              cesarb         zip                         311 lines
    cargo vet diff sha1 0.10.5 0.10.6                     newpavlov      zip                         7 files changed, 302 insertions(+), 20 deletions(-)
    cargo vet inspect rawpointer 0.2.1                    bluss          ndarray and matrixmultiply  559 lines
    cargo vet diff zip 0.6.4 0.6.6                        Plecra         tch and torch-sys           14 files changed, 604 insertions(+), 109 deletions(-)
    cargo vet inspect inout 0.1.3                         newpavlov      cipher                      1112 lines
      NOTE: cargo vet import zcash would eliminate this
    cargo vet inspect pbkdf2 0.9.0                        tarcieri       zip                         1120 lines
    cargo vet inspect bzip2 0.4.4                         alexcrichton   zip                         2094 lines
      NOTE: this project trusts Alex Crichton (alexcrichton) - consider cargo vet trust bzip2 or cargo vet trust --all alexcrichton
    cargo vet inspect safetensors 0.3.3                   Narsil         tch                         2200 lines
    cargo vet inspect cipher 0.4.4                        newpavlov      aes                         2635 lines
      NOTE: cargo vet import zcash would reduce this to a [130](https://github.com/bytecodealliance/wasmtime/actions/runs/10836457564/job/30070281197?pr=9234#step:6:131)0-line diff
    cargo vet inspect password-hash 0.3.2                 tarcieri       pbkdf2                      3139 lines
    cargo vet inspect base64ct 1.6.0                      tarcieri       password-hash               3381 lines
    cargo vet diff half 1.8.2 2.4.1                       starkat99      tch                         19 files changed, 2546 insertions(+), 958 deletions(-)
    cargo vet inspect time 0.1.44                         jhpratt        zip                         3915 lines
    cargo vet inspect aes 0.7.5                           tarcieri       zip                         6822 lines
    cargo vet inspect matrixmultiply 0.3.8                bluss          ndarray                     7934 lines
    cargo vet inspect ndarray 0.15.6                      jturner314     tch                         41996 lines
    cargo vet inspect torch-sys 0.17.0                    LaurentMazare  tch                         52119 lines
    cargo vet inspect bzip2-sys 0.1.11+1.0.8              alexcrichton   bzip2                       264[133](https://github.com/bytecodealliance/wasmtime/actions/runs/10836457564/job/30070281197?pr=9234#step:6:134) lines
      NOTE: this project trusts Alex Crichton (alexcrichton) - consider cargo vet trust bzip2-sys or cargo vet trust --all alexcrichton
    cargo vet inspect tch 0.17.0                          LaurentMazare  wasmtime-wasi-nn            2287297 lines

[rahulchaphalkar]

This is a good start. The main thing to fix is the handling of the input and output tensors.

Thanks for the review, Andrew. I've marked smaller Nits as resolved, and I've addressed other comments as well, but kept them 'unresolved' as of now until you take a look.

[abrown]

This is almost there: we're still ignoring the index, though. What might need to happen is that we set up inputs as a Vec<Option<TchTensor>> filled with None and then set the right item based on the index. We are able to retrieve the number of inputs, right?

[rahulchaphalkar]

The index is being ignored because jit-compiled pytorch models expect tensors in order, similar to what is achieved by vector of tensors in set_input(). Do we need to use index here to keep things consistent with wasi-nn? I don't think I saw a direct way to retrieve the number of inputs to a model, I can look into that further. However, someone using pytorch would probably not intend on using index.
This was my previous comment on ignoring the index:

This is one of the differences for this backend. The module's forward method should handle multiple inputs appropriately if it does support multiple inputs. The vector of input tensors being passed to forward should be sufficient, no index or name is needed.

[rahulchaphalkar]

Did you mean use the index as the position/index for the vector of tensors? That sounds good. I'll see if there's a way to determine the max number of inputs to the given model.

[rahulchaphalkar]

Alright, so here's what I've changed -

• set_input() now uses the id passed to it. If it is u32, it assigns the tensor to the vector inputs[id], assigning None if there are any empty spots below the given id. This will give the user flexibility to set input tensors non sequentially.
• compute() will check the vector for any None values, and give an error if present, before calling forward_ts().
• There is no reliable way at this time to get max inputs for a model. However, assigning more inputs than available returns an error message detailing the max inputs the model expects at run time, so this is helpful for the user. The expectation, similar to other backends, is for the user to be aware of input size/shape etc.
• If id is a string, I've currently permitted only single input to be set, effectively ignoring the index. If more inputs are assigned, or if u32 and String indexes are used together, we give out an error.

Let me know if this looks fine. Thanks.

[abrown]

Yeah, looking at 727a4aa, I think that makes more sense. I was kind of hoping that CModule::named_parameters would return the list of inputs by name, but, if that's not the case, then let's just return an error for the Id::Name side of things. In any case, we probably don't need id_type: if we get an Id::Index we know were to put it in the vector and if we get an Id::Name then we either (a) look up its index in named_parameters or (b) if that's not possible, return an error.

[rahulchaphalkar]

I had added the Id:Name case because the wit test cases (nn_wit_image_classification_pytorch in this case) need a Name instead of Index to pass - nn.rs and wasi-nn.wit. Although it looks like set-input might be going away anyway.

[abrown]

Can we find a smaller model or download this instead? Not all Wasmtime users probably want to download this file... twice.

[rahulchaphalkar]

I've made the following changes -

• Created a repository to generate jit-compiled libtorch models, and download models from Release for tests and examples.
  https://github.com/rahulchaphalkar/libtorch-models
  https://github.com/rahulchaphalkar/libtorch-models/releases/tag/v0.1
• Switched from Resnet18(45MB) to Squeezenet1.1(4.5MB) as the file to download from above repo.
• This libtorch-models repo can issue new releases if we update the pytorch/libtorch version in this backend from current 2.4.0 to ensure model compatibility.

[abrown]

Here's the identical 44.7MB file checked in again.

[rahulchaphalkar]

Downloading from external repo, same as tests

[abrown]

I think 51930fa and da1f467 have unintentionally updated some crates. Can you remove those commits?

[rahulchaphalkar]

I rolled back 2 commits, but there's an issue with the lock file. I had previously attempted to fix this by deleting my lockfile, rebasing off of latest main, and then doing a cargo build to generate any of my lock file changes on top of that.

[rahulchaphalkar]

The failing tests fail due to

error: failed retrieving file 'mingw-w64-x86_64-headers-git-12.0.0.r329.g8f7b5ce36-1-any.pkg.tar.zst.sig' from mirror.clarkson.edu : Operation too slow. Less than 1 bytes/sec transferred the last 10 seconds
 mingw-w64-x86_64-gcc-14.2.0-1-any downloading...
error: failed retrieving file 'mingw-w64-x86_64-mpfr-4.2.1-2-any.pkg.tar.zst' from mirror.msys2.org : Operation too slow. Less than 1 bytes/sec transferred the last 10 seconds

I'm hoping a rerun of CI would help.

[rahulchaphalkar]

@abrown can you take a look