llvm / llvm/torch-mlir

Lowering program that mutates its inputs - how to lower `torch.copy.to_vtensor` and `torch.overwrite.tensor.contents` ops

Open
#3,944 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
1.9k
Forks
736
Avg merge
5d 22h
Merged PRs (30d)
15

Description

Given a program `model` that mutates its inputs,
```
class Basic(torch.nn.Module):
def forward(self, x):
x.mul_(2)
return x
```
the fx_importer `fx.export_and_import(model, output_type=output = OutputType.TORCH, experimental_support_mutation=True)` produces the following mlir.
```
func.func @main(%arg0: !torch.tensor<[3,4],f32>) -> !torch.vtensor<[3,4],f32> attributes {torch.assume_strict_symbolic_shapes} {
%int2 = torch.constant.int 2
%0 = torch.copy.to_vtensor %arg0 : !torch.vtensor<[3,4],f32>
%1 = torch.aten.mul.Scalar %0, %int2 : !torch.vtensor<[3,4],f32>, !torch.int -> !torch.vtensor<[3,4],f32>
torch.overwrite.tensor.contents %1 overwrites %arg0 : !torch.vtensor<[3,4],f32>, !torch.tensor<[3,4],f32>
return %1 : !torch.vtensor<[3,4],f32>
}
```
The `torch.copy.to_vtensor` and `torch.overwrite.tensor.contents` ops model the mutation semantics in the `torch` level. But how to lower the two ops?

I can think of something in the `memref` dialect that can represent these semantics. But when the output is set to be `linalg_on_tensors` or `tosa`, how the lowering should look like?

Contributor guide

No contributing guide indexed for this repository

Research direction

Start at the fx_importer path that emits torch.copy.to_vtensor and torch.overwrite.tensor.contents, then trace the lowering pipeline for the linalg_on_tensors and tosa output modes. Done means both operations have defined lowerings that preserve input-mutation semantics in each requested output mode; the issue names no tests or files to run.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.