Enzyme: simplify via mixedduplicated #483
Conversation
| @@ -150,7 +126,7 @@ module EnzymeExt | |||
| args2 = ntuple(Val(N)) do i | |||
| Base.@_inline_meta | |||
| if args[i] isa Active | |||
| Duplicated(Ref(args[i].val), arg_refs[i]) | |||
| EnzymeCore.MixedDuplicated(args[i].val, arg_refs[i]) | |||
There was a problem hiding this comment.
This is correct for the CPU, but I think we won't be able to pass MixedDuplicated to the GPU.
Or at least we need an Adapt rule.
There was a problem hiding this comment.
Do we not have those already like we do for the other annotations?
There was a problem hiding this comment.
Apparently not, now we do here: EnzymeAD/Enzyme.jl#1551
There was a problem hiding this comment.
We should double check that this adapts to what we want on CUDA. I don't know if CuRef is right here...
There was a problem hiding this comment.
Sure but this is code wouldn’t affect any currently landed cuda support (this is only reverse mode).
So we should likely explore that in michel’s PR?
|
@vchuravy in interim could we get this merged/tagged (is blocking CliMA/Oceananigans.jl#3618 / https://buildkite.com/clima/oceananigans/builds/16126#01903c25-aa5f-4e96-9088-70984b90ebe4 seemingly) |
|
@wsmoses just so that you aware, this is what |
No description provided.