Cannot export model with torch.arange about torch-mlir HOT 4 CLOSED

I came up a simpler test case.

import torch
import torch_mlir

torch_mlir.debug_trace_to_stderr()

mb = torch_mlir.ModuleBuilder()
N = 10
with mb.capture_function("arange_test", []) as f:
    x = torch.arange(N)
    f.returns([x])

It seems that even though the dispatch is going to the fallback kernel, when in the dispatcher::callBoxed function, the dispatchKey is c10::DispatchKey::Undefined which means that pytorch is unable to find a boxed implementation for this op. It seems like we might have to implement this op just like you implemented the copyUnderKernel. What do you think?

from torch-mlir.

Thanks - I believe that this gets routed to the "backendSelect" dispatch key, which is used for kernels that don't have tensor arguments (a heuristic on tensor arguments is what drives dispatch key selection). I know roughly what goes on there but this is the first I've run in to fixing it. I've got a version of PyTorch that dumps all of the dispatch metadata, which helps in these "where'd the kernel go" diagnoses. I just sent out/landed the other PRs and my queue is clear now. Let me take some time tomorrow to triage these (great!) test cases you've come up with.

from torch-mlir.

This appears to be a specific problem with kernels that have out parameters: aten::arange.start_out(Scalar start, Scalar end, Scalar step=1, *, Tensor(a!) out) -> (Tensor(a!))

The sequence on the stack is not as expected and something is getting inverted.

from torch-mlir.

Actually, looking at blame logs, may have been fixed two days ago (yay for procrastination). Rebasing and seeing...

from torch-mlir.