Commit Graph

2 Commits (4199feffed3f9b89e321a8bd65f2b78bc5be2700)

Author SHA1 Message Date
Sean Silva 3695ca83e6 [torch_mlir.compile] Handle the case of already-scripted models better
Closes #1582
2022-11-16 10:47:13 -08:00
Sean Silva 64914603fa [torch_mlir.compile] Add support for multiple exported methods
For AoT deployments models often have multiple exported methods.
This patch enables something like this:

```
class TwoMethodsModule(torch.nn.Module):
    def sin(self, x):
        return torch.ops.aten.sin(x)

    def cos(self, x):
        return torch.ops.aten.cos(x)

example_args = torch_mlir.ExampleArgs()
example_args.add_method("sin", torch.ones(2, 3))
example_args.add_method("cos", torch.ones(2, 4))
print(torch_mlir.compile(TwoMethodsModule(), example_args))
```

In the
[long-term](https://github.com/llvm/torch-mlir/blob/main/docs/long_term_roadmap.md#tools-for-advanced-aot-deployments)
we will need to reconcile this with our story for stateful models and the
backend contract being purely functional. For now, this provides some basic
infra that seems harmless. Arguably, we could tighten up the backend contract
even more to only allow a single compiled function which would prohibit this or
require building out a layer above.

Fixes #1557
2022-11-10 02:10:22 -08:00