Failed to lower 'torch.aten.scatter_reduce.two' to stablehlo
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 1.9k
- Forks
- 736
- Avg merge
- 5d 22h
- Merged PRs (30d)
- 15
Description
Hi, i got some problem when trying to lower `torch.scatter_reduce()` to stablehlo
The IR that can reproduce this problem is as follows:
```
#loc1 = loc("":9:0)
module {
func.func @forward(%arg0: !torch.vtensor<[6],f32> loc("":9:0), %arg1: !torch.vtensor<[6],si64> loc("":9:0), %arg2: !torch.vtensor<[4],f32> loc("":9:0)) -> !torch.vtensor<[4],f32> {
%int0 = torch.constant.int 0 loc(#loc1)
%str = torch.constant.str "amax" loc(#loc1)
%false = torch.constant.bool false loc(#loc1)
%0 = torch.aten.scatter_reduce.two %arg2, %int0, %arg1, %arg0, %str, %false : !torch.vtensor<[4],f32>, !torch.int, !torch.vtensor<[6],si64>, !torch.vtensor<[6],f32>, !torch.str, !torch.bool -> !torch.vtensor<[4],f32> loc(#loc1)
return %0 : !torch.vtensor<[4],f32> loc(#loc1)
} loc(#loc1)
} loc(#loc)
#loc = loc(unknown)
```
The corresponding torch operation is very simple, being:
```
src = torch.tensor([1., 2., 3., 4., 5., 6.])
index = torch.tensor([0, 1, 0, 1, 2, 1])
input = torch.tensor([1., 2., 3., 4.])
output = input.scatter_reduce(0, index, src, reduce="amax", include_self=False)
```
The error is:
```
error: unknown: failed to legalize operation 'torch.constant.bool'
note: unknown: see current operation: %0 = "torch.constant.bool"() <{value = false}> : () -> !torch.bool
```
and can be produced by:
```
torch-mlir-opt -pass-pipeline='builtin.module(torch-backend-to-stablehlo-backend-pipeline)' -mlir-print-ir-after-all -mlir-disable-threading --debug
```
I think this may due to stablehlo does not have a corresponding `scatter_reduce` op?
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Reproduce the failure with torch-mlir-opt using the provided scatter_reduce.mlir input and the torch-backend-to-stablehlo-backend-pipeline. Trace legalization of torch.aten.scatter_reduce.two and torch.constant.bool; done means the reproducer lowers through the StableHLO backend without the reported legalization error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100