EnzymeAD / EnzymeAD/Enzyme-JAX

long chain of self additions to mul

Open
#1,836 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.