Inductor CPP codegen for WebAssembly target

vadimkantorov · September 27, 2023, 9:32pm

Can PyTorch generator AOT kernels for the wasm/browser target (e.g. generate AITemplate style a self-contained, minimal ggml/llama.cpp-style C++ program to benefit from wasm-simd and link it only to XNNPACK which also exists for wasm)?

Related to Small depthwise Conv1d: maximum perf on CPU? - #4 by smth

marksaroufim · September 27, 2023, 10:07pm

github.com

pytorch/pytorch/blob/main/test/cpp/aot_inductor/test.py

import shutil

import torch
from torch._export import aot_compile, dynamic_dim


class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(64, 10)
        weights = torch.arange(640)
        weights = torch.reshape(weights, (10, 64))

        with torch.no_grad():
            self.fc.weight.copy_(weights)
            self.fc.bias.copy_(torch.zeros(10))

    def forward(self, x, y):
        return self.fc(torch.sin(x) + torch.cos(y))

This file has been truncated. show original

vadimkantorov · September 27, 2023, 10:34pm

I guess for this to work, we’d need to get raw self-contained C++ files as output + maybe Makefiles / compilation commands / CMake / etc, so that we can then build them using Emscripten

Maybe some libtorch.so-mobile version can be built for wasm as well, but super-aggressive tree-shaking is needed to reduce the size. Ideally we’d just need to have only used ops wasm-simd C++ code + XNNPACK.