EnzymeAD / EnzymeAD/Reactant.jl
Big compile times on small Tenet functions
- Dominant language
- Julia
- Stars
- 370
- Forks
- 74
- Avg merge
- 18h 47m
- Merged PRs (30d)
- 30
Description
I'm trying to get down the time taken to compile Tenet.jl.
Right now I'm using this script:
```julia
using Pkg
pkg"activate --temp"
pkg"dev Tenet Reactant"
pkg"add Adapt Enzyme"
using Tenet, Reactant, Adapt, Enzyme
a = Tensor(rand(2, 2), (:i, :k))
b = Tensor(rand(2, 2), (:k, :j))
a_re = adapt(ConcreteRArray, a)
b_re = adapt(ConcreteRArray, b)
function f(x, y)
return Enzyme.gradient(ReverseWithPrimal, Tenet.contract, x, y)
end
f_re = @compile f(a_re, b_re)
```
I'm uploading the profiles of `@compile` here: [tenet-profile-reactant-compile-run.zip](https://github.com/user-attachments/files/18531676/tenet-profile-reactant-compile-run.zip)
What surprises me most is the 2nd run. It should be instantaneous yet it shows a ~80% compile time. The profile shows that a large part of this time is spent in set operations (intersect, union and setdiff), which I can confirm that these are always set operations between `Vector{Symbol}`s, `NTuple{N,Symbol}` or `Set{Symbol}`. **All the leaves** of the flamegraph are calls to `make_oc_dict`, so @wsmoses could it be that we are not correctly caching the `OpaqueClosure`s?
I tried adding `@nospecializeinfer`, `@noinline` and `Base.inferencebarrier` to my `Tenet.contract` specialization for `TracedRArray`s, but nothing. `@reactant_overlay` seems to make it disappear from the profile, but compile time remains the same.
It would be nice to have a better way to stop abstract interpretation at some points where we know that it's not worth it because we are already in a Reactant-specific method and we are just emitting MLIR ops, not tracing below.
CC @gbaraldi since I believe you're in charge of getting compile times down within the GB project.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.