How to support self-define torch op?

[KangHe000]

Hi, I have some special self-define torch ops, and want to lower them to torch dialect.
Does torch-mlir provide the export mechanism to support self-define torch ops ？
When export to onnx, I can use torch.autograd.function to meet this kind of needs, but torch-mlir seems don't support that.
I got the error below:

python builtin <built-in method apply of FunctionMeta object at 0x56353a8494f0> is currently not supported in Torchscript:
  File "my_symbolic.py", line 20
    def forward(self, x):
        return TestKernel.apply(x)
               ~~~~~~~~~~~~~~~~ <--- HERE

When I try to run my compile code:

import torch
import torch.nn as nn
import torch.autograd as AF
import torch_mlir

class TestKernel(AF.Function):
    @staticmethod
    def symbolic(g, inputs):
        return g.op("torch::symbolic", inputs)
    @staticmethod
    def forward(ctx, inputs):
        out = torch.exp(-0.5 * ((inputs - 1) ** 2) / 4)
        return out

class TestKernelTest(nn.Module):
    def __init__(self):
        super(TestKernelTest, self).__init__()

    def forward(self, x):
        return TestKernel.apply(x)

if __name__ == "__main__":
    model = TestKernelTest()
    x = torch.randn(2, 3)
    res = model(x)
    torch.onnx.export(model, x, "test_symbolic.onnx", verbose=True)
    traced = torch.jit.trace(model, x)
    module = torch_mlir.compile(model, x, output_type=torch_mlir.OutputType.TORCH)
    #module = torch_mlir.compile(model, x, output_type=torch_mlir.OutputType.TORCH, use_tracing=True)

[ramiro050]

Hi @KangHe000,

The error you're getting is from PyTorch, not torch-mlir. Namely, the PyTorch function that is in charge of compiling a module into Torchscript does not seem to support AF.Functions.

In torch-mlir there is currently an open PR that adds custom op support through a different method. The approach is to register a custom op in PyTorch, and then add support for it in torch-mlir.

[KangHe000]

@ramiro050 Sorry for replying so late, I will try register a custom op in PyTorch. Thanks!

[silvasean]

Closing this as it appears we found a solution.
FYI we have a custom ops RFC as well: #1462