Support for torch.fft

[gorkemalgan]

I have a custom model that uses torch.fft.rfftn and torch.fft.irfftn. I can successfully run capture_pre_autograd_graph and export (only with static sizes though). But, when I run to_edge I get the following error:

Operator torch._ops.aten._fft_r2c.default is not Aten Canonical.

[cccclai]

To temporarily unblock, try to_edge(exir.EdgeCompileConfig(_check_ir_validity=False)) while we're working on address it.

cc: @larryliu0820

[gorkemalgan]

Thanks for the suggestion but now it throws the following error:

Exception has occurred: SpecViolationError
These operators are taking Tensor inputs with mismatched dtypes: defaultdict(<class 'dict'>, {<EdgeOpOverload: aten._fft_r2c.default>: schema = aten::_fft_r2c(Tensor self, int[] dim, int normalization, bool onesided) -> Tensor: {'self': torch.float32, '__ret_0': torch.complex64}, <EdgeOpOverload: aten.view_as_real_copy.default>: schema = aten::view_as_real_copy(Tensor self) -> Tensor: {'self': torch.complex64, '__ret_0': torch.float32}, <EdgeOpOverload: aten.complex.default>: schema = aten::complex(Tensor real, Tensor imag) -> Tensor: {'real': torch.float32, 'imag': torch.float32, '__ret_0': torch.complex64}, <EdgeOpOverload: aten._fft_c2r.default>: schema = aten::_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor: {'self': torch.complex64, '__ret_0': torch.float32}})

[larryliu0820]

From which step did you get this error?

[gorkemalgan]

I get it in the last step

example_args = (torch.randn(1, 4, 1000, 1504),)
pre_autograd_aten_dialect = capture_pre_autograd_graph(model, example_args)
aten_dialect: ExportedProgram = export(pre_autograd_aten_dialect, example_args)
edge_program: EdgeProgramManager = to_edge(aten_dialect, compile_config=EdgeCompileConfig(_check_ir_validity=False)) # to_edge(aten_dialect)

[larryliu0820]

I think this is an valid error. It is saying tensor dtype torch.complex64 is not accepted by edge dialect (hence not accepted by ExecuTorch). We are not supporting this dtype now and is working on to get it supported.

[gorkemalgan]

torch.fft.rfft returns torch.complex64 dtype tensor so as I understand there is no way to run it with ExecuTorch for now? Correct me if I am wrong, but I will need to wait till torch.complex64 type is supported by ExecuTorch

[larryliu0820]

torch.fft.rfft returns torch.complex64 dtype tensor so as I understand there is no way to run it with ExecuTorch for now? Correct me if I am wrong, but I will need to wait till torch.complex64 type is supported by ExecuTorch

Yep I don’t think we have complex number support in ExecuTorch runtime right now

[kimishpatel]

@guangy10 if we have triaging and followup meeting, I would like to follow up on this issue specifically.

@gorkemalgan note that there are a few issues here. 1) complex dtype not supported 2) implementation of *fft not available

[cccclai]

I think we already have the complex data type issue filed separately here: #886

[cbilgin]

@gorkemalgan Closing because it's a duplicate of #886. If this is a separate issue, please do open it again