@@ -44,12 +44,12 @@ EnzymeRules.inactive(::Type{StaticSize}, x...) = nothing
44
44
# https://github.com/EnzymeAD/Enzyme.jl/issues/1516
45
45
# On the CPU `autodiff_deferred` can deadlock.
46
46
# Hence a specialized CPU version
47
- function cpu_fwd (config, ctx , f, args... )
47
+ function cpu_fwd (ctx, config , f, args... )
48
48
EnzymeCore. autodiff (EnzymeCore. set_runtime_activity (Forward, config), Const (f), Const{Nothing}, Const (ctx), args... )
49
49
return nothing
50
50
end
51
51
52
- function gpu_fwd (ctx, f, args... )
52
+ function gpu_fwd (ctx, config, f, args... )
53
53
EnzymeCore. autodiff_deferred (EnzymeCore. set_runtime_activity (Forward, config), Const (f), Const{Nothing}, Const (ctx), args... )
54
54
return nothing
55
55
end
@@ -66,7 +66,7 @@ function EnzymeRules.forward(
66
66
f = kernel. f
67
67
fwd_kernel = similar (config, kernel, cpu_fwd)
68
68
69
- fwd_kernel (f, args... ; ndrange, workgroupsize)
69
+ fwd_kernel (config, f, args... ; ndrange, workgroupsize)
70
70
end
71
71
72
72
function EnzymeRules. forward (
@@ -81,7 +81,7 @@ function EnzymeRules.forward(
81
81
f = kernel. f
82
82
fwd_kernel = similar (config, kernel, gpu_fwd)
83
83
84
- fwd_kernel (f, args... ; ndrange, workgroupsize)
84
+ fwd_kernel (config, f, args... ; ndrange, workgroupsize)
85
85
end
86
86
87
87
_enzyme_mkcontext (kernel:: Kernel{CPU} , ndrange, iterspace, dynamic) =
0 commit comments