Big compile times on small Tenet functions

[mofeing]

I'm trying to get down the time taken to compile Tenet.jl.

Right now I'm using this script:

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

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 OpaqueClosures?

I tried adding @nospecializeinfer, @noinline and Base.inferencebarrier to my Tenet.contract specialization for TracedRArrays, 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.