llvm / llvm/torch-mlir

LowerToBackendContract fails due to index put style operation

Open
#1,925 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
C++
Stars
1.9k
Forks
736
Avg merge
5d 22h
Merged PRs (30d)
15

Description

I have the following pytorch model:
```
import torch
import torch_mlir

class Model(torch.nn.Module):
def __init__(self):
super().__init__()

def forward(self, x : torch.Tensor, y : torch.Tensor):
x[0, :, :] = y
return x.sum()

model = Model()
model.train(False)
inputs = [
torch.rand((3, 224, 224), dtype=torch.float32),
torch.rand((224, 224), dtype=torch.float32),
]
module = torch_mlir.compile(model, inputs, use_tracing=True, verbose=True)
print(module)
```

When I run this, it gives me the following error:
```
:0: error: unsupported by backend contract: tensor with unknown rank
:0: note: see current operation: %5 = "torch.tensor_static_info_cast"(%arg0) : (!torch.vtensor<[3,224,224],f32>) -> !torch.vtensor<*,f32>
:0: note: this is likely due to a missing transfer function in abstract_interp_lib_gen.py
```
This is the IR at the start:
```
module attributes {torch.debug_module_name = "Model"} {
func.func private @__torch__.Model.forward(%arg0: !torch.nn.Module<"__torch__.Model">, %arg1: !torch.tensor {torch.type_bound = !torch.vtensor<[3,224,224],f32>}, %arg2: !torch.tensor {torch.type_bound = !torch.vtensor<[224,224],f32>}) -> !torch.tensor {
%1 = torch.tensor_static_info_cast %arg1 : !torch.tensor to !torch.tensor<[3,224,224],f32>
%2 = torch.tensor_static_info_cast %arg2 : !torch.tensor to !torch.tensor<[224,224],f32>
%int0 = torch.constant.int 0
%int0_0 = torch.constant.int 0
%3 = torch.aten.select.int %1, %int0, %int0_0 : !torch.tensor<[3,224,224],f32>, !torch.int, !torch.int -> !torch.tensor<[224,224],f32>
%int0_1 = torch.constant.int 0
%int0_2 = torch.constant.int 0
%int9223372036854775807 = torch.constant.int 9223372036854775807
%int1 = torch.constant.int 1
%4 = torch.aten.slice.Tensor %3, %int0_1, %int0_2, %int9223372036854775807, %int1 : !torch.tensor<[224,224],f32>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.tensor<[224,224],f32>
%int1_3 = torch.constant.int 1
%int0_4 = torch.constant.int 0
%int9223372036854775807_5 = torch.constant.int 9223372036854775807
%int1_6 = torch.constant.int 1
%5 = torch.aten.slice.Tensor %4, %int1_3, %int0_4, %int9223372036854775807_5, %int1_6 : !torch.tensor<[224,224],f32>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.tensor<[224,224],f32>
%false_7 = torch.constant.bool false
%6 = torch.aten.copy_ %5, %2, %false_7 : !torch.tensor<[224,224],f32>, !torch.tensor<[224,224],f32>, !torch.bool -> !torch.tensor<[224,224],f32>
%none_8 = torch.constant.none
%7 = torch.aten.sum %1, %none_8 : !torch.tensor<[3,224,224],f32>, !torch.none -> !torch.tensor<[],f32>
%8 = torch.tensor_static_info_cast %7 : !torch.tensor<[],f32> to !torch.tensor
return %8 : !torch.tensor
}
torch.class_type @__torch__.Model {
torch.attr private "training" : !torch.bool
torch.attr private "_is_full_backward_hook" : !torch.optional
torch.method "forward", @__torch__.Model.forward
}
%false = torch.constant.bool false
%none = torch.constant.none
%0 = torch.nn_module {
torch.slot "training", %false : !torch.bool
torch.slot "_is_full_backward_hook", %none : !torch.none
} : !torch.nn.Module<"__torch__.Model">
}
```

The is the IR at LowerToBackendContract (which fails):
```
module attributes {torch.debug_module_name = "Model"} {
func.func @forward(%arg0: !torch.vtensor<[3,224,224],f32>, %arg1: !torch.vtensor<[224,224],f32>) -> !torch.vtensor<[],f32> {
%int224 = torch.constant.int 224
%int1 = torch.constant.int 1
%int0 = torch.constant.int 0
%int9223372036854775807 = torch.constant.int 9223372036854775807
%none = torch.constant.none
%0 = torch.tensor_static_info_cast %arg0 : !torch.vtensor<[3,224,224],f32> to !torch.vtensor<*,f32>
%1 = torch.copy.to_tensor %0 : !torch.tensor<*,f32>
%2 = torch.tensor_static_info_cast %1 : !torch.tensor<*,f32> to !torch.tensor<[3,224,224],f32>
%3 = torch.aten.slice.Tensor %2, %int0, %int0, %int1, %int1 : !torch.tensor<[3,224,224],f32>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.tensor<[1,224,224],f32>
%4 = torch.aten.squeeze.dim %3, %int0 : !torch.tensor<[1,224,224],f32>, !torch.int -> !torch.tensor<[224,224],f32>
%5 = torch.aten.slice.Tensor %4, %int0, %int0, %int9223372036854775807, %int1 : !torch.tensor<[224,224],f32>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.tensor<[224,224],f32>
%6 = torch.aten.slice.Tensor %5, %int1, %int0, %int9223372036854775807, %int1 : !torch.tensor<[224,224],f32>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.tensor<[224,224],f32>
%7 = torch.prim.ListConstruct %int224, %int224 : (!torch.int, !torch.int) -> !torch.list
%8 = torch.aten.broadcast_to %arg1, %7 : !torch.vtensor<[224,224],f32>, !torch.list -> !torch.vtensor<[224,224],f32>
torch.overwrite.tensor.contents %8 overwrites %6 : !torch.vtensor<[224,224],f32>, !torch.tensor<[224,224],f32>
%9 = torch.copy.to_vtensor %2 : !torch.vtensor<[3,224,224],f32>
%10 = torch.aten.sum %9, %none : !torch.vtensor<[3,224,224],f32>, !torch.none -> !torch.vtensor<[],f32>
return %10 : !torch.vtensor<[],f32>
}
}
```

After talking to @ramiro050, he mentioned that this is because pytorch breaks this index put operation into a bunch of slice + copy operations and then because the sum is performed on the mutated tensor, value semantics cannot be achieved. He mentioned that https://github.com/llvm/torch-mlir/pull/1901 is trying to fix this, but that this is a more complicated case which has multiple slice operations that should be matched.

Thanks in advance for your help! 😄

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the LowerToBackendContract failure and the generated IR in this issue, then inspect the transfer logic mentioned in abstract_interp_lib_gen.py. Compare the slice-and-copy handling with PR #1901 and determine how multiple slice operations should be matched so the mutated tensor can satisfy the backend contract; done means the reproducer lowers without an unknown-rank error.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.