EnzymeAD / EnzymeAD/Enzyme-JAX
Incorrect batching for StableHLO `map`
- 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