[Torch] Add patterns to convert complex conversions to bitcasts #19396
Conversation
The `torch.view_as_complex` and `torch.view_as_real` ops are better suited to be represented as bitcasts, however there isn't a suitable bitcast op upstream. Currently torch-mlir lowers to a difficult to analyze linalg.generic, but in IREE we can instead lower to `flow.tensor.bitcast`.
qedawkins
requested review from
MaheshRavishankar,
rsuderman,
zjgarvey and
IanWood1
There was a problem hiding this comment.
Do we want to convert it to bitcasts? I am asking cause I am guessing you are planning to fuse across these bitcasts into a dispatch. Cant do that with flow.tensor.bitcast. From previous conversation tensor.bitcast wont cut it/needs to be adapted. To avoid protracted discussion upstream on bitcast semantics maybe we add a util.tensor.bitcast which we can chose to fuse into dispatches if needed.
There was a problem hiding this comment.
I don't see anything wrong about the code or tests here, but I don't have enough context to know if this is the right approach overall. In my opinion, anything is better than linalg.generic ops that use tensor.extract instead of iterating over the input tensor (which is the case for Aten::ViewAsComplexOp).
Regarding the comment by Mahesh, why can't the flow tensor bitcast ops be fused into a dispatch?
The
torch.view_as_complexandtorch.view_as_realops are better suited to be represented as bitcasts, however there isn't a suitable bitcast op upstream. Currently torch-mlir lowers to a difficult to analyze linalg.generic, but in IREE we can instead lower toflow.tensor.bitcast.