EnzymeAD / EnzymeAD/Enzyme-JAX

Reuse common power op

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.