Lowering program that mutates its inputs - how to lower `torch.copy.to_vtensor` and `torch.overwrite.tensor.contents` ops
- 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