EnzymeAD / EnzymeAD/Reactant.jl

Big compile times on small Tenet functions

Open
#605 0 comments 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.