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

[dixinzhou]

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?