EnzymeAD / EnzymeAD/Enzyme-JAX

Incorrect batching for StableHLO `map`

Open
#2,887 3 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
MLIR
Stars
131
Forks
53
Avg merge
1d 10h
Merged PRs (30d)
193

Description

I believe Enzyme batches the following minimal `map` example incorrectly
```mlir
module @root {
func.func @func1(%arg0: tensor<1xf64>) -> tensor<1xf64> {
%0 = "stablehlo.map"(%arg0) <{dimensions = array}> ({
^bb0(%arg1: tensor):
%1 = stablehlo.sqrt %arg1 : tensor
stablehlo.return %1 : tensor
}) : (tensor<1xf64>) -> tensor<1xf64>
return %0 : tensor<1xf64>
}
func.func @main() -> tensor<1x1xf64> {
%cst = stablehlo.constant dense<2.000000e+00> : tensor<1x1xf64>
%0 = enzyme.batch @func1(%cst) {batch_shape = array} : (tensor<1x1xf64>) -> tensor<1x1xf64>
return %0 : tensor<1x1xf64>
}
}
```
I see the following error
```
$ ./enzymexlamlir-opt test.mlir --enzyme-batch
test.mlir:3:10: error: computation arguments must be 0-rank tensor, but got: arg #0 of type 'tensor<1xf64>'
%0 = "stablehlo.map"(%arg0) <{dimensions = array}> ({
^
test.mlir:3:10: error: 'stablehlo.map' op failed to infer returned types
%0 = "stablehlo.map"(%arg0) <{dimensions = array}> ({
^
test.mlir:3:10: note: see current operation:
%0 = "stablehlo.map"(%arg0) <{dimensions = array}> ({
^bb0(%arg1: tensor<1xf64>):
%1 = "stablehlo.sqrt"(%arg1) : (tensor<1xf64>) -> tensor<1xf64>
"stablehlo.return"(%1) : (tensor<1xf64>) -> ()
}) : (tensor<1x1xf64>) -> tensor<1x1xf64>
```
suggesting it's trying to batch the function passed to `map`. I believe the correct output would be closer to
```mlir
module @root {
func.func @main() -> tensor<1x1xf64> {
%cst = stablehlo.constant dense<2.000000e+00> : tensor<1x1xf64>
%0 = call @batched_func1(%cst) : (tensor<1x1xf64>) -> tensor<1x1xf64>
return %0 : tensor<1x1xf64>
}
func.func @batched_func1(%arg0: tensor<1x1xf64>) -> tensor<1x1xf64> {
%0 = "stablehlo.map"(%arg0) <{dimensions = array}> ({
^bb0(%arg1: tensor):
%1 = stablehlo.sqrt %arg1 : tensor
stablehlo.return %1 : tensor
}) : (tensor<1x1xf64>) -> tensor<1x1xf64>
return %0 : tensor<1x1xf64>
}
}
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the minimal reproducer in test.mlir and run enzymexlamlir-opt with --enzyme-batch to observe the StableHLO map failure. Trace the batching entry point that handles enzyme.batch around stablehlo.map. Done means the example lowers without the computation-argument error and produces the expected batched map dimensions shown in the issue.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
65/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.