michel2323
commented
4 months ago MWE without KA and only CUDA because CUDABackend is required to trigger the segfault:
using Enzyme
using CUDA
using EnzymeCore
using EnzymeCore.EnzymeRules
# Commented lines would work
# struct Backend end
# my_synchronize(::Backend) = synchronize()
my_synchronize(::CUDABackend) = synchronize()
function EnzymeRules.augmented_primal(
config::Config,
func::Const{typeof(my_synchronize)},
::Type{Const{Nothing}},
backend::T
) where T <: EnzymeCore.Annotation
my_synchronize(backend.val)
println("Forward synchronize")
return AugmentedReturn(
nothing, nothing, nothing
)
end
function EnzymeRules.reverse(config::Config, func::Const{typeof(my_synchronize)}, ::Type{Const{Nothing}}, tape, backend)
return (nothing,)
end
function square_caller(backend)
my_synchronize(backend)
return nothing
end
function enzyme_testsuite(backend)
Enzyme.autodiff(
Reverse, square_caller, Const(backend())
)
end
# enzyme_testsuite(Backend)
enzyme_testsuite(CUDABackend)
wsmoses
I'm using the
customconstbranch https://github.com/EnzymeAD/Enzyme.jl/pull/1371 to add a synchronize rule. I did a rebase of the branch onmain. I know this worked some time before the ABI fix and whatnot like a month ago pre-0.12 release. Now it segaults inEnzymeRegisterDiffUseCallHandler(name, handle)which was added through this PR.