EnzymeAD / EnzymeAD/Enzyme-JAX
long chain of self additions to mul
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```mlir
module @jit_fn attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
func.func public @main(%arg0: tensor<1000x1000xf32>) -> (tensor<1000x1000xf32> {jax.result_info = "result"}) {
%cst = stablehlo.constant dense<0.000000e+00> : tensor
%0 = stablehlo.reduce(%arg0 init: %cst) applies stablehlo.add across dimensions = [0, 1] : (tensor<1000x1000xf32>, tensor) -> tensor
%1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor) -> tensor<1000x1000xf32>
%2 = stablehlo.add %1, %arg0 : tensor<1000x1000xf32>
%3 = stablehlo.add %2, %arg0 : tensor<1000x1000xf32>
%4 = stablehlo.add %3, %arg0 : tensor<1000x1000xf32>
%5 = stablehlo.add %4, %arg0 : tensor<1000x1000xf32>
%6 = stablehlo.add %5, %arg0 : tensor<1000x1000xf32>
%7 = stablehlo.add %6, %arg0 : tensor<1000x1000xf32>
%8 = stablehlo.add %7, %arg0 : tensor<1000x1000xf32>
%9 = stablehlo.add %8, %arg0 : tensor<1000x1000xf32>
%10 = stablehlo.add %9, %arg0 : tensor<1000x1000xf32>
return %10 : tensor<1000x1000xf32>
}
}
module @jit_fn attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
func.func public @main(%arg0: tensor<1000x1000xf32>) -> (tensor<1000x1000xf32> {jax.result_info = "result"}) {
%cst = stablehlo.constant dense<0.000000e+00> : tensor
%0 = stablehlo.reduce(%arg0 init: %cst) applies stablehlo.add across dimensions = [0, 1] : (tensor<1000x1000xf32>, tensor) -> tensor
%1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor) -> tensor<1000x1000xf32>
%2 = stablehlo.add %1, %arg0 : tensor<1000x1000xf32>
%3 = stablehlo.add %arg0, %arg0 : tensor<1000x1000xf32>
%4 = stablehlo.add %3, %3 : tensor<1000x1000xf32>
%5 = stablehlo.add %4, %4 : tensor<1000x1000xf32>
%6 = stablehlo.add %2, %5 : tensor<1000x1000xf32>
return %6 : tensor<1000x1000xf32>
}
}
```
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the two MLIR modules in the issue and compare the long chain of additions with the shorter self-addition chain. Determine which transformation or behavior is unexpected and identify the relevant compiler entry point before defining a regression test. Done means the behavior is explained and covered by a focused test.
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
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100