EnzymeAD / EnzymeAD/Enzyme.jl

Add trait to skip conservative f const check?

Open
#2,974 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

**tldr;** it would be nice if ther would be a way skipping the conservative is const check:
```julia
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](https://github.com/SciML/PreallocationTools.jl/pull/141) 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.

```julia
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))
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.