EnzymeAD / EnzymeAD/Enzyme-JAX
conditionally perform convert(concat) -> concat(convert)
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```julia
julia> function fn(x)
Float32.(vcat(x[1:4], -x[5:8]))
end
fn (generic function with 1 method)
julia> @code_hlo fn(Reactant.to_rarray(rand(BFloat16, 8)))
module @reactant_fn attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<8xbf16>) -> tensor<8xf32> {
%0 = stablehlo.slice %arg0 [0:4] : (tensor<8xbf16>) -> tensor<4xbf16>
%1 = stablehlo.slice %arg0 [4:8] : (tensor<8xbf16>) -> tensor<4xbf16>
%2 = stablehlo.negate %1 : tensor<4xbf16>
%3 = stablehlo.convert %0 : (tensor<4xbf16>) -> tensor<4xf32>
%4 = stablehlo.convert %2 : (tensor<4xbf16>) -> tensor<4xf32>
%5 = stablehlo.concatenate %3, %4, dim = 0 : (tensor<4xf32>, tensor<4xf32>) -> tensor<8xf32>
return %5 : tensor<8xf32>
}
}
```
This might be good with we are reducing precision but when increasing precision we will end up increasing memory requirements
Contributor guide
No contributing guide indexed for this repository
Research direction
Reproduce the Julia fn example using Reactant.to_rarray and inspect its @code_hlo output. Start by tracing the handling of stablehlo.convert and stablehlo.concatenate; completion should conditionally use the requested ordering while avoiding increased memory when precision is widened, with coverage for the shown case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- compilers, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 32/100