JuliaSIMD / JuliaSIMD/LoopVectorization.jl

Performance issue with Zygote.jl-generated function within `@turbo`

Open
#489 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
789
Forks
73
PR merge metrics
No merged PRs in 30d

Description

I am trying to understand a performance issue I am seeing in [DynamicExpressions.jl](https://github.com/SymbolicML/DynamicExpressions.jl) where using `@turbo` makes the evaluation kernels 4x faster, but makes the derivative kernels 10% _slower_. See the detailed benchmarks here: https://github.com/SymbolicML/DynamicExpressions.jl/pull/28#issuecomment-1529108964

My derivative kernels look like [this](https://github.com/SymbolicML/DynamicExpressions.jl/blob/1fa158687f3d3a0e9c53d13fa99ac7b04a2c6aed/src/EvaluateEquationDerivative.jl#L128-L134):

```julia
@maybe_turbo turbo for j in indices((cumulator, dcumulator))
x = op(cumulator[j])::T
dx = diff_op(cumulator[j])::T * dcumulator[j]

cumulator[j] = x
dcumulator[j] = dx
end
```

(The `@maybe_turbo turbo ...` will turn into `@turbo ...` when `turbo=true`, but just `@inbounds @simd ...` otherwise. It will also remove the various type assertions in the scope.)

To create the `diff_op`, I generate it using Zygote.jl [here](https://github.com/SymbolicML/DynamicExpressions.jl/blob/1fa158687f3d3a0e9c53d13fa99ac7b04a2c6aed/src/OperatorEnumConstruction.jl#L229-L232):

```julia
for op in unary_operators
diff_op(x) = gradient(op, x)[1]
push!(diff_unary_operators, diff_op)
end
```

I can try to create a MWE for this, but I quickly wanted to check if anything was obvious in how I am using `@turbo` here that might hurt performance rather than help it. For example, perhaps this `diff_op` is not being inlined correctly, and therefore not being optimized by `@turbo`? For the record I am not seeing any warnings about the derivative operator being incompatible, so I'm not quite sure why this is occurring.

Also - the `diff_op` in the benchmark is the derivative of one of `+, -, *, /, cos, exp` so nothing too crazy.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.