EnzymeAD / EnzymeAD/Enzyme-JAX
Reuse common power op
- Dominant language
- MLIR
- Stars
- 133
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
> ```mlir
> julia> @code_hlo optimize=true Enzyme.gradient(ReverseWithPrimal, f8pow, x)
> module @reactant_gradient attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
> func.func @main(%arg0: tensor) -> (tensor, tensor) {
> %cst = stablehlo.constant dense<7.000000e+00> : tensor
> %cst_0 = stablehlo.constant dense<8.000000e+00> : tensor
> %0 = stablehlo.power %arg0, %cst_0 : tensor
> %1 = stablehlo.power %arg0, %cst : tensor
> %2 = stablehlo.multiply %cst_0, %1 : tensor
> return %2, %0 : tensor, tensor
> }
> }
> ```
>
> We also need a pass to rewrite the first power as a multiply
_Originally posted by @avik-pal in [#1077](https://github.com/EnzymeAD/Enzyme-JAX/issues/1077#issuecomment-2993227741)_
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by reproducing the shown `@code_hlo optimize=true Enzyme.gradient(ReverseWithPrimal, f8pow, x)` output and inspect the generated StableHLO operations. Review the context from issue #1077, then identify where a pass can reuse the common power operation and rewrite the first `stablehlo.power` as a multiply. Done means the output no longer contains the redundant power operation and the required power-to-multiply rewrite is covered.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- compilers
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100