tldr; it would be nice if ther would be a way skipping the conservative is const check:
Enzyme.trust_me_its_const(::Type(typeof(f)) = true
in caseses where f without annotation fails to prove that f is const to make rules work with default AutoEnzyme() AD type.
Context
I am trying to make a function which encloses a PreallocationTools cache Enzyme diffable.
There is a stalled attempt in that repo, I am attempting something different here. My goal is to ignore the preallocated cache and allocate a new one within an Enzyme context, making my f truly constant.
I did so by adding a top level Enzyme rule which skips the problematic PreallocationTools.get_tmp call to get the enclosed cache and instead allocates a fresh cache. I then apply Enzyme to the rest of the function.
Question: is it even fine to use Enzyme.autodiff inside a rule? If you're not supposed to do that stop reading and close the issue.
The rule seems to work fine if the function is annotated as const. The problem is that in default AutoEnzyme() that is not the case, instead (as discussed in #1669) Enzyme will perform a conservative check on f to proof its const.
It would be nice to "mark" Callable as const, such that the conservative check does not fail anymore and the function gets compatible with pure AutoEnzyme without passing explicit function annotations.
using Enzyme, PreallocationTools
struct Callable
cache::DiffCache{Vector{Float64}, Vector{Float64}}
end
Callable(n::Int) = Callable(DiffCache(zeros(n)))
dim(f::Callable) = length(f.cache.du)
function (f::Callable)(res, x)
cache = get_tmp(f.cache, eltype(x))
inner(res, x, cache)
end
# Shared computation kernel — used by ALL paths (primal, ForwardDiff, Enzyme).
# cache is a plain working buffer; caller is responsible for providing it.
function inner(res, x, cache)
cache .= x
for i in eachindex(cache)
cache[i] = cache[i]^i
end
res .= cache
nothing
end
# enzym forward rule, which allocates a fresh cache and then runs enzyme on the inner function
function Enzyme.EnzymeRules.forward(
config,
f::Const{<:Callable},
fret, # allways nothing
res::Duplicated,
x,
)
n = length(f.val.cache.du)
dup_cache = Duplicated(zeros(eltype(res.val), n), zeros(eltype(res.val), n))
Enzyme.autodiff(
Enzyme.Forward # <- in my actual code i also enable runtime activity here
Const(inner),
res,
x,
dup_cache,
)
return nothing
end
N = 4
f = Callable(N)
x0 = ones(N)
dx = [0.0, 1.0, 0.0, 0.0] # seed direction: d/dx[1]
res = zeros(N)
dres = zeros(N)
mode = Enzyme.Forward
# this one works
Enzyme.autodiff(mode, Const(f), Const, Duplicated(res, dres), Duplicated(x0, dx))
# this one doesnt
Enzyme.autodiff(mode, f, Const, Duplicated(res, dres), Duplicated(x0, dx))
tldr; it would be nice if ther would be a way skipping the conservative is const check:
in caseses where
fwithout annotation fails to prove thatfis const to make rules work with defaultAutoEnzyme()AD type.Context
I am trying to make a function which encloses a
PreallocationToolscache Enzyme diffable.There is a stalled attempt in that repo, I am attempting something different here. My goal is to ignore the preallocated cache and allocate a new one within an Enzyme context, making my
ftruly constant.I did so by adding a top level Enzyme rule which skips the problematic
PreallocationTools.get_tmpcall to get the enclosed cache and instead allocates a fresh cache. I then apply Enzyme to the rest of the function.Question: is it even fine to use
Enzyme.autodiffinside a rule? If you're not supposed to do that stop reading and close the issue.The rule seems to work fine if the function is annotated as const. The problem is that in default
AutoEnzyme()that is not the case, instead (as discussed in #1669) Enzyme will perform a conservative check onfto proof its const.It would be nice to "mark"
Callableas const, such that the conservative check does not fail anymore and the function gets compatible with pureAutoEnzymewithout passing explicit function annotations.