Adding KA.synchronize rule #1464
Comments
|
MWE without KA and only CUDA because 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) |
|
@michel2323 this runs for me |
|
@michel2323 please reopen if it persists, closing |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
I'm using the
customconstbranch #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.The text was updated successfully, but these errors were encountered: