Skip to content

Add trait to skip conservative f const check? #2974

Description

@hexaeder

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions