Skip to content

Commit 925e8f3

Browse files
committed
Update EnzymeExt.jl
1 parent 7516ccb commit 925e8f3

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

ext/EnzymeExt.jl

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,12 +44,12 @@ EnzymeRules.inactive(::Type{StaticSize}, x...) = nothing
4444
# https://github.com/EnzymeAD/Enzyme.jl/issues/1516
4545
# On the CPU `autodiff_deferred` can deadlock.
4646
# Hence a specialized CPU version
47-
function cpu_fwd(config, ctx, f, args...)
47+
function cpu_fwd(ctx, config, f, args...)
4848
EnzymeCore.autodiff(EnzymeCore.set_runtime_activity(Forward, config), Const(f), Const{Nothing}, Const(ctx), args...)
4949
return nothing
5050
end
5151

52-
function gpu_fwd(ctx, f, args...)
52+
function gpu_fwd(ctx, config, f, args...)
5353
EnzymeCore.autodiff_deferred(EnzymeCore.set_runtime_activity(Forward, config), Const(f), Const{Nothing}, Const(ctx), args...)
5454
return nothing
5555
end
@@ -66,7 +66,7 @@ function EnzymeRules.forward(
6666
f = kernel.f
6767
fwd_kernel = similar(config, kernel, cpu_fwd)
6868

69-
fwd_kernel(f, args...; ndrange, workgroupsize)
69+
fwd_kernel(config, f, args...; ndrange, workgroupsize)
7070
end
7171

7272
function EnzymeRules.forward(
@@ -81,7 +81,7 @@ function EnzymeRules.forward(
8181
f = kernel.f
8282
fwd_kernel = similar(config, kernel, gpu_fwd)
8383

84-
fwd_kernel(f, args...; ndrange, workgroupsize)
84+
fwd_kernel(config, f, args...; ndrange, workgroupsize)
8585
end
8686

8787
_enzyme_mkcontext(kernel::Kernel{CPU}, ndrange, iterspace, dynamic) =

0 commit comments

Comments
 (0)